hgweb/sshserver: extract capabilities for easier modification

mercurial/hgweb/protocol.py

@@ -1,207 +1,208
 #
 # 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, incorporated herein by reference.
 
-import cStringIO, zlib, tempfile, errno, os, sys, urllib
+import cStringIO, zlib, tempfile, errno, os, sys, urllib, copy
 from mercurial import util, streamclone
 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',
 ]
 
 HGTYPE = 'application/mercurial-0.1'
+basecaps = 'lookup changegroupsubset branchmap'.split()
 
 def lookup(repo, req):
     try:
         r = hex(repo.lookup(req.form['key'][0]))
         success = 1
     except Exception, inst:
         r = str(inst)
         success = 0
     resp = "%s %s\n" % (success, r)
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     yield resp
 
 def heads(repo, req):
     resp = " ".join(map(hex, repo.heads())) + "\n"
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     yield resp
 
 def branchmap(repo, req):
     branches = repo.branchmap()
     heads = []
     for branch, nodes in branches.iteritems():
         branchname = urllib.quote(branch)
         branchnodes = [hex(node) for node in nodes]
         heads.append('%s %s' % (branchname, ' '.join(branchnodes)))
     resp = '\n'.join(heads)
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     yield resp
 
 def branches(repo, req):
     nodes = []
     if 'nodes' in req.form:
         nodes = map(bin, req.form['nodes'][0].split(" "))
     resp = cStringIO.StringIO()
     for b in repo.branches(nodes):
         resp.write(" ".join(map(hex, b)) + "\n")
     resp = resp.getvalue()
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     yield resp
 
 def between(repo, req):
     if 'pairs' in req.form:
         pairs = [map(bin, p.split("-"))
                  for p in req.form['pairs'][0].split(" ")]
     resp = cStringIO.StringIO()
     for b in repo.between(pairs):
         resp.write(" ".join(map(hex, b)) + "\n")
     resp = resp.getvalue()
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     yield resp
 
 def changegroup(repo, req):
     req.respond(HTTP_OK, HGTYPE)
     nodes = []
 
     if 'roots' in req.form:
         nodes = map(bin, req.form['roots'][0].split(" "))
 
     z = zlib.compressobj()
     f = repo.changegroup(nodes, 'serve')
     while 1:
         chunk = f.read(4096)
         if not chunk:
             break
         yield z.compress(chunk)
 
     yield z.flush()
 
 def changegroupsubset(repo, req):
     req.respond(HTTP_OK, HGTYPE)
     bases = []
     heads = []
 
     if 'bases' in req.form:
         bases = [bin(x) for x in req.form['bases'][0].split(' ')]
     if 'heads' in req.form:
         heads = [bin(x) for x in req.form['heads'][0].split(' ')]
 
     z = zlib.compressobj()
     f = repo.changegroupsubset(bases, heads, 'serve')
     while 1:
         chunk = f.read(4096)
         if not chunk:
             break
         yield z.compress(chunk)
 
     yield z.flush()
 
 def capabilities(repo, req):
-    caps = ['lookup', 'changegroupsubset', 'branchmap']
+    caps = copy.copy(basecaps)
     if repo.ui.configbool('server', 'uncompressed', untrusted=True):
         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
 
 def unbundle(repo, req):
 
     proto = req.env.get('wsgi.url_scheme') or 'http'
     their_heads = req.form['heads'][0].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():
         req.drain()
         raise ErrorResponse(HTTP_OK, 'unsynced changes')
 
     # do not lock repo until all changegroup data is
     # streamed. save to temporary file.
 
     fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
     fp = os.fdopen(fd, 'wb+')
     try:
         length = int(req.env['CONTENT_LENGTH'])
         for s in util.filechunkiter(req, limit=length):
             fp.write(s)
 
         try:
             lock = repo.lock()
             try:
                 if not check_heads():
                     raise ErrorResponse(HTTP_OK, 'unsynced changes')
 
                 fp.seek(0)
                 header = fp.read(6)
                 if header.startswith('HG') and 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)
 
                 # send addchangegroup output to client
 
                 oldio = sys.stdout, sys.stderr
                 sys.stderr = sys.stdout = cStringIO.StringIO()
 
                 try:
                     url = 'remote:%s:%s:%s' % (
                           proto,
                           urllib.quote(req.env.get('REMOTE_HOST', '')),
                           urllib.quote(req.env.get('REMOTE_USER', '')))
                     try:
                         ret = repo.addchangegroup(gen, 'serve', url)
                     except util.Abort, inst:
                         sys.stdout.write("abort: %s\n" % inst)
                         ret = 0
                 finally:
                     val = sys.stdout.getvalue()
                     sys.stdout, sys.stderr = oldio
                 req.respond(HTTP_OK, HGTYPE)
                 return '%d\n%s' % (ret, val),
             finally:
                 lock.release()
         except ValueError, inst:
             raise ErrorResponse(HTTP_OK, inst)
         except (OSError, IOError), inst:
             error = getattr(inst, 'strerror', 'Unknown error')
             if inst.errno == errno.ENOENT:
                 code = HTTP_NOT_FOUND
             else:
                 code = HTTP_SERVER_ERROR
             filename = getattr(inst, 'filename', '')
             # Don't send our filesystem layout to the client
             if filename and filename.startswith(repo.root):
                 filename = filename[len(repo.root)+1:]
                 text = '%s: %s' % (error, filename)
             else:
                 text = error.replace(repo.root + os.path.sep, '')
             raise ErrorResponse(code, text)
     finally:
         fp.close()
         os.unlink(tempname)
 
 def stream_out(repo, req):
     req.respond(HTTP_OK, HGTYPE)
     try:
         for chunk in streamclone.stream_out(repo, untrusted=True):
             yield chunk
     except streamclone.StreamException, inst:
         yield str(inst)

mercurial/sshserver.py

@@ -1,225 +1,227
 # 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, incorporated herein by reference.
 
 from i18n import _
 from node import bin, hex
 import streamclone, util, hook
-import os, sys, tempfile, urllib
+import os, sys, tempfile, urllib, copy
 
 class sshserver(object):
+
+    caps = 'unbundle lookup changegroupsubset branchmap'.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 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']
+        caps = copy.copy(self.caps)
         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()