wireproto: move command registration types to wireprototypes

These are shared across wire protocol implementations. wireprototypes
is our module for common code.

Differential Revision: https://phab.mercurial-scm.org/D3396

mercurial/wireproto.py

@@ -1,809 +1,730
 # 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.
 
 from __future__ import absolute_import
 
 import os
 import tempfile
 
 from .i18n import _
 from .node import (
     hex,
     nullid,
 )
 
 from . import (
     bundle2,
     changegroup as changegroupmod,
     discovery,
     encoding,
     error,
     exchange,
     pushkey as pushkeymod,
     pycompat,
     streamclone,
     util,
     wireprototypes,
 )
 
 from .utils import (
     procutil,
     stringutil,
 )
 
 urlerr = util.urlerr
 urlreq = util.urlreq
 
 bundle2requiredmain = _('incompatible Mercurial client; bundle2 required')
 bundle2requiredhint = _('see https://www.mercurial-scm.org/wiki/'
                         'IncompatibleClient')
 bundle2required = '%s\n(%s)\n' % (bundle2requiredmain, bundle2requiredhint)
 
 def clientcompressionsupport(proto):
     """Returns a list of compression methods supported by the client.
 
     Returns a list of the compression methods supported by the client
     according to the protocol capabilities. If no such capability has
     been announced, fallback to the default of zlib and uncompressed.
     """
     for cap in proto.getprotocaps():
         if cap.startswith('comp='):
             return cap[5:].split(',')
     return ['zlib', 'none']
 
 # wire protocol command can either return a string or one of these classes.
 
 def getdispatchrepo(repo, proto, command):
     """Obtain the repo used for processing wire protocol commands.
 
     The intent of this function is to serve as a monkeypatch point for
     extensions that need commands to operate on different repo views under
     specialized circumstances.
     """
     return repo.filtered('served')
 
 def dispatch(repo, proto, command):
     repo = getdispatchrepo(repo, proto, command)
 
     transportversion = wireprototypes.TRANSPORTS[proto.name]['version']
     commandtable = commandsv2 if transportversion == 2 else commands
     func, spec = commandtable[command]
 
     args = proto.getargs(spec)
 
     # Version 1 protocols define arguments as a list. Version 2 uses a dict.
     if isinstance(args, list):
         return func(repo, proto, *args)
     elif isinstance(args, dict):
         return func(repo, proto, **args)
     else:
         raise error.ProgrammingError('unexpected type returned from '
                                      'proto.getargs(): %s' % type(args))
 
 def options(cmd, keys, others):
     opts = {}
     for k in keys:
         if k in others:
             opts[k] = others[k]
             del others[k]
     if others:
         procutil.stderr.write("warning: %s ignored unexpected arguments %s\n"
                               % (cmd, ",".join(others)))
     return opts
 
 def bundle1allowed(repo, action):
     """Whether a bundle1 operation is allowed from the server.
 
     Priority is:
 
     1. server.bundle1gd.<action> (if generaldelta active)
     2. server.bundle1.<action>
     3. server.bundle1gd (if generaldelta active)
     4. server.bundle1
     """
     ui = repo.ui
     gd = 'generaldelta' in repo.requirements
 
     if gd:
         v = ui.configbool('server', 'bundle1gd.%s' % action)
         if v is not None:
             return v
 
     v = ui.configbool('server', 'bundle1.%s' % action)
     if v is not None:
         return v
 
     if gd:
         v = ui.configbool('server', 'bundle1gd')
         if v is not None:
             return v
 
     return ui.configbool('server', 'bundle1')
 
 def supportedcompengines(ui, role):
     """Obtain the list of supported compression engines for a request."""
     assert role in (util.CLIENTROLE, util.SERVERROLE)
 
     compengines = util.compengines.supportedwireengines(role)
 
     # Allow config to override default list and ordering.
     if role == util.SERVERROLE:
         configengines = ui.configlist('server', 'compressionengines')
         config = 'server.compressionengines'
     else:
         # This is currently implemented mainly to facilitate testing. In most
         # cases, the server should be in charge of choosing a compression engine
         # because a server has the most to lose from a sub-optimal choice. (e.g.
         # CPU DoS due to an expensive engine or a network DoS due to poor
         # compression ratio).
         configengines = ui.configlist('experimental',
                                       'clientcompressionengines')
         config = 'experimental.clientcompressionengines'
 
     # No explicit config. Filter out the ones that aren't supposed to be
     # advertised and return default ordering.
     if not configengines:
         attr = 'serverpriority' if role == util.SERVERROLE else 'clientpriority'
         return [e for e in compengines
                 if getattr(e.wireprotosupport(), attr) > 0]
 
     # If compression engines are listed in the config, assume there is a good
     # reason for it (like server operators wanting to achieve specific
     # performance characteristics). So fail fast if the config references
     # unusable compression engines.
     validnames = set(e.name() for e in compengines)
     invalidnames = set(e for e in configengines if e not in validnames)
     if invalidnames:
         raise error.Abort(_('invalid compression engine defined in %s: %s') %
                           (config, ', '.join(sorted(invalidnames))))
 
     compengines = [e for e in compengines if e.name() in configengines]
     compengines = sorted(compengines,
                          key=lambda e: configengines.index(e.name()))
 
     if not compengines:
         raise error.Abort(_('%s config option does not specify any known '
                             'compression engines') % config,
                           hint=_('usable compression engines: %s') %
                           ', '.sorted(validnames))
 
     return compengines
 
-class commandentry(object):
-    """Represents a declared wire protocol command."""
-    def __init__(self, func, args='', transports=None,
-                 permission='push'):
-        self.func = func
-        self.args = args
-        self.transports = transports or set()
-        self.permission = permission
-
-    def _merge(self, func, args):
-        """Merge this instance with an incoming 2-tuple.
-
-        This is called when a caller using the old 2-tuple API attempts
-        to replace an instance. The incoming values are merged with
-        data not captured by the 2-tuple and a new instance containing
-        the union of the two objects is returned.
-        """
-        return commandentry(func, args=args, transports=set(self.transports),
-                            permission=self.permission)
-
-    # Old code treats instances as 2-tuples. So expose that interface.
-    def __iter__(self):
-        yield self.func
-        yield self.args
-
-    def __getitem__(self, i):
-        if i == 0:
-            return self.func
-        elif i == 1:
-            return self.args
-        else:
-            raise IndexError('can only access elements 0 and 1')
-
-class commanddict(dict):
-    """Container for registered wire protocol commands.
-
-    It behaves like a dict. But __setitem__ is overwritten to allow silent
-    coercion of values from 2-tuples for API compatibility.
-    """
-    def __setitem__(self, k, v):
-        if isinstance(v, commandentry):
-            pass
-        # Cast 2-tuples to commandentry instances.
-        elif isinstance(v, tuple):
-            if len(v) != 2:
-                raise ValueError('command tuples must have exactly 2 elements')
-
-            # It is common for extensions to wrap wire protocol commands via
-            # e.g. ``wireproto.commands[x] = (newfn, args)``. Because callers
-            # doing this aren't aware of the new API that uses objects to store
-            # command entries, we automatically merge old state with new.
-            if k in self:
-                v = self[k]._merge(v[0], v[1])
-            else:
-                # Use default values from @wireprotocommand.
-                v = commandentry(v[0], args=v[1],
-                                 transports=set(wireprototypes.TRANSPORTS),
-                                 permission='push')
-        else:
-            raise ValueError('command entries must be commandentry instances '
-                             'or 2-tuples')
-
-        return super(commanddict, self).__setitem__(k, v)
-
-    def commandavailable(self, command, proto):
-        """Determine if a command is available for the requested protocol."""
-        assert proto.name in wireprototypes.TRANSPORTS
-
-        entry = self.get(command)
-
-        if not entry:
-            return False
-
-        if proto.name not in entry.transports:
-            return False
-
-        return True
-
 # For version 1 transports.
-commands = commanddict()
+commands = wireprototypes.commanddict()
 
 # For version 2 transports.
-commandsv2 = commanddict()
+commandsv2 = wireprototypes.commanddict()
 
 def wireprotocommand(name, args=None, permission='push'):
     """Decorator to declare a wire protocol command.
 
     ``name`` is the name of the wire protocol command being provided.
 
     ``args`` defines the named arguments accepted by the command. It is
     a space-delimited list of argument names. ``*`` denotes a special value
     that says to accept all named arguments.
 
     ``permission`` defines the permission type needed to run this command.
     Can be ``push`` or ``pull``. These roughly map to read-write and read-only,
     respectively. Default is to assume command requires ``push`` permissions
     because otherwise commands not declaring their permissions could modify
     a repository that is supposed to be read-only.
     """
     transports = {k for k, v in wireprototypes.TRANSPORTS.items()
                   if v['version'] == 1}
 
     # Because SSHv2 is a mirror of SSHv1, we allow "batch" commands through to
     # SSHv2.
     # TODO undo this hack when SSH is using the unified frame protocol.
     if name == b'batch':
         transports.add(wireprototypes.SSHV2)
 
     if permission not in ('push', 'pull'):
         raise error.ProgrammingError('invalid wire protocol permission; '
                                      'got %s; expected "push" or "pull"' %
                                      permission)
 
     if args is None:
         args = ''
 
     if not isinstance(args, bytes):
         raise error.ProgrammingError('arguments for version 1 commands '
                                      'must be declared as bytes')
 
     def register(func):
         if name in commands:
             raise error.ProgrammingError('%s command already registered '
                                          'for version 1' % name)
-        commands[name] = commandentry(func, args=args,
-                                      transports=transports,
-                                      permission=permission)
+        commands[name] = wireprototypes.commandentry(
+            func, args=args, transports=transports, permission=permission)
 
         return func
     return register
 
 # TODO define a more appropriate permissions type to use for this.
 @wireprotocommand('batch', 'cmds *', permission='pull')
 def batch(repo, proto, cmds, others):
     unescapearg = wireprototypes.unescapebatcharg
     repo = repo.filtered("served")
     res = []
     for pair in cmds.split(';'):
         op, args = pair.split(' ', 1)
         vals = {}
         for a in args.split(','):
             if a:
                 n, v = a.split('=')
                 vals[unescapearg(n)] = unescapearg(v)
         func, spec = commands[op]
 
         # Validate that client has permissions to perform this command.
         perm = commands[op].permission
         assert perm in ('push', 'pull')
         proto.checkperm(perm)
 
         if spec:
             keys = spec.split()
             data = {}
             for k in keys:
                 if k == '*':
                     star = {}
                     for key in vals.keys():
                         if key not in keys:
                             star[key] = vals[key]
                     data['*'] = star
                 else:
                     data[k] = vals[k]
             result = func(repo, proto, *[data[k] for k in keys])
         else:
             result = func(repo, proto)
         if isinstance(result, wireprototypes.ooberror):
             return result
 
         # For now, all batchable commands must return bytesresponse or
         # raw bytes (for backwards compatibility).
         assert isinstance(result, (wireprototypes.bytesresponse, bytes))
         if isinstance(result, wireprototypes.bytesresponse):
             result = result.data
         res.append(wireprototypes.escapebatcharg(result))
 
     return wireprototypes.bytesresponse(';'.join(res))
 
 @wireprotocommand('between', 'pairs', permission='pull')
 def between(repo, proto, pairs):
     pairs = [wireprototypes.decodelist(p, '-') for p in pairs.split(" ")]
     r = []
     for b in repo.between(pairs):
         r.append(wireprototypes.encodelist(b) + "\n")
 
     return wireprototypes.bytesresponse(''.join(r))
 
 @wireprotocommand('branchmap', permission='pull')
 def branchmap(repo, proto):
     branchmap = repo.branchmap()
     heads = []
     for branch, nodes in branchmap.iteritems():
         branchname = urlreq.quote(encoding.fromlocal(branch))
         branchnodes = wireprototypes.encodelist(nodes)
         heads.append('%s %s' % (branchname, branchnodes))
 
     return wireprototypes.bytesresponse('\n'.join(heads))
 
 @wireprotocommand('branches', 'nodes', permission='pull')
 def branches(repo, proto, nodes):
     nodes = wireprototypes.decodelist(nodes)
     r = []
     for b in repo.branches(nodes):
         r.append(wireprototypes.encodelist(b) + "\n")
 
     return wireprototypes.bytesresponse(''.join(r))
 
 @wireprotocommand('clonebundles', '', permission='pull')
 def clonebundles(repo, proto):
     """Server command for returning info for available bundles to seed clones.
 
     Clients will parse this response and determine what bundle to fetch.
 
     Extensions may wrap this command to filter or dynamically emit data
     depending on the request. e.g. you could advertise URLs for the closest
     data center given the client's IP address.
     """
     return wireprototypes.bytesresponse(
         repo.vfs.tryread('clonebundles.manifest'))
 
 wireprotocaps = ['lookup', 'branchmap', 'pushkey',
                  'known', 'getbundle', 'unbundlehash']
 
 def _capabilities(repo, proto):
     """return a list of capabilities for a repo
 
     This function exists to allow extensions to easily wrap capabilities
     computation
 
     - returns a lists: easy to alter
     - change done here will be propagated to both `capabilities` and `hello`
       command without any other action needed.
     """
     # copy to prevent modification of the global list
     caps = list(wireprotocaps)
 
     # Command of same name as capability isn't exposed to version 1 of
     # transports. So conditionally add it.
     if commands.commandavailable('changegroupsubset', proto):
         caps.append('changegroupsubset')
 
     if streamclone.allowservergeneration(repo):
         if repo.ui.configbool('server', 'preferuncompressed'):
             caps.append('stream-preferred')
         requiredformats = repo.requirements & repo.supportedformats
         # if our local revlogs are just revlogv1, add 'stream' cap
         if not requiredformats - {'revlogv1'}:
             caps.append('stream')
         # otherwise, add 'streamreqs' detailing our local revlog format
         else:
             caps.append('streamreqs=%s' % ','.join(sorted(requiredformats)))
     if repo.ui.configbool('experimental', 'bundle2-advertise'):
         capsblob = bundle2.encodecaps(bundle2.getrepocaps(repo, role='server'))
         caps.append('bundle2=' + urlreq.quote(capsblob))
     caps.append('unbundle=%s' % ','.join(bundle2.bundlepriority))
 
     return proto.addcapabilities(repo, caps)
 
 # If you are writing an extension and consider wrapping this function. Wrap
 # `_capabilities` instead.
 @wireprotocommand('capabilities', permission='pull')
 def capabilities(repo, proto):
     caps = _capabilities(repo, proto)
     return wireprototypes.bytesresponse(' '.join(sorted(caps)))
 
 @wireprotocommand('changegroup', 'roots', permission='pull')
 def changegroup(repo, proto, roots):
     nodes = wireprototypes.decodelist(roots)
     outgoing = discovery.outgoing(repo, missingroots=nodes,
                                   missingheads=repo.heads())
     cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
     gen = iter(lambda: cg.read(32768), '')
     return wireprototypes.streamres(gen=gen)
 
 @wireprotocommand('changegroupsubset', 'bases heads',
                   permission='pull')
 def changegroupsubset(repo, proto, bases, heads):
     bases = wireprototypes.decodelist(bases)
     heads = wireprototypes.decodelist(heads)
     outgoing = discovery.outgoing(repo, missingroots=bases,
                                   missingheads=heads)
     cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
     gen = iter(lambda: cg.read(32768), '')
     return wireprototypes.streamres(gen=gen)
 
 @wireprotocommand('debugwireargs', 'one two *',
                   permission='pull')
 def debugwireargs(repo, proto, one, two, others):
     # only accept optional args from the known set
     opts = options('debugwireargs', ['three', 'four'], others)
     return wireprototypes.bytesresponse(repo.debugwireargs(
         one, two, **pycompat.strkwargs(opts)))
 
 def find_pullbundle(repo, proto, opts, clheads, heads, common):
     """Return a file object for the first matching pullbundle.
 
     Pullbundles are specified in .hg/pullbundles.manifest similar to
     clonebundles.
     For each entry, the bundle specification is checked for compatibility:
     - Client features vs the BUNDLESPEC.
     - Revisions shared with the clients vs base revisions of the bundle.
       A bundle can be applied only if all its base revisions are known by
       the client.
     - At least one leaf of the bundle's DAG is missing on the client.
     - Every leaf of the bundle's DAG is part of node set the client wants.
       E.g. do not send a bundle of all changes if the client wants only
       one specific branch of many.
     """
     def decodehexstring(s):
         return set([h.decode('hex') for h in s.split(';')])
 
     manifest = repo.vfs.tryread('pullbundles.manifest')
     if not manifest:
         return None
     res = exchange.parseclonebundlesmanifest(repo, manifest)
     res = exchange.filterclonebundleentries(repo, res)
     if not res:
         return None
     cl = repo.changelog
     heads_anc = cl.ancestors([cl.rev(rev) for rev in heads], inclusive=True)
     common_anc = cl.ancestors([cl.rev(rev) for rev in common], inclusive=True)
     compformats = clientcompressionsupport(proto)
     for entry in res:
         if 'COMPRESSION' in entry and entry['COMPRESSION'] not in compformats:
             continue
         # No test yet for VERSION, since V2 is supported by any client
         # that advertises partial pulls
         if 'heads' in entry:
             try:
                 bundle_heads = decodehexstring(entry['heads'])
             except TypeError:
                 # Bad heads entry
                 continue
             if bundle_heads.issubset(common):
                 continue # Nothing new
             if all(cl.rev(rev) in common_anc for rev in bundle_heads):
                 continue # Still nothing new
             if any(cl.rev(rev) not in heads_anc and
                    cl.rev(rev) not in common_anc for rev in bundle_heads):
                 continue
         if 'bases' in entry:
             try:
                 bundle_bases = decodehexstring(entry['bases'])
             except TypeError:
                 # Bad bases entry
                 continue
             if not all(cl.rev(rev) in common_anc for rev in bundle_bases):
                 continue
         path = entry['URL']
         repo.ui.debug('sending pullbundle "%s"\n' % path)
         try:
             return repo.vfs.open(path)
         except IOError:
             repo.ui.debug('pullbundle "%s" not accessible\n' % path)
             continue
     return None
 
 @wireprotocommand('getbundle', '*', permission='pull')
 def getbundle(repo, proto, others):
     opts = options('getbundle', wireprototypes.GETBUNDLE_ARGUMENTS.keys(),
                    others)
     for k, v in opts.iteritems():
         keytype = wireprototypes.GETBUNDLE_ARGUMENTS[k]
         if keytype == 'nodes':
             opts[k] = wireprototypes.decodelist(v)
         elif keytype == 'csv':
             opts[k] = list(v.split(','))
         elif keytype == 'scsv':
             opts[k] = set(v.split(','))
         elif keytype == 'boolean':
             # Client should serialize False as '0', which is a non-empty string
             # so it evaluates as a True bool.
             if v == '0':
                 opts[k] = False
             else:
                 opts[k] = bool(v)
         elif keytype != 'plain':
             raise KeyError('unknown getbundle option type %s'
                            % keytype)
 
     if not bundle1allowed(repo, 'pull'):
         if not exchange.bundle2requested(opts.get('bundlecaps')):
             if proto.name == 'http-v1':
                 return wireprototypes.ooberror(bundle2required)
             raise error.Abort(bundle2requiredmain,
                               hint=bundle2requiredhint)
 
     prefercompressed = True
 
     try:
         clheads = set(repo.changelog.heads())
         heads = set(opts.get('heads', set()))
         common = set(opts.get('common', set()))
         common.discard(nullid)
         if (repo.ui.configbool('server', 'pullbundle') and
             'partial-pull' in proto.getprotocaps()):
             # Check if a pre-built bundle covers this request.
             bundle = find_pullbundle(repo, proto, opts, clheads, heads, common)
             if bundle:
                 return wireprototypes.streamres(gen=util.filechunkiter(bundle),
                                                 prefer_uncompressed=True)
 
         if repo.ui.configbool('server', 'disablefullbundle'):
             # Check to see if this is a full clone.
             changegroup = opts.get('cg', True)
             if changegroup and not common and clheads == heads:
                 raise error.Abort(
                     _('server has pull-based clones disabled'),
                     hint=_('remove --pull if specified or upgrade Mercurial'))
 
         info, chunks = exchange.getbundlechunks(repo, 'serve',
                                                 **pycompat.strkwargs(opts))
         prefercompressed = info.get('prefercompressed', True)
     except error.Abort as exc:
         # cleanly forward Abort error to the client
         if not exchange.bundle2requested(opts.get('bundlecaps')):
             if proto.name == 'http-v1':
                 return wireprototypes.ooberror(pycompat.bytestr(exc) + '\n')
             raise # cannot do better for bundle1 + ssh
         # bundle2 request expect a bundle2 reply
         bundler = bundle2.bundle20(repo.ui)
         manargs = [('message', pycompat.bytestr(exc))]
         advargs = []
         if exc.hint is not None:
             advargs.append(('hint', exc.hint))
         bundler.addpart(bundle2.bundlepart('error:abort',
                                            manargs, advargs))
         chunks = bundler.getchunks()
         prefercompressed = False
 
     return wireprototypes.streamres(
         gen=chunks, prefer_uncompressed=not prefercompressed)
 
 @wireprotocommand('heads', permission='pull')
 def heads(repo, proto):
     h = repo.heads()
     return wireprototypes.bytesresponse(wireprototypes.encodelist(h) + '\n')
 
 @wireprotocommand('hello', permission='pull')
 def hello(repo, proto):
     """Called as part of SSH handshake to obtain server info.
 
     Returns a list of lines describing interesting things about the
     server, in an RFC822-like format.
 
     Currently, the only one defined is ``capabilities``, which consists of a
     line of space separated tokens describing server abilities:
 
         capabilities: <token0> <token1> <token2>
     """
     caps = capabilities(repo, proto).data
     return wireprototypes.bytesresponse('capabilities: %s\n' % caps)
 
 @wireprotocommand('listkeys', 'namespace', permission='pull')
 def listkeys(repo, proto, namespace):
     d = sorted(repo.listkeys(encoding.tolocal(namespace)).items())
     return wireprototypes.bytesresponse(pushkeymod.encodekeys(d))
 
 @wireprotocommand('lookup', 'key', permission='pull')
 def lookup(repo, proto, key):
     try:
         k = encoding.tolocal(key)
         n = repo.lookup(k)
         r = hex(n)
         success = 1
     except Exception as inst:
         r = stringutil.forcebytestr(inst)
         success = 0
     return wireprototypes.bytesresponse('%d %s\n' % (success, r))
 
 @wireprotocommand('known', 'nodes *', permission='pull')
 def known(repo, proto, nodes, others):
     v = ''.join(b and '1' or '0'
                 for b in repo.known(wireprototypes.decodelist(nodes)))
     return wireprototypes.bytesresponse(v)
 
 @wireprotocommand('protocaps', 'caps', permission='pull')
 def protocaps(repo, proto, caps):
     if proto.name == wireprototypes.SSHV1:
         proto._protocaps = set(caps.split(' '))
     return wireprototypes.bytesresponse('OK')
 
 @wireprotocommand('pushkey', 'namespace key old new', permission='push')
 def pushkey(repo, proto, namespace, key, old, new):
     # compatibility with pre-1.8 clients which were accidentally
     # sending raw binary nodes rather than utf-8-encoded hex
     if len(new) == 20 and stringutil.escapestr(new) != new:
         # looks like it could be a binary node
         try:
             new.decode('utf-8')
             new = encoding.tolocal(new) # but cleanly decodes as UTF-8
         except UnicodeDecodeError:
             pass # binary, leave unmodified
     else:
         new = encoding.tolocal(new) # normal path
 
     with proto.mayberedirectstdio() as output:
         r = repo.pushkey(encoding.tolocal(namespace), encoding.tolocal(key),
                          encoding.tolocal(old), new) or False
 
     output = output.getvalue() if output else ''
     return wireprototypes.bytesresponse('%d\n%s' % (int(r), output))
 
 @wireprotocommand('stream_out', permission='pull')
 def stream(repo, proto):
     '''If the server supports streaming clone, it advertises the "stream"
     capability with a value representing the version and flags of the repo
     it is serving. Client checks to see if it understands the format.
     '''
     return wireprototypes.streamreslegacy(
         streamclone.generatev1wireproto(repo))
 
 @wireprotocommand('unbundle', 'heads', permission='push')
 def unbundle(repo, proto, heads):
     their_heads = wireprototypes.decodelist(heads)
 
     with proto.mayberedirectstdio() as output:
         try:
             exchange.check_heads(repo, their_heads, 'preparing changes')
             cleanup = lambda: None
             try:
                 payload = proto.getpayload()
                 if repo.ui.configbool('server', 'streamunbundle'):
                     def cleanup():
                         # Ensure that the full payload is consumed, so
                         # that the connection doesn't contain trailing garbage.
                         for p in payload:
                             pass
                     fp = util.chunkbuffer(payload)
                 else:
                     # write bundle data to temporary file as it can be big
                     fp, tempname = None, None
                     def cleanup():
                         if fp:
                             fp.close()
                         if tempname:
                             os.unlink(tempname)
                     fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
                     repo.ui.debug('redirecting incoming bundle to %s\n' %
                         tempname)
                     fp = os.fdopen(fd, pycompat.sysstr('wb+'))
                     r = 0
                     for p in payload:
                         fp.write(p)
                     fp.seek(0)
 
                 gen = exchange.readbundle(repo.ui, fp, None)
                 if (isinstance(gen, changegroupmod.cg1unpacker)
                     and not bundle1allowed(repo, 'push')):
                     if proto.name == 'http-v1':
                         # need to special case http because stderr do not get to
                         # the http client on failed push so we need to abuse
                         # some other error type to make sure the message get to
                         # the user.
                         return wireprototypes.ooberror(bundle2required)
                     raise error.Abort(bundle2requiredmain,
                                       hint=bundle2requiredhint)
 
                 r = exchange.unbundle(repo, gen, their_heads, 'serve',
                                       proto.client())
                 if util.safehasattr(r, 'addpart'):
                     # The return looks streamable, we are in the bundle2 case
                     # and should return a stream.
                     return wireprototypes.streamreslegacy(gen=r.getchunks())
                 return wireprototypes.pushres(
                     r, output.getvalue() if output else '')
 
             finally:
                 cleanup()
 
         except (error.BundleValueError, error.Abort, error.PushRaced) as exc:
             # handle non-bundle2 case first
             if not getattr(exc, 'duringunbundle2', False):
                 try:
                     raise
                 except error.Abort:
                     # The old code we moved used procutil.stderr directly.
                     # We did not change it to minimise code change.
                     # This need to be moved to something proper.
                     # Feel free to do it.
                     procutil.stderr.write("abort: %s\n" % exc)
                     if exc.hint is not None:
                         procutil.stderr.write("(%s)\n" % exc.hint)
                     procutil.stderr.flush()
                     return wireprototypes.pushres(
                         0, output.getvalue() if output else '')
                 except error.PushRaced:
                     return wireprototypes.pusherr(
                         pycompat.bytestr(exc),
                         output.getvalue() if output else '')
 
             bundler = bundle2.bundle20(repo.ui)
             for out in getattr(exc, '_bundle2salvagedoutput', ()):
                 bundler.addpart(out)
             try:
                 try:
                     raise
                 except error.PushkeyFailed as exc:
                     # check client caps
                     remotecaps = getattr(exc, '_replycaps', None)
                     if (remotecaps is not None
                             and 'pushkey' not in remotecaps.get('error', ())):
                         # no support remote side, fallback to Abort handler.
                         raise
                     part = bundler.newpart('error:pushkey')
                     part.addparam('in-reply-to', exc.partid)
                     if exc.namespace is not None:
                         part.addparam('namespace', exc.namespace,
                                       mandatory=False)
                     if exc.key is not None:
                         part.addparam('key', exc.key, mandatory=False)
                     if exc.new is not None:
                         part.addparam('new', exc.new, mandatory=False)
                     if exc.old is not None:
                         part.addparam('old', exc.old, mandatory=False)
                     if exc.ret is not None:
                         part.addparam('ret', exc.ret, mandatory=False)
             except error.BundleValueError as exc:
                 errpart = bundler.newpart('error:unsupportedcontent')
                 if exc.parttype is not None:
                     errpart.addparam('parttype', exc.parttype)
                 if exc.params:
                     errpart.addparam('params', '\0'.join(exc.params))
             except error.Abort as exc:
                 manargs = [('message', stringutil.forcebytestr(exc))]
                 advargs = []
                 if exc.hint is not None:
                     advargs.append(('hint', exc.hint))
                 bundler.addpart(bundle2.bundlepart('error:abort',
                                                    manargs, advargs))
             except error.PushRaced as exc:
                 bundler.newpart('error:pushraced',
                                 [('message', stringutil.forcebytestr(exc))])
             return wireprototypes.streamreslegacy(gen=bundler.getchunks())

mercurial/wireprototypes.py

@@ -1,243 +1,321
 # Copyright 2018 Gregory Szorc <[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 __future__ import absolute_import
 
 from .node import (
     bin,
     hex,
 )
 from .thirdparty.zope import (
     interface as zi,
 )
 
 # Names of the SSH protocol implementations.
 SSHV1 = 'ssh-v1'
 # These are advertised over the wire. Increment the counters at the end
 # to reflect BC breakages.
 SSHV2 = 'exp-ssh-v2-0001'
 HTTP_WIREPROTO_V2 = 'exp-http-v2-0001'
 
 # All available wire protocol transports.
 TRANSPORTS = {
     SSHV1: {
         'transport': 'ssh',
         'version': 1,
     },
     SSHV2: {
         'transport': 'ssh',
         # TODO mark as version 2 once all commands are implemented.
         'version': 1,
     },
     'http-v1': {
         'transport': 'http',
         'version': 1,
     },
     HTTP_WIREPROTO_V2: {
         'transport': 'http',
         'version': 2,
     }
 }
 
 class bytesresponse(object):
     """A wire protocol response consisting of raw bytes."""
     def __init__(self, data):
         self.data = data
 
 class ooberror(object):
     """wireproto reply: failure of a batch of operation
 
     Something failed during a batch call. The error message is stored in
     `self.message`.
     """
     def __init__(self, message):
         self.message = message
 
 class pushres(object):
     """wireproto reply: success with simple integer return
 
     The call was successful and returned an integer contained in `self.res`.
     """
     def __init__(self, res, output):
         self.res = res
         self.output = output
 
 class pusherr(object):
     """wireproto reply: failure
 
     The call failed. The `self.res` attribute contains the error message.
     """
     def __init__(self, res, output):
         self.res = res
         self.output = output
 
 class streamres(object):
     """wireproto reply: binary stream
 
     The call was successful and the result is a stream.
 
     Accepts a generator containing chunks of data to be sent to the client.
 
     ``prefer_uncompressed`` indicates that the data is expected to be
     uncompressable and that the stream should therefore use the ``none``
     engine.
     """
     def __init__(self, gen=None, prefer_uncompressed=False):
         self.gen = gen
         self.prefer_uncompressed = prefer_uncompressed
 
 class streamreslegacy(object):
     """wireproto reply: uncompressed binary stream
 
     The call was successful and the result is a stream.
 
     Accepts a generator containing chunks of data to be sent to the client.
 
     Like ``streamres``, but sends an uncompressed data for "version 1" clients
     using the application/mercurial-0.1 media type.
     """
     def __init__(self, gen=None):
         self.gen = gen
 
 class cborresponse(object):
     """Encode the response value as CBOR."""
     def __init__(self, v):
         self.value = v
 
 class v2errorresponse(object):
     """Represents a command error for version 2 transports."""
     def __init__(self, message, args=None):
         self.message = message
         self.args = args
 
 class v2streamingresponse(object):
     """A response whose data is supplied by a generator.
 
     The generator can either consist of data structures to CBOR
     encode or a stream of already-encoded bytes.
     """
     def __init__(self, gen, compressible=True):
         self.gen = gen
         self.compressible = compressible
 
 # list of nodes encoding / decoding
 def decodelist(l, sep=' '):
     if l:
         return [bin(v) for v in  l.split(sep)]
     return []
 
 def encodelist(l, sep=' '):
     try:
         return sep.join(map(hex, l))
     except TypeError:
         raise
 
 # batched call argument encoding
 
 def escapebatcharg(plain):
     return (plain
             .replace(':', ':c')
             .replace(',', ':o')
             .replace(';', ':s')
             .replace('=', ':e'))
 
 def unescapebatcharg(escaped):
     return (escaped
             .replace(':e', '=')
             .replace(':s', ';')
             .replace(':o', ',')
             .replace(':c', ':'))
 
 # mapping of options accepted by getbundle and their types
 #
 # Meant to be extended by extensions. It is extensions responsibility to ensure
 # such options are properly processed in exchange.getbundle.
 #
 # supported types are:
 #
 # :nodes: list of binary nodes
 # :csv:   list of comma-separated values
 # :scsv:  list of comma-separated values return as set
 # :plain: string with no transformation needed.
 GETBUNDLE_ARGUMENTS = {
     'heads':  'nodes',
     'bookmarks': 'boolean',
     'common': 'nodes',
     'obsmarkers': 'boolean',
     'phases': 'boolean',
     'bundlecaps': 'scsv',
     'listkeys': 'csv',
     'cg': 'boolean',
     'cbattempted': 'boolean',
     'stream': 'boolean',
 }
 
 class baseprotocolhandler(zi.Interface):
     """Abstract base class for wire protocol handlers.
 
     A wire protocol handler serves as an interface between protocol command
     handlers and the wire protocol transport layer. Protocol handlers provide
     methods to read command arguments, redirect stdio for the duration of
     the request, handle response types, etc.
     """
 
     name = zi.Attribute(
         """The name of the protocol implementation.
 
         Used for uniquely identifying the transport type.
         """)
 
     def getargs(args):
         """return the value for arguments in <args>
 
         For version 1 transports, returns a list of values in the same
         order they appear in ``args``. For version 2 transports, returns
         a dict mapping argument name to value.
         """
 
     def getprotocaps():
         """Returns the list of protocol-level capabilities of client
 
         Returns a list of capabilities as declared by the client for
         the current request (or connection for stateful protocol handlers)."""
 
     def getpayload():
         """Provide a generator for the raw payload.
 
         The caller is responsible for ensuring that the full payload is
         processed.
         """
 
     def mayberedirectstdio():
         """Context manager to possibly redirect stdio.
 
         The context manager yields a file-object like object that receives
         stdout and stderr output when the context manager is active. Or it
         yields ``None`` if no I/O redirection occurs.
 
         The intent of this context manager is to capture stdio output
         so it may be sent in the response. Some transports support streaming
         stdio to the client in real time. For these transports, stdio output
         won't be captured.
         """
 
     def client():
         """Returns a string representation of this client (as bytes)."""
 
     def addcapabilities(repo, caps):
         """Adds advertised capabilities specific to this protocol.
 
         Receives the list of capabilities collected so far.
 
         Returns a list of capabilities. The passed in argument can be returned.
         """
 
     def checkperm(perm):
         """Validate that the client has permissions to perform a request.
 
         The argument is the permission required to proceed. If the client
         doesn't have that permission, the exception should raise or abort
         in a protocol specific manner.
         """
+
+class commandentry(object):
+    """Represents a declared wire protocol command."""
+    def __init__(self, func, args='', transports=None,
+                 permission='push'):
+        self.func = func
+        self.args = args
+        self.transports = transports or set()
+        self.permission = permission
+
+    def _merge(self, func, args):
+        """Merge this instance with an incoming 2-tuple.
+
+        This is called when a caller using the old 2-tuple API attempts
+        to replace an instance. The incoming values are merged with
+        data not captured by the 2-tuple and a new instance containing
+        the union of the two objects is returned.
+        """
+        return commandentry(func, args=args, transports=set(self.transports),
+                            permission=self.permission)
+
+    # Old code treats instances as 2-tuples. So expose that interface.
+    def __iter__(self):
+        yield self.func
+        yield self.args
+
+    def __getitem__(self, i):
+        if i == 0:
+            return self.func
+        elif i == 1:
+            return self.args
+        else:
+            raise IndexError('can only access elements 0 and 1')
+
+class commanddict(dict):
+    """Container for registered wire protocol commands.
+
+    It behaves like a dict. But __setitem__ is overwritten to allow silent
+    coercion of values from 2-tuples for API compatibility.
+    """
+    def __setitem__(self, k, v):
+        if isinstance(v, commandentry):
+            pass
+        # Cast 2-tuples to commandentry instances.
+        elif isinstance(v, tuple):
+            if len(v) != 2:
+                raise ValueError('command tuples must have exactly 2 elements')
+
+            # It is common for extensions to wrap wire protocol commands via
+            # e.g. ``wireproto.commands[x] = (newfn, args)``. Because callers
+            # doing this aren't aware of the new API that uses objects to store
+            # command entries, we automatically merge old state with new.
+            if k in self:
+                v = self[k]._merge(v[0], v[1])
+            else:
+                # Use default values from @wireprotocommand.
+                v = commandentry(v[0], args=v[1],
+                                 transports=set(TRANSPORTS),
+                                 permission='push')
+        else:
+            raise ValueError('command entries must be commandentry instances '
+                             'or 2-tuples')
+
+        return super(commanddict, self).__setitem__(k, v)
+
+    def commandavailable(self, command, proto):
+        """Determine if a command is available for the requested protocol."""
+        assert proto.name in TRANSPORTS
+
+        entry = self.get(command)
+
+        if not entry:
+            return False
+
+        if proto.name not in entry.transports:
+            return False
+
+        return True

mercurial/wireprotov2server.py

@@ -1,522 +1,522
 # 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.
 
 from __future__ import absolute_import
 
 import contextlib
 
 from .i18n import _
 from .thirdparty import (
     cbor,
 )
 from .thirdparty.zope import (
     interface as zi,
 )
 from . import (
     encoding,
     error,
     pycompat,
     streamclone,
     util,
     wireproto,
     wireprotoframing,
     wireprototypes,
 )
 
 FRAMINGTYPE = b'application/mercurial-exp-framing-0005'
 
 HTTP_WIREPROTO_V2 = wireprototypes.HTTP_WIREPROTO_V2
 
 def handlehttpv2request(rctx, req, res, checkperm, urlparts):
     from .hgweb import common as hgwebcommon
 
     # URL space looks like: <permissions>/<command>, where <permission> can
     # be ``ro`` or ``rw`` to signal read-only or read-write, respectively.
 
     # Root URL does nothing meaningful... yet.
     if not urlparts:
         res.status = b'200 OK'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('HTTP version 2 API handler'))
         return
 
     if len(urlparts) == 1:
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('do not know how to process %s\n') %
                          req.dispatchpath)
         return
 
     permission, command = urlparts[0:2]
 
     if permission not in (b'ro', b'rw'):
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('unknown permission: %s') % permission)
         return
 
     if req.method != 'POST':
         res.status = b'405 Method Not Allowed'
         res.headers[b'Allow'] = b'POST'
         res.setbodybytes(_('commands require POST requests'))
         return
 
     # At some point we'll want to use our own API instead of recycling the
     # behavior of version 1 of the wire protocol...
     # TODO return reasonable responses - not responses that overload the
     # HTTP status line message for error reporting.
     try:
         checkperm(rctx, req, 'pull' if permission == b'ro' else 'push')
     except hgwebcommon.ErrorResponse as e:
         res.status = hgwebcommon.statusmessage(e.code, pycompat.bytestr(e))
         for k, v in e.headers:
             res.headers[k] = v
         res.setbodybytes('permission denied')
         return
 
     # We have a special endpoint to reflect the request back at the client.
     if command == b'debugreflect':
         _processhttpv2reflectrequest(rctx.repo.ui, rctx.repo, req, res)
         return
 
     # Extra commands that we handle that aren't really wire protocol
     # commands. Think extra hard before making this hackery available to
     # extension.
     extracommands = {'multirequest'}
 
     if command not in wireproto.commandsv2 and command not in extracommands:
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('unknown wire protocol command: %s\n') % command)
         return
 
     repo = rctx.repo
     ui = repo.ui
 
     proto = httpv2protocolhandler(req, ui)
 
     if (not wireproto.commandsv2.commandavailable(command, proto)
         and command not in extracommands):
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('invalid wire protocol command: %s') % command)
         return
 
     # TODO consider cases where proxies may add additional Accept headers.
     if req.headers.get(b'Accept') != FRAMINGTYPE:
         res.status = b'406 Not Acceptable'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('client MUST specify Accept header with value: %s\n')
                            % FRAMINGTYPE)
         return
 
     if req.headers.get(b'Content-Type') != FRAMINGTYPE:
         res.status = b'415 Unsupported Media Type'
         # TODO we should send a response with appropriate media type,
         # since client does Accept it.
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('client MUST send Content-Type header with '
                            'value: %s\n') % FRAMINGTYPE)
         return
 
     _processhttpv2request(ui, repo, req, res, permission, command, proto)
 
 def _processhttpv2reflectrequest(ui, repo, req, res):
     """Reads unified frame protocol request and dumps out state to client.
 
     This special endpoint can be used to help debug the wire protocol.
 
     Instead of routing the request through the normal dispatch mechanism,
     we instead read all frames, decode them, and feed them into our state
     tracker. We then dump the log of all that activity back out to the
     client.
     """
     import json
 
     # Reflection APIs have a history of being abused, accidentally disclosing
     # sensitive data, etc. So we have a config knob.
     if not ui.configbool('experimental', 'web.api.debugreflect'):
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('debugreflect service not available'))
         return
 
     # We assume we have a unified framing protocol request body.
 
     reactor = wireprotoframing.serverreactor()
     states = []
 
     while True:
         frame = wireprotoframing.readframe(req.bodyfh)
 
         if not frame:
             states.append(b'received: <no frame>')
             break
 
         states.append(b'received: %d %d %d %s' % (frame.typeid, frame.flags,
                                                   frame.requestid,
                                                   frame.payload))
 
         action, meta = reactor.onframerecv(frame)
         states.append(json.dumps((action, meta), sort_keys=True,
                                  separators=(', ', ': ')))
 
     action, meta = reactor.oninputeof()
     meta['action'] = action
     states.append(json.dumps(meta, sort_keys=True, separators=(', ',': ')))
 
     res.status = b'200 OK'
     res.headers[b'Content-Type'] = b'text/plain'
     res.setbodybytes(b'\n'.join(states))
 
 def _processhttpv2request(ui, repo, req, res, authedperm, reqcommand, proto):
     """Post-validation handler for HTTPv2 requests.
 
     Called when the HTTP request contains unified frame-based protocol
     frames for evaluation.
     """
     # TODO Some HTTP clients are full duplex and can receive data before
     # the entire request is transmitted. Figure out a way to indicate support
     # for that so we can opt into full duplex mode.
     reactor = wireprotoframing.serverreactor(deferoutput=True)
     seencommand = False
 
     outstream = reactor.makeoutputstream()
 
     while True:
         frame = wireprotoframing.readframe(req.bodyfh)
         if not frame:
             break
 
         action, meta = reactor.onframerecv(frame)
 
         if action == 'wantframe':
             # Need more data before we can do anything.
             continue
         elif action == 'runcommand':
             sentoutput = _httpv2runcommand(ui, repo, req, res, authedperm,
                                            reqcommand, reactor, outstream,
                                            meta, issubsequent=seencommand)
 
             if sentoutput:
                 return
 
             seencommand = True
 
         elif action == 'error':
             # TODO define proper error mechanism.
             res.status = b'200 OK'
             res.headers[b'Content-Type'] = b'text/plain'
             res.setbodybytes(meta['message'] + b'\n')
             return
         else:
             raise error.ProgrammingError(
                 'unhandled action from frame processor: %s' % action)
 
     action, meta = reactor.oninputeof()
     if action == 'sendframes':
         # We assume we haven't started sending the response yet. If we're
         # wrong, the response type will raise an exception.
         res.status = b'200 OK'
         res.headers[b'Content-Type'] = FRAMINGTYPE
         res.setbodygen(meta['framegen'])
     elif action == 'noop':
         pass
     else:
         raise error.ProgrammingError('unhandled action from frame processor: %s'
                                      % action)
 
 def _httpv2runcommand(ui, repo, req, res, authedperm, reqcommand, reactor,
                       outstream, command, issubsequent):
     """Dispatch a wire protocol command made from HTTPv2 requests.
 
     The authenticated permission (``authedperm``) along with the original
     command from the URL (``reqcommand``) are passed in.
     """
     # We already validated that the session has permissions to perform the
     # actions in ``authedperm``. In the unified frame protocol, the canonical
     # command to run is expressed in a frame. However, the URL also requested
     # to run a specific command. We need to be careful that the command we
     # run doesn't have permissions requirements greater than what was granted
     # by ``authedperm``.
     #
     # Our rule for this is we only allow one command per HTTP request and
     # that command must match the command in the URL. However, we make
     # an exception for the ``multirequest`` URL. This URL is allowed to
     # execute multiple commands. We double check permissions of each command
     # as it is invoked to ensure there is no privilege escalation.
     # TODO consider allowing multiple commands to regular command URLs
     # iff each command is the same.
 
     proto = httpv2protocolhandler(req, ui, args=command['args'])
 
     if reqcommand == b'multirequest':
         if not wireproto.commandsv2.commandavailable(command['command'], proto):
             # TODO proper error mechanism
             res.status = b'200 OK'
             res.headers[b'Content-Type'] = b'text/plain'
             res.setbodybytes(_('wire protocol command not available: %s') %
                              command['command'])
             return True
 
         # TODO don't use assert here, since it may be elided by -O.
         assert authedperm in (b'ro', b'rw')
         wirecommand = wireproto.commandsv2[command['command']]
         assert wirecommand.permission in ('push', 'pull')
 
         if authedperm == b'ro' and wirecommand.permission != 'pull':
             # TODO proper error mechanism
             res.status = b'403 Forbidden'
             res.headers[b'Content-Type'] = b'text/plain'
             res.setbodybytes(_('insufficient permissions to execute '
                                'command: %s') % command['command'])
             return True
 
         # TODO should we also call checkperm() here? Maybe not if we're going
         # to overhaul that API. The granted scope from the URL check should
         # be good enough.
 
     else:
         # Don't allow multiple commands outside of ``multirequest`` URL.
         if issubsequent:
             # TODO proper error mechanism
             res.status = b'200 OK'
             res.headers[b'Content-Type'] = b'text/plain'
             res.setbodybytes(_('multiple commands cannot be issued to this '
                                'URL'))
             return True
 
         if reqcommand != command['command']:
             # TODO define proper error mechanism
             res.status = b'200 OK'
             res.headers[b'Content-Type'] = b'text/plain'
             res.setbodybytes(_('command in frame must match command in URL'))
             return True
 
     rsp = wireproto.dispatch(repo, proto, command['command'])
 
     res.status = b'200 OK'
     res.headers[b'Content-Type'] = FRAMINGTYPE
 
     if isinstance(rsp, wireprototypes.cborresponse):
         encoded = cbor.dumps(rsp.value, canonical=True)
         action, meta = reactor.oncommandresponseready(outstream,
                                                       command['requestid'],
                                                       encoded)
     elif isinstance(rsp, wireprototypes.v2streamingresponse):
         action, meta = reactor.oncommandresponsereadygen(outstream,
                                                          command['requestid'],
                                                          rsp.gen)
     elif isinstance(rsp, wireprototypes.v2errorresponse):
         action, meta = reactor.oncommanderror(outstream,
                                               command['requestid'],
                                               rsp.message,
                                               rsp.args)
     else:
         action, meta = reactor.onservererror(
             _('unhandled response type from wire proto command'))
 
     if action == 'sendframes':
         res.setbodygen(meta['framegen'])
         return True
     elif action == 'noop':
         return False
     else:
         raise error.ProgrammingError('unhandled event from reactor: %s' %
                                      action)
 
 @zi.implementer(wireprototypes.baseprotocolhandler)
 class httpv2protocolhandler(object):
     def __init__(self, req, ui, args=None):
         self._req = req
         self._ui = ui
         self._args = args
 
     @property
     def name(self):
         return HTTP_WIREPROTO_V2
 
     def getargs(self, args):
         data = {}
         for k, typ in args.items():
             if k == '*':
                 raise NotImplementedError('do not support * args')
             elif k in self._args:
                 # TODO consider validating value types.
                 data[k] = self._args[k]
 
         return data
 
     def getprotocaps(self):
         # Protocol capabilities are currently not implemented for HTTP V2.
         return set()
 
     def getpayload(self):
         raise NotImplementedError
 
     @contextlib.contextmanager
     def mayberedirectstdio(self):
         raise NotImplementedError
 
     def client(self):
         raise NotImplementedError
 
     def addcapabilities(self, repo, caps):
         return caps
 
     def checkperm(self, perm):
         raise NotImplementedError
 
 def httpv2apidescriptor(req, repo):
     proto = httpv2protocolhandler(req, repo.ui)
 
     return _capabilitiesv2(repo, proto)
 
 def _capabilitiesv2(repo, proto):
     """Obtain the set of capabilities for version 2 transports.
 
     These capabilities are distinct from the capabilities for version 1
     transports.
     """
     compression = []
     for engine in wireproto.supportedcompengines(repo.ui, util.SERVERROLE):
         compression.append({
             b'name': engine.wireprotosupport().name,
         })
 
     caps = {
         'commands': {},
         'compression': compression,
         'framingmediatypes': [FRAMINGTYPE],
     }
 
     for command, entry in wireproto.commandsv2.items():
         caps['commands'][command] = {
             'args': entry.args,
             'permissions': [entry.permission],
         }
 
     if streamclone.allowservergeneration(repo):
         caps['rawrepoformats'] = sorted(repo.requirements &
                                         repo.supportedformats)
 
     return proto.addcapabilities(repo, caps)
 
 def wireprotocommand(name, args=None, permission='push'):
     """Decorator to declare a wire protocol command.
 
     ``name`` is the name of the wire protocol command being provided.
 
     ``args`` is a dict of argument names to example values.
 
     ``permission`` defines the permission type needed to run this command.
     Can be ``push`` or ``pull``. These roughly map to read-write and read-only,
     respectively. Default is to assume command requires ``push`` permissions
     because otherwise commands not declaring their permissions could modify
     a repository that is supposed to be read-only.
     """
     transports = {k for k, v in wireprototypes.TRANSPORTS.items()
                   if v['version'] == 2}
 
     if permission not in ('push', 'pull'):
         raise error.ProgrammingError('invalid wire protocol permission; '
                                      'got %s; expected "push" or "pull"' %
                                      permission)
 
     if args is None:
         args = {}
 
     if not isinstance(args, dict):
         raise error.ProgrammingError('arguments for version 2 commands '
                                      'must be declared as dicts')
 
     def register(func):
         if name in wireproto.commandsv2:
             raise error.ProgrammingError('%s command already registered '
                                          'for version 2' % name)
 
-        wireproto.commandsv2[name] = wireproto.commandentry(
+        wireproto.commandsv2[name] = wireprototypes.commandentry(
             func, args=args, transports=transports, permission=permission)
 
         return func
 
     return register
 
 @wireprotocommand('branchmap', permission='pull')
 def branchmapv2(repo, proto):
     branchmap = {encoding.fromlocal(k): v
                  for k, v in repo.branchmap().iteritems()}
 
     return wireprototypes.cborresponse(branchmap)
 
 @wireprotocommand('capabilities', permission='pull')
 def capabilitiesv2(repo, proto):
     caps = _capabilitiesv2(repo, proto)
 
     return wireprototypes.cborresponse(caps)
 
 @wireprotocommand('heads',
                   args={
                       'publiconly': False,
                   },
                   permission='pull')
 def headsv2(repo, proto, publiconly=False):
     if publiconly:
         repo = repo.filtered('immutable')
 
     return wireprototypes.cborresponse(repo.heads())
 
 @wireprotocommand('known',
                   args={
                       'nodes': [b'deadbeef'],
                   },
                   permission='pull')
 def knownv2(repo, proto, nodes=None):
     nodes = nodes or []
     result = b''.join(b'1' if n else b'0' for n in repo.known(nodes))
     return wireprototypes.cborresponse(result)
 
 @wireprotocommand('listkeys',
                   args={
                       'namespace': b'ns',
                   },
                   permission='pull')
 def listkeysv2(repo, proto, namespace=None):
     keys = repo.listkeys(encoding.tolocal(namespace))
     keys = {encoding.fromlocal(k): encoding.fromlocal(v)
             for k, v in keys.iteritems()}
 
     return wireprototypes.cborresponse(keys)
 
 @wireprotocommand('lookup',
                   args={
                       'key': b'foo',
                   },
                   permission='pull')
 def lookupv2(repo, proto, key):
     key = encoding.tolocal(key)
 
     # TODO handle exception.
     node = repo.lookup(key)
 
     return wireprototypes.cborresponse(node)
 
 @wireprotocommand('pushkey',
                   args={
                       'namespace': b'ns',
                       'key': b'key',
                       'old': b'old',
                       'new': b'new',
                   },
                   permission='push')
 def pushkeyv2(repo, proto, namespace, key, old, new):
     # TODO handle ui output redirection
     r = repo.pushkey(encoding.tolocal(namespace),
                      encoding.tolocal(key),
                      encoding.tolocal(old),
                      encoding.tolocal(new))
 
     return wireprototypes.cborresponse(r)