ref: caf8df647886963848a2e266124de54fcc737247
dir: /sys/lib/python/mercurial/sshserver.py/
# sshserver.py - ssh protocol server support for mercurial # # Copyright 2005-2007 Matt Mackall <[email protected]> # Copyright 2006 Vadim Gelfer <[email protected]> # # This software may be used and distributed according to the terms of the # GNU General Public License version 2, incorporated herein by reference. from i18n import _ from node import bin, hex import streamclone, util, hook import os, sys, tempfile, urllib class sshserver(object): def __init__(self, ui, repo): self.ui = ui self.repo = repo self.lock = None self.fin = sys.stdin self.fout = sys.stdout hook.redirect(True) sys.stdout = sys.stderr # Prevent insertion/deletion of CRs util.set_binary(self.fin) util.set_binary(self.fout) def getarg(self): argline = self.fin.readline()[:-1] arg, l = argline.split() val = self.fin.read(int(l)) return arg, val def respond(self, v): self.fout.write("%d\n" % len(v)) self.fout.write(v) self.fout.flush() def serve_forever(self): try: while self.serve_one(): pass finally: if self.lock is not None: self.lock.release() sys.exit(0) def serve_one(self): cmd = self.fin.readline()[:-1] if cmd: impl = getattr(self, 'do_' + cmd, None) if impl: impl() else: self.respond("") return cmd != '' def do_lookup(self): arg, key = self.getarg() assert arg == 'key' try: r = hex(self.repo.lookup(key)) success = 1 except Exception, inst: r = str(inst) success = 0 self.respond("%s %s\n" % (success, r)) def do_branchmap(self): branchmap = self.repo.branchmap() heads = [] for branch, nodes in branchmap.iteritems(): branchname = urllib.quote(branch) branchnodes = [hex(node) for node in nodes] heads.append('%s %s' % (branchname, ' '.join(branchnodes))) self.respond('\n'.join(heads)) def do_heads(self): h = self.repo.heads() self.respond(" ".join(map(hex, h)) + "\n") def do_hello(self): '''the hello command returns a set of lines describing various interesting things about the server, in an RFC822-like format. Currently the only one defined is "capabilities", which consists of a line in the form: capabilities: space separated list of tokens ''' caps = ['unbundle', 'lookup', 'changegroupsubset', 'branchmap'] if self.ui.configbool('server', 'uncompressed'): caps.append('stream=%d' % self.repo.changelog.version) self.respond("capabilities: %s\n" % (' '.join(caps),)) def do_lock(self): '''DEPRECATED - allowing remote client to lock repo is not safe''' self.lock = self.repo.lock() self.respond("") def do_unlock(self): '''DEPRECATED''' if self.lock: self.lock.release() self.lock = None self.respond("") def do_branches(self): arg, nodes = self.getarg() nodes = map(bin, nodes.split(" ")) r = [] for b in self.repo.branches(nodes): r.append(" ".join(map(hex, b)) + "\n") self.respond("".join(r)) def do_between(self): arg, pairs = self.getarg() pairs = [map(bin, p.split("-")) for p in pairs.split(" ")] r = [] for b in self.repo.between(pairs): r.append(" ".join(map(hex, b)) + "\n") self.respond("".join(r)) def do_changegroup(self): nodes = [] arg, roots = self.getarg() nodes = map(bin, roots.split(" ")) cg = self.repo.changegroup(nodes, 'serve') while True: d = cg.read(4096) if not d: break self.fout.write(d) self.fout.flush() def do_changegroupsubset(self): argmap = dict([self.getarg(), self.getarg()]) bases = [bin(n) for n in argmap['bases'].split(' ')] heads = [bin(n) for n in argmap['heads'].split(' ')] cg = self.repo.changegroupsubset(bases, heads, 'serve') while True: d = cg.read(4096) if not d: break self.fout.write(d) self.fout.flush() def do_addchangegroup(self): '''DEPRECATED''' if not self.lock: self.respond("not locked") return self.respond("") r = self.repo.addchangegroup(self.fin, 'serve', self.client_url()) self.respond(str(r)) def client_url(self): client = os.environ.get('SSH_CLIENT', '').split(' ', 1)[0] return 'remote:ssh:' + client def do_unbundle(self): their_heads = self.getarg()[1].split() def check_heads(): heads = map(hex, self.repo.heads()) return their_heads == [hex('force')] or their_heads == heads # fail early if possible if not check_heads(): self.respond(_('unsynced changes')) return self.respond('') # write bundle data to temporary file because it can be big tempname = fp = None try: fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-') fp = os.fdopen(fd, 'wb+') count = int(self.fin.readline()) while count: fp.write(self.fin.read(count)) count = int(self.fin.readline()) was_locked = self.lock is not None if not was_locked: self.lock = self.repo.lock() try: if not check_heads(): # someone else committed/pushed/unbundled while we # were transferring data self.respond(_('unsynced changes')) return self.respond('') # push can proceed fp.seek(0) r = self.repo.addchangegroup(fp, 'serve', self.client_url()) self.respond(str(r)) finally: if not was_locked: self.lock.release() self.lock = None finally: if fp is not None: fp.close() if tempname is not None: os.unlink(tempname) def do_stream_out(self): try: for chunk in streamclone.stream_out(self.repo): self.fout.write(chunk) self.fout.flush() except streamclone.StreamException, inst: self.fout.write(str(inst)) self.fout.flush()