protocol: unify server-side capabilities functions

mercurial/hgweb/protocol.py

@@ -1,34 +1,23
 #
 # Copyright 21 May 2005 - (c) 2005 Jake Edge <[e-mail]>
 # Copyright 2005-2007 Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 import cStringIO, zlib, tempfile, errno, os, sys, urllib, copy
 from mercurial import util, streamclone, pushkey
 from mercurial.node import bin, hex
 from mercurial import changegroup as changegroupmod
 from common import ErrorResponse, HTTP_OK, HTTP_NOT_FOUND, HTTP_SERVER_ERROR
 
 # __all__ is populated with the allowed commands. Be sure to add to it if
 # you're adding a new command, or the new command won't work.
 
 __all__ = [
    'lookup', 'heads', 'branches', 'between', 'changegroup',
    'changegroupsubset', 'capabilities', 'unbundle', 'stream_out',
    'branchmap', 'pushkey', 'listkeys'
 ]
 
 HGTYPE = 'application/mercurial-0.1'
-basecaps = 'lookup changegroupsubset branchmap pushkey'.split()
-
-def capabilities(repo, req):
-    caps = copy.copy(basecaps)
-    if streamclone.allowed(repo.ui):
-        caps.append('stream=%d' % repo.changelog.version)
-    if changegroupmod.bundlepriority:
-        caps.append('unbundle=%s' % ','.join(changegroupmod.bundlepriority))
-    rsp = ' '.join(caps)
-    req.respond(HTTP_OK, HGTYPE, length=len(rsp))
-    yield rsp

mercurial/sshserver.py

@@ -1,150 +1,134
 # sshserver.py - ssh protocol server support for mercurial
 #
 # Copyright 2005-2007 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 version 2 or any later version.
 
 from i18n import _
 from node import bin, hex
 import streamclone, util, hook, pushkey, wireproto
 import os, sys, tempfile, urllib, copy
 
 class sshserver(object):
-
-    caps = 'unbundle lookup changegroupsubset branchmap pushkey'.split()
-
     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 getargs(self, args):
         data = {}
         keys = args.split()
         count = len(keys)
         for n in xrange(len(keys)):
             argline = self.fin.readline()[:-1]
             arg, l = argline.split()
             val = self.fin.read(int(l))
             if arg not in keys:
                 raise util.Abort("unexpected parameter %r" % arg)
             if arg == '*':
                 star = {}
                 for n in xrange(int(l)):
                     arg, l = argline.split()
                     val = self.fin.read(int(l))
                     star[arg] = val
                 data['*'] = star
             else:
                 data[arg] = val
         return [data[k] for k in keys]
 
     def getarg(self, name):
         return self.getargs(name)[0]
 
     def respond(self, v):
         self.fout.write("%d\n" % len(v))
         self.fout.write(v)
         self.fout.flush()
 
     def sendchangegroup(self, changegroup):
         while True:
             d = changegroup.read(4096)
             if not d:
                 break
             self.fout.write(d)
 
         self.fout.flush()
 
     def sendstream(self, source):
         for chunk in source:
             self.fout.write(chunk)
         self.fout.flush()
 
     def getfile(self, fpout):
         self.respond('')
         count = int(self.fin.readline())
         while count:
             fpout.write(self.fin.read(count))
             count = int(self.fin.readline())
 
     def redirect(self):
         pass
 
     def respondpush(self, ret):
         self.respond('')
         self.respond(str(ret))
 
     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 and not wireproto.dispatch(self.repo, self, cmd):
             impl = getattr(self, 'do_' + cmd, None)
             if impl:
                 r = impl()
                 if r is not None:
                     self.respond(r)
             else: self.respond("")
         return cmd != ''
 
-    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 = copy.copy(self.caps)
-        if streamclone.allowed(self.repo.ui):
-            caps.append('stream=%d' % self.repo.changelog.version)
-        return "capabilities: %s\n" % (' '.join(caps),)
-
     def do_lock(self):
         '''DEPRECATED - allowing remote client to lock repo is not safe'''
 
         self.lock = self.repo.lock()
         return ""
 
     def do_unlock(self):
         '''DEPRECATED'''
 
         if self.lock:
             self.lock.release()
         self.lock = None
         return ""
 
     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(),
                                      lock=self.lock)
         return str(r)
 
     def _client(self):
         client = os.environ.get('SSH_CLIENT', '').split(' ', 1)[0]
         return 'remote:ssh:' + client

mercurial/wireproto.py

@@ -1,264 +1,283
 # wireproto.py - generic wire protocol support functions
 #
 # Copyright 2005-2010 Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 import urllib, tempfile, os
 from i18n import _
 from node import bin, hex
 import changegroup as changegroupmod
 import streamclone, repo, error, encoding, util
 import pushkey as pushkey_
 
 # client side
 
 class wirerepository(repo.repository):
     def lookup(self, key):
         self.requirecap('lookup', _('look up remote revision'))
         d = self._call("lookup", key=key)
         success, data = d[:-1].split(" ", 1)
         if int(success):
             return bin(data)
         self._abort(error.RepoError(data))
 
     def heads(self):
         d = self._call("heads")
         try:
             return map(bin, d[:-1].split(" "))
         except:
             self.abort(error.ResponseError(_("unexpected response:"), d))
 
     def branchmap(self):
         d = self._call("branchmap")
         try:
             branchmap = {}
             for branchpart in d.splitlines():
                 branchheads = branchpart.split(' ')
                 branchname = urllib.unquote(branchheads[0])
                 # Earlier servers (1.3.x) send branch names in (their) local
                 # charset. The best we can do is assume it's identical to our
                 # own local charset, in case it's not utf-8.
                 try:
                     branchname.decode('utf-8')
                 except UnicodeDecodeError:
                     branchname = encoding.fromlocal(branchname)
                 branchheads = [bin(x) for x in branchheads[1:]]
                 branchmap[branchname] = branchheads
             return branchmap
         except TypeError:
             self._abort(error.ResponseError(_("unexpected response:"), d))
 
     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._abort(error.ResponseError(_("unexpected response:"), d))
 
     def between(self, pairs):
         batch = 8 # avoid giant requests
         r = []
         for i in xrange(0, len(pairs), batch):
             n = " ".join(["-".join(map(hex, p)) for p in pairs[i:i + batch]])
             d = self._call("between", pairs=n)
             try:
                 r += [l and map(bin, l.split(" ")) or []
                       for l in d.splitlines()]
             except:
                 self._abort(error.ResponseError(_("unexpected response:"), d))
         return r
 
     def pushkey(self, namespace, key, old, new):
         if not self.capable('pushkey'):
             return False
         d = self._call("pushkey",
                       namespace=namespace, key=key, old=old, new=new)
         return bool(int(d))
 
     def listkeys(self, namespace):
         if not self.capable('pushkey'):
             return {}
         d = self._call("listkeys", namespace=namespace)
         r = {}
         for l in d.splitlines():
             k, v = l.split('\t')
             r[k.decode('string-escape')] = v.decode('string-escape')
         return r
 
     def stream_out(self):
         return self._callstream('stream_out')
 
     def changegroup(self, nodes, kind):
         n = " ".join(map(hex, nodes))
         f = self._callstream("changegroup", roots=n)
         return self._decompress(f)
 
     def changegroupsubset(self, bases, heads, kind):
         self.requirecap('changegroupsubset', _('look up remote changes'))
         bases = " ".join(map(hex, bases))
         heads = " ".join(map(hex, heads))
         return self._decompress(self._callstream("changegroupsubset",
                                                  bases=bases, heads=heads))
 
     def unbundle(self, cg, heads, source):
         '''Send cg (a readable file-like object representing the
         changegroup to push, typically a chunkbuffer object) to the
         remote server as a bundle. Return an integer indicating the
         result of the push (see localrepository.addchangegroup()).'''
 
         ret, output = self._callpush("unbundle", cg, heads=' '.join(map(hex, heads)))
         if ret == "":
             raise error.ResponseError(
                 _('push failed:'), output)
         try:
             ret = int(ret)
         except ValueError, err:
             raise error.ResponseError(
                 _('push failed (unexpected response):'), ret)
 
         for l in output.splitlines(True):
             self.ui.status(_('remote: '), l)
         return ret
 
 # server side
 
 def dispatch(repo, proto, command):
     if command not in commands:
         return False
     func, spec = commands[command]
     args = proto.getargs(spec)
     r = func(repo, proto, *args)
     if r != None:
         proto.respond(r)
     return True
 
 def between(repo, proto, pairs):
     pairs = [map(bin, p.split("-")) for p in pairs.split(" ")]
     r = []
     for b in repo.between(pairs):
         r.append(" ".join(map(hex, b)) + "\n")
     return "".join(r)
 
 def branchmap(repo, proto):
     branchmap = 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)))
     return '\n'.join(heads)
 
 def branches(repo, proto, nodes):
     nodes = map(bin, nodes.split(" "))
     r = []
     for b in repo.branches(nodes):
         r.append(" ".join(map(hex, b)) + "\n")
     return "".join(r)
 
+def capabilities(repo, proto):
+    caps = 'lookup changegroupsubset branchmap pushkey'.split()
+    if streamclone.allowed(repo.ui):
+        caps.append('stream=%d' % repo.changelog.version)
+    caps.append('unbundle=%s' % ','.join(changegroupmod.bundlepriority))
+    return ' '.join(caps)
+
 def changegroup(repo, proto, roots):
     nodes = map(bin, roots.split(" "))
     cg = repo.changegroup(nodes, 'serve')
     proto.sendchangegroup(cg)
 
 def changegroupsubset(repo, proto, bases, heads):
     bases = [bin(n) for n in bases.split(' ')]
     heads = [bin(n) for n in heads.split(' ')]
     cg = repo.changegroupsubset(bases, heads, 'serve')
     proto.sendchangegroup(cg)
 
 def heads(repo, proto):
     h = repo.heads()
     return " ".join(map(hex, h)) + "\n"
 
+def hello(repo, proto):
+    '''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
+    '''
+    return "capabilities: %s\n" % (capabilities(repo, proto))
+
 def listkeys(repo, proto, namespace):
     d = pushkey_.list(repo, namespace).items()
     t = '\n'.join(['%s\t%s' % (k.encode('string-escape'),
                                v.encode('string-escape')) for k, v in d])
     return t
 
 def lookup(repo, proto, key):
     try:
         r = hex(repo.lookup(key))
         success = 1
     except Exception, inst:
         r = str(inst)
         success = 0
     return "%s %s\n" % (success, r)
 
 def pushkey(repo, proto, namespace, key, old, new):
     r = pushkey_.push(repo, namespace, key, old, new)
     return '%s\n' % int(r)
 
 def stream(repo, proto):
     try:
         proto.sendstream(streamclone.stream_out(repo))
     except streamclone.StreamException, inst:
         return str(inst)
 
 def unbundle(repo, proto, heads):
     their_heads = heads.split()
 
     def check_heads():
         heads = map(hex, repo.heads())
         return their_heads == [hex('force')] or their_heads == heads
 
     # fail early if possible
     if not check_heads():
         repo.respond(_('unsynced changes'))
         return
 
     # write bundle data to temporary file because it can be big
     fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
     fp = os.fdopen(fd, 'wb+')
     r = 0
     proto.redirect()
     try:
         proto.getfile(fp)
         lock = repo.lock()
         try:
             if not check_heads():
                 # someone else committed/pushed/unbundled while we
                 # were transferring data
                 proto.respond(_('unsynced changes'))
                 return
 
             # push can proceed
             fp.seek(0)
             header = fp.read(6)
             if header.startswith('HG'):
                 if not header.startswith('HG10'):
                     raise ValueError('unknown bundle version')
                 elif header not in changegroupmod.bundletypes:
                     raise ValueError('unknown bundle compression type')
             gen = changegroupmod.unbundle(header, fp)
 
             try:
                 r = repo.addchangegroup(gen, 'serve', proto._client(),
                                         lock=lock)
             except util.Abort, inst:
                 sys.stderr.write("abort: %s\n" % inst)
         finally:
             lock.release()
             proto.respondpush(r)
 
     finally:
         fp.close()
         os.unlink(tempname)
 
 commands = {
     'between': (between, 'pairs'),
     'branchmap': (branchmap, ''),
     'branches': (branches, 'nodes'),
+    'capabilities': (capabilities, ''),
     'changegroup': (changegroup, 'roots'),
     'changegroupsubset': (changegroupsubset, 'bases heads'),
     'heads': (heads, ''),
+    'hello': (hello, ''),
     'listkeys': (listkeys, 'namespace'),
     'lookup': (lookup, 'key'),
     'pushkey': (pushkey, 'namespace key old new'),
     'stream_out': (stream, ''),
     'unbundle': (unbundle, 'heads'),
 }