Adding changegroupsubset and lookup to ssh protocol so pull -r and

clone -r can be supported.

mercurial/sshrepo.py

@@ -1,211 +1,223
 # sshrepo.py - ssh repository proxy class for mercurial
 #
 # Copyright 2005, 2006 Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms
 # of the GNU General Public License, incorporated herein by reference.
 
 from node import *
 from remoterepo import *
 from i18n import gettext as _
 from demandload import *
 demandload(globals(), "hg os re stat util")
 
 class sshrepository(remoterepository):
     def __init__(self, ui, path, create=0):
         self._url = path
         self.ui = ui
 
         m = re.match(r'ssh://(([^@]+)@)?([^:/]+)(:(\d+))?(/(.*))?', path)
         if not m:
             self.repoerror(_("couldn't parse location %s") % path)
 
         self.user = m.group(2)
         self.host = m.group(3)
         self.port = m.group(5)
         self.path = m.group(7) or "."
 
         args = self.user and ("%s@%s" % (self.user, self.host)) or self.host
         args = self.port and ("%s -p %s") % (args, self.port) or args
 
         sshcmd = self.ui.config("ui", "ssh", "ssh")
         remotecmd = self.ui.config("ui", "remotecmd", "hg")
 
         if create:
             cmd = '%s %s "%s init %s"'
             cmd = cmd % (sshcmd, args, remotecmd, self.path)
 
             ui.note('running %s\n' % cmd)
             res = os.system(cmd)
             if res != 0:
                 self.repoerror(_("could not create remote repo"))
 
         self.validate_repo(ui, sshcmd, args, remotecmd)
 
     def url(self):
         return self._url
 
     def validate_repo(self, ui, sshcmd, args, remotecmd):
         # cleanup up previous run
         self.cleanup()
 
         cmd = '%s %s "%s -R %s serve --stdio"'
         cmd = cmd % (sshcmd, args, remotecmd, self.path)
 
         ui.note('running %s\n' % cmd)
         self.pipeo, self.pipei, self.pipee = os.popen3(cmd, 'b')
 
         # skip any noise generated by remote shell
         self.do_cmd("hello")
         r = self.do_cmd("between", pairs=("%s-%s" % ("0"*40, "0"*40)))
         lines = ["", "dummy"]
         max_noise = 500
         while lines[-1] and max_noise:
             l = r.readline()
             self.readerr()
             if lines[-1] == "1\n" and l == "\n":
                 break
             if l:
                 ui.debug(_("remote: "), l)
             lines.append(l)
             max_noise -= 1
         else:
             self.repoerror(_("no suitable response from remote hg"))
 
         self.capabilities = ()
         lines.reverse()
         for l in lines:
             if l.startswith("capabilities:"):
                 self.capabilities = l[:-1].split(":")[1].split()
                 break
 
     def readerr(self):
         while 1:
             size = util.fstat(self.pipee).st_size
             if size == 0: break
             l = self.pipee.readline()
             if not l: break
             self.ui.status(_("remote: "), l)
 
     def repoerror(self, msg):
         self.cleanup()
         raise hg.RepoError(msg)
 
     def cleanup(self):
         try:
             self.pipeo.close()
             self.pipei.close()
             # read the error descriptor until EOF
             for l in self.pipee:
                 self.ui.status(_("remote: "), l)
             self.pipee.close()
         except:
             pass
 
     __del__ = cleanup
 
     def do_cmd(self, cmd, **args):
         self.ui.debug(_("sending %s command\n") % cmd)
         self.pipeo.write("%s\n" % cmd)
         for k, v in args.items():
             self.pipeo.write("%s %d\n" % (k, len(v)))
             self.pipeo.write(v)
         self.pipeo.flush()
 
         return self.pipei
 
     def call(self, cmd, **args):
         r = self.do_cmd(cmd, **args)
         l = r.readline()
         self.readerr()
         try:
             l = int(l)
         except:
             self.repoerror(_("unexpected response '%s'") % l)
         return r.read(l)
 
     def lock(self):
         self.call("lock")
         return remotelock(self)
 
     def unlock(self):
         self.call("unlock")
 
+    def lookup(self, key):
+        d = self.call("lookup", key=key)
+        try:
+            return bin(d[:-1])
+        except:
+            raise hg.RepoError("unexpected response '%s'" % (d[:400] + "..."))
+
     def heads(self):
         d = self.call("heads")
         try:
             return map(bin, d[:-1].split(" "))
         except:
             self.repoerror(_("unexpected response '%s'") % (d[:400] + "..."))
 
     def branches(self, nodes):
         n = " ".join(map(hex, nodes))
         d = self.call("branches", nodes=n)
         try:
             br = [ tuple(map(bin, b.split(" "))) for b in d.splitlines() ]
             return br
         except:
             self.repoerror(_("unexpected response '%s'") % (d[:400] + "..."))
 
     def between(self, pairs):
         n = "\n".join(["-".join(map(hex, p)) for p in pairs])
         d = self.call("between", pairs=n)
         try:
             p = [ l and map(bin, l.split(" ")) or [] for l in d.splitlines() ]
             return p
         except:
             self.repoerror(_("unexpected response '%s'") % (d[:400] + "..."))
 
     def changegroup(self, nodes, kind):
         n = " ".join(map(hex, nodes))
         return self.do_cmd("changegroup", roots=n)
 
+    def changegroupsubset(self, bases, heads, kind):
+        bases = " ".join(map(hex, bases))
+        heads = " ".join(map(hex, heads))
+        return self.do_cmd("changegroupsubset", bases=bases, heads=heads)
+
     def unbundle(self, cg, heads, source):
         d = self.call("unbundle", heads=' '.join(map(hex, heads)))
         if d:
             self.repoerror(_("push refused: %s") % d)
 
         while 1:
             d = cg.read(4096)
             if not d: break
             self.pipeo.write(str(len(d)) + '\n')
             self.pipeo.write(d)
             self.readerr()
 
         self.pipeo.write('0\n')
         self.pipeo.flush()
 
         self.readerr()
         d = self.pipei.readline()
         if d != '\n':
             return 1
 
         l = int(self.pipei.readline())
         r = self.pipei.read(l)
         if not r:
             return 1
         return int(r)
 
     def addchangegroup(self, cg, source, url):
         d = self.call("addchangegroup")
         if d:
             self.repoerror(_("push refused: %s") % d)
         while 1:
             d = cg.read(4096)
             if not d: break
             self.pipeo.write(d)
             self.readerr()
 
         self.pipeo.flush()
 
         self.readerr()
         l = int(self.pipei.readline())
         r = self.pipei.read(l)
         if not r:
             return 1
         return int(r)
 
     def stream_out(self):
         return self.do_cmd('stream_out')
 
 instance = sshrepository

mercurial/sshserver.py

@@ -1,178 +1,199
 # sshserver.py - ssh protocol server support for mercurial
 #
 # Copyright 2005 Matt Mackall <[e-mail]>
 # Copyright 2006 Vadim Gelfer <[e-mail]>
 #
 # This software may be used and distributed according to the terms
 # of the GNU General Public License, incorporated herein by reference.
 
 from demandload import demandload
 from i18n import gettext as _
 from node import *
 demandload(globals(), "os streamclone sys tempfile util")
 
 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
 
         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):
         while self.serve_one(): pass
         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'
+        self.respond(hex(self.repo.lookup(key)) + "\n")
+
     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']
+        caps = ['unbundle', 'lookup', 'changegroupsubset']
         if self.ui.configbool('server', 'uncompressed'):
             caps.append('stream=%d' % self.repo.revlogversion)
         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):
+        bases = []
+        heads = []
+        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
 
         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:
             fp.close()
             os.unlink(tempname)
 
     def do_stream_out(self):
         streamclone.stream_out(self.repo, self.fout)