wireproto: separate commands tables for version 1 and 2 commands

We can't easily reuse existing command handlers for version 2
commands because the response types will be different. e.g. many
commands return nodes encoded as hex. Our new wire protocol is
binary safe, so we'll wish to encode nodes as binary.

We /could/ teach each command handler to look at the protocol
handler and change behavior based on the version in use. However,
this would make logic a bit unwieldy over time and would make
it harder to design a unified protocol handler interface. I think
it's better to create a clean break between version 1 and version 2
of commands on the server.

What I imagine happening is we will have separate @wireprotocommand
functions for each protocol generation. Those functions will parse the
request, dispatch to a common function to process it, then generate
the response in its own, transport-specific manner.

This commit establishes a separate table for tracking version 1
commands from version 2 commands. The HTTP server pieces have been
updated to use this new table.

Most commands are marked as both version 1 and version 2, so there is
little practical impact to this change.

A side-effect of this change is we now rely on transport registration
in wireprototypes.TRANSPORTS and certain properties of the protocol
interface. So a test had to be updated to conform.

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

hgext/largefiles/uisetup.py

@@ -1,202 +1,203 b''
 # Copyright 2009-2010 Gregory P. Ward
 # Copyright 2009-2010 Intelerad Medical Systems Incorporated
 # Copyright 2010-2011 Fog Creek Software
 # Copyright 2010-2011 Unity Technologies
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 '''setup for largefiles extension: uisetup'''
 from __future__ import absolute_import
 
 from mercurial.i18n import _
 
 from mercurial.hgweb import (
     webcommands,
 )
 
 from mercurial import (
     archival,
     cmdutil,
     commands,
     copies,
     exchange,
     extensions,
     filemerge,
     hg,
     httppeer,
     merge,
     scmutil,
     sshpeer,
     subrepo,
     upgrade,
     url,
     wireproto,
 )
 
 from . import (
     overrides,
     proto,
 )
 
 def uisetup(ui):
     # Disable auto-status for some commands which assume that all
     # files in the result are under Mercurial's control
 
     entry = extensions.wrapcommand(commands.table, 'add',
                                    overrides.overrideadd)
     addopt = [('', 'large', None, _('add as largefile')),
               ('', 'normal', None, _('add as normal file')),
               ('', 'lfsize', '', _('add all files above this size '
                                    '(in megabytes) as largefiles '
                                    '(default: 10)'))]
     entry[1].extend(addopt)
 
     # The scmutil function is called both by the (trivial) addremove command,
     # and in the process of handling commit -A (issue3542)
     extensions.wrapfunction(scmutil, 'addremove', overrides.scmutiladdremove)
     extensions.wrapfunction(cmdutil, 'add', overrides.cmdutiladd)
     extensions.wrapfunction(cmdutil, 'remove', overrides.cmdutilremove)
     extensions.wrapfunction(cmdutil, 'forget', overrides.cmdutilforget)
 
     extensions.wrapfunction(copies, 'pathcopies', overrides.copiespathcopies)
 
     extensions.wrapfunction(upgrade, 'preservedrequirements',
                             overrides.upgraderequirements)
 
     extensions.wrapfunction(upgrade, 'supporteddestrequirements',
                             overrides.upgraderequirements)
 
     # Subrepos call status function
     entry = extensions.wrapcommand(commands.table, 'status',
                                    overrides.overridestatus)
     extensions.wrapfunction(subrepo.hgsubrepo, 'status',
                             overrides.overridestatusfn)
 
     entry = extensions.wrapcommand(commands.table, 'log',
                                    overrides.overridelog)
     entry = extensions.wrapcommand(commands.table, 'rollback',
                                    overrides.overriderollback)
     entry = extensions.wrapcommand(commands.table, 'verify',
                                    overrides.overrideverify)
 
     verifyopt = [('', 'large', None,
                   _('verify that all largefiles in current revision exists')),
                  ('', 'lfa', None,
                   _('verify largefiles in all revisions, not just current')),
                  ('', 'lfc', None,
                   _('verify local largefile contents, not just existence'))]
     entry[1].extend(verifyopt)
 
     entry = extensions.wrapcommand(commands.table, 'debugstate',
                                    overrides.overridedebugstate)
     debugstateopt = [('', 'large', None, _('display largefiles dirstate'))]
     entry[1].extend(debugstateopt)
 
     outgoing = lambda orgfunc, *arg, **kwargs: orgfunc(*arg, **kwargs)
     entry = extensions.wrapcommand(commands.table, 'outgoing', outgoing)
     outgoingopt = [('', 'large', None, _('display outgoing largefiles'))]
     entry[1].extend(outgoingopt)
     cmdutil.outgoinghooks.add('largefiles', overrides.outgoinghook)
     entry = extensions.wrapcommand(commands.table, 'summary',
                                    overrides.overridesummary)
     summaryopt = [('', 'large', None, _('display outgoing largefiles'))]
     entry[1].extend(summaryopt)
     cmdutil.summaryremotehooks.add('largefiles', overrides.summaryremotehook)
 
     entry = extensions.wrapcommand(commands.table, 'pull',
                                    overrides.overridepull)
     pullopt = [('', 'all-largefiles', None,
                  _('download all pulled versions of largefiles (DEPRECATED)')),
                ('', 'lfrev', [],
                 _('download largefiles for these revisions'), _('REV'))]
     entry[1].extend(pullopt)
 
     entry = extensions.wrapcommand(commands.table, 'push',
                                    overrides.overridepush)
     pushopt = [('', 'lfrev', [],
                 _('upload largefiles for these revisions'), _('REV'))]
     entry[1].extend(pushopt)
     extensions.wrapfunction(exchange, 'pushoperation',
                             overrides.exchangepushoperation)
 
     entry = extensions.wrapcommand(commands.table, 'clone',
                                    overrides.overrideclone)
     cloneopt = [('', 'all-largefiles', None,
                  _('download all versions of all largefiles'))]
     entry[1].extend(cloneopt)
     extensions.wrapfunction(hg, 'clone', overrides.hgclone)
     extensions.wrapfunction(hg, 'postshare', overrides.hgpostshare)
 
     entry = extensions.wrapcommand(commands.table, 'cat',
                                    overrides.overridecat)
     extensions.wrapfunction(merge, '_checkunknownfile',
                             overrides.overridecheckunknownfile)
     extensions.wrapfunction(merge, 'calculateupdates',
                             overrides.overridecalculateupdates)
     extensions.wrapfunction(merge, 'recordupdates',
                             overrides.mergerecordupdates)
     extensions.wrapfunction(merge, 'update', overrides.mergeupdate)
     extensions.wrapfunction(filemerge, '_filemerge',
                             overrides.overridefilemerge)
     extensions.wrapfunction(cmdutil, 'copy', overrides.overridecopy)
 
     # Summary calls dirty on the subrepos
     extensions.wrapfunction(subrepo.hgsubrepo, 'dirty', overrides.overridedirty)
 
     extensions.wrapfunction(cmdutil, 'revert', overrides.overriderevert)
 
     extensions.wrapcommand(commands.table, 'archive',
                            overrides.overridearchivecmd)
     extensions.wrapfunction(archival, 'archive', overrides.overridearchive)
     extensions.wrapfunction(subrepo.hgsubrepo, 'archive',
                             overrides.hgsubrepoarchive)
     extensions.wrapfunction(webcommands, 'archive', overrides.hgwebarchive)
     extensions.wrapfunction(cmdutil, 'bailifchanged',
                             overrides.overridebailifchanged)
 
     extensions.wrapfunction(cmdutil, 'postcommitstatus',
                             overrides.postcommitstatus)
     extensions.wrapfunction(scmutil, 'marktouched',
                             overrides.scmutilmarktouched)
 
     extensions.wrapfunction(url, 'open',
                             overrides.openlargefile)
 
     # create the new wireproto commands ...
     wireproto.wireprotocommand('putlfile', 'sha', permission='push')(
         proto.putlfile)
     wireproto.wireprotocommand('getlfile', 'sha', permission='pull')(
         proto.getlfile)
     wireproto.wireprotocommand('statlfile', 'sha', permission='pull')(
         proto.statlfile)
     wireproto.wireprotocommand('lheads', '', permission='pull')(
         wireproto.heads)
 
     # ... and wrap some existing ones
     wireproto.commands['heads'].func = proto.heads
+    # TODO also wrap wireproto.commandsv2 once heads is implemented there.
 
     extensions.wrapfunction(webcommands, 'decodepath', overrides.decodepath)
 
     extensions.wrapfunction(wireproto, '_capabilities', proto._capabilities)
 
     # can't do this in reposetup because it needs to have happened before
     # wirerepo.__init__ is called
     proto.ssholdcallstream = sshpeer.sshv1peer._callstream
     proto.httpoldcallstream = httppeer.httppeer._callstream
     sshpeer.sshv1peer._callstream = proto.sshrepocallstream
     httppeer.httppeer._callstream = proto.httprepocallstream
 
     # override some extensions' stuff as well
     for name, module in extensions.extensions():
         if name == 'purge':
             extensions.wrapcommand(getattr(module, 'cmdtable'), 'purge',
                 overrides.overridepurge)
         if name == 'rebase':
             extensions.wrapcommand(getattr(module, 'cmdtable'), 'rebase',
                 overrides.overriderebase)
             extensions.wrapfunction(module, 'rebase',
                                     overrides.overriderebase)
         if name == 'transplant':
             extensions.wrapcommand(getattr(module, 'cmdtable'), 'transplant',
                 overrides.overridetransplant)

mercurial/wireproto.py

@@ -1,1139 +1,1163 b''
 # 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 hashlib
 import os
 import tempfile
 
 from .i18n import _
 from .node import (
     bin,
     hex,
     nullid,
 )
 
 from . import (
     bundle2,
     changegroup as changegroupmod,
     discovery,
     encoding,
     error,
     exchange,
     peer,
     pushkey as pushkeymod,
     pycompat,
     repository,
     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)
 
 class remoteiterbatcher(peer.iterbatcher):
     def __init__(self, remote):
         super(remoteiterbatcher, self).__init__()
         self._remote = remote
 
     def __getattr__(self, name):
         # Validate this method is batchable, since submit() only supports
         # batchable methods.
         fn = getattr(self._remote, name)
         if not getattr(fn, 'batchable', None):
             raise error.ProgrammingError('Attempted to batch a non-batchable '
                                          'call to %r' % name)
 
         return super(remoteiterbatcher, self).__getattr__(name)
 
     def submit(self):
         """Break the batch request into many patch calls and pipeline them.
 
         This is mostly valuable over http where request sizes can be
         limited, but can be used in other places as well.
         """
         # 2-tuple of (command, arguments) that represents what will be
         # sent over the wire.
         requests = []
 
         # 4-tuple of (command, final future, @batchable generator, remote
         # future).
         results = []
 
         for command, args, opts, finalfuture in self.calls:
             mtd = getattr(self._remote, command)
             batchable = mtd.batchable(mtd.__self__, *args, **opts)
 
             commandargs, fremote = next(batchable)
             assert fremote
             requests.append((command, commandargs))
             results.append((command, finalfuture, batchable, fremote))
 
         if requests:
             self._resultiter = self._remote._submitbatch(requests)
 
         self._results = results
 
     def results(self):
         for command, finalfuture, batchable, remotefuture in self._results:
             # Get the raw result, set it in the remote future, feed it
             # back into the @batchable generator so it can be decoded, and
             # set the result on the final future to this value.
             remoteresult = next(self._resultiter)
             remotefuture.set(remoteresult)
             finalfuture.set(next(batchable))
 
             # Verify our @batchable generators only emit 2 values.
             try:
                 next(batchable)
             except StopIteration:
                 pass
             else:
                 raise error.ProgrammingError('%s @batchable generator emitted '
                                              'unexpected value count' % command)
 
             yield finalfuture.value
 
 # Forward a couple of names from peer to make wireproto interactions
 # slightly more sensible.
 batchable = peer.batchable
 future = peer.future
 
 # 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 escapearg(plain):
     return (plain
             .replace(':', ':c')
             .replace(',', ':o')
             .replace(';', ':s')
             .replace('=', ':e'))
 
 def unescapearg(escaped):
     return (escaped
             .replace(':e', '=')
             .replace(':s', ';')
             .replace(':o', ',')
             .replace(':c', ':'))
 
 def encodebatchcmds(req):
     """Return a ``cmds`` argument value for the ``batch`` command."""
     cmds = []
     for op, argsdict in req:
         # Old servers didn't properly unescape argument names. So prevent
         # the sending of argument names that may not be decoded properly by
         # servers.
         assert all(escapearg(k) == k for k in argsdict)
 
         args = ','.join('%s=%s' % (escapearg(k), escapearg(v))
                         for k, v in argsdict.iteritems())
         cmds.append('%s %s' % (op, args))
 
     return ';'.join(cmds)
 
 # 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.
 gboptsmap = {'heads':  'nodes',
              'bookmarks': 'boolean',
              'common': 'nodes',
              'obsmarkers': 'boolean',
              'phases': 'boolean',
              'bundlecaps': 'scsv',
              'listkeys': 'csv',
              'cg': 'boolean',
              'cbattempted': 'boolean',
              'stream': 'boolean',
 }
 
 # client side
 
 class wirepeer(repository.legacypeer):
     """Client-side interface for communicating with a peer repository.
 
     Methods commonly call wire protocol commands of the same name.
 
     See also httppeer.py and sshpeer.py for protocol-specific
     implementations of this interface.
     """
     # Begin of basewirepeer interface.
 
     def iterbatch(self):
         return remoteiterbatcher(self)
 
     @batchable
     def lookup(self, key):
         self.requirecap('lookup', _('look up remote revision'))
         f = future()
         yield {'key': encoding.fromlocal(key)}, f
         d = f.value
         success, data = d[:-1].split(" ", 1)
         if int(success):
             yield bin(data)
         else:
             self._abort(error.RepoError(data))
 
     @batchable
     def heads(self):
         f = future()
         yield {}, f
         d = f.value
         try:
             yield decodelist(d[:-1])
         except ValueError:
             self._abort(error.ResponseError(_("unexpected response:"), d))
 
     @batchable
     def known(self, nodes):
         f = future()
         yield {'nodes': encodelist(nodes)}, f
         d = f.value
         try:
             yield [bool(int(b)) for b in d]
         except ValueError:
             self._abort(error.ResponseError(_("unexpected response:"), d))
 
     @batchable
     def branchmap(self):
         f = future()
         yield {}, f
         d = f.value
         try:
             branchmap = {}
             for branchpart in d.splitlines():
                 branchname, branchheads = branchpart.split(' ', 1)
                 branchname = encoding.tolocal(urlreq.unquote(branchname))
                 branchheads = decodelist(branchheads)
                 branchmap[branchname] = branchheads
             yield branchmap
         except TypeError:
             self._abort(error.ResponseError(_("unexpected response:"), d))
 
     @batchable
     def listkeys(self, namespace):
         if not self.capable('pushkey'):
             yield {}, None
         f = future()
         self.ui.debug('preparing listkeys for "%s"\n' % namespace)
         yield {'namespace': encoding.fromlocal(namespace)}, f
         d = f.value
         self.ui.debug('received listkey for "%s": %i bytes\n'
                       % (namespace, len(d)))
         yield pushkeymod.decodekeys(d)
 
     @batchable
     def pushkey(self, namespace, key, old, new):
         if not self.capable('pushkey'):
             yield False, None
         f = future()
         self.ui.debug('preparing pushkey for "%s:%s"\n' % (namespace, key))
         yield {'namespace': encoding.fromlocal(namespace),
                'key': encoding.fromlocal(key),
                'old': encoding.fromlocal(old),
                'new': encoding.fromlocal(new)}, f
         d = f.value
         d, output = d.split('\n', 1)
         try:
             d = bool(int(d))
         except ValueError:
             raise error.ResponseError(
                 _('push failed (unexpected response):'), d)
         for l in output.splitlines(True):
             self.ui.status(_('remote: '), l)
         yield d
 
     def stream_out(self):
         return self._callstream('stream_out')
 
     def getbundle(self, source, **kwargs):
         kwargs = pycompat.byteskwargs(kwargs)
         self.requirecap('getbundle', _('look up remote changes'))
         opts = {}
         bundlecaps = kwargs.get('bundlecaps')
         if bundlecaps is not None:
             kwargs['bundlecaps'] = sorted(bundlecaps)
         else:
             bundlecaps = () # kwargs could have it to None
         for key, value in kwargs.iteritems():
             if value is None:
                 continue
             keytype = gboptsmap.get(key)
             if keytype is None:
                 raise error.ProgrammingError(
                     'Unexpectedly None keytype for key %s' % key)
             elif keytype == 'nodes':
                 value = encodelist(value)
             elif keytype in ('csv', 'scsv'):
                 value = ','.join(value)
             elif keytype == 'boolean':
                 value = '%i' % bool(value)
             elif keytype != 'plain':
                 raise KeyError('unknown getbundle option type %s'
                                % keytype)
             opts[key] = value
         f = self._callcompressable("getbundle", **pycompat.strkwargs(opts))
         if any((cap.startswith('HG2') for cap in bundlecaps)):
             return bundle2.getunbundler(self.ui, f)
         else:
             return changegroupmod.cg1unpacker(f, 'UN')
 
     def unbundle(self, cg, heads, url):
         '''Send cg (a readable file-like object representing the
         changegroup to push, typically a chunkbuffer object) to the
         remote server as a bundle.
 
         When pushing a bundle10 stream, return an integer indicating the
         result of the push (see changegroup.apply()).
 
         When pushing a bundle20 stream, return a bundle20 stream.
 
         `url` is the url the client thinks it's pushing to, which is
         visible to hooks.
         '''
 
         if heads != ['force'] and self.capable('unbundlehash'):
             heads = encodelist(['hashed',
                                 hashlib.sha1(''.join(sorted(heads))).digest()])
         else:
             heads = encodelist(heads)
 
         if util.safehasattr(cg, 'deltaheader'):
             # this a bundle10, do the old style call sequence
             ret, output = self._callpush("unbundle", cg, heads=heads)
             if ret == "":
                 raise error.ResponseError(
                     _('push failed:'), output)
             try:
                 ret = int(ret)
             except ValueError:
                 raise error.ResponseError(
                     _('push failed (unexpected response):'), ret)
 
             for l in output.splitlines(True):
                 self.ui.status(_('remote: '), l)
         else:
             # bundle2 push. Send a stream, fetch a stream.
             stream = self._calltwowaystream('unbundle', cg, heads=heads)
             ret = bundle2.getunbundler(self.ui, stream)
         return ret
 
     # End of basewirepeer interface.
 
     # Begin of baselegacywirepeer interface.
 
     def branches(self, nodes):
         n = encodelist(nodes)
         d = self._call("branches", nodes=n)
         try:
             br = [tuple(decodelist(b)) for b in d.splitlines()]
             return br
         except ValueError:
             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([encodelist(p, '-') for p in pairs[i:i + batch]])
             d = self._call("between", pairs=n)
             try:
                 r.extend(l and decodelist(l) or [] for l in d.splitlines())
             except ValueError:
                 self._abort(error.ResponseError(_("unexpected response:"), d))
         return r
 
     def changegroup(self, nodes, kind):
         n = encodelist(nodes)
         f = self._callcompressable("changegroup", roots=n)
         return changegroupmod.cg1unpacker(f, 'UN')
 
     def changegroupsubset(self, bases, heads, kind):
         self.requirecap('changegroupsubset', _('look up remote changes'))
         bases = encodelist(bases)
         heads = encodelist(heads)
         f = self._callcompressable("changegroupsubset",
                                    bases=bases, heads=heads)
         return changegroupmod.cg1unpacker(f, 'UN')
 
     # End of baselegacywirepeer interface.
 
     def _submitbatch(self, req):
         """run batch request <req> on the server
 
         Returns an iterator of the raw responses from the server.
         """
         ui = self.ui
         if ui.debugflag and ui.configbool('devel', 'debug.peer-request'):
             ui.debug('devel-peer-request: batched-content\n')
             for op, args in req:
                 msg = 'devel-peer-request:    - %s (%d arguments)\n'
                 ui.debug(msg % (op, len(args)))
 
         rsp = self._callstream("batch", cmds=encodebatchcmds(req))
         chunk = rsp.read(1024)
         work = [chunk]
         while chunk:
             while ';' not in chunk and chunk:
                 chunk = rsp.read(1024)
                 work.append(chunk)
             merged = ''.join(work)
             while ';' in merged:
                 one, merged = merged.split(';', 1)
                 yield unescapearg(one)
             chunk = rsp.read(1024)
             work = [merged, chunk]
         yield unescapearg(''.join(work))
 
     def _submitone(self, op, args):
         return self._call(op, **pycompat.strkwargs(args))
 
     def debugwireargs(self, one, two, three=None, four=None, five=None):
         # don't pass optional arguments left at their default value
         opts = {}
         if three is not None:
             opts[r'three'] = three
         if four is not None:
             opts[r'four'] = four
         return self._call('debugwireargs', one=one, two=two, **opts)
 
     def _call(self, cmd, **args):
         """execute <cmd> on the server
 
         The command is expected to return a simple string.
 
         returns the server reply as a string."""
         raise NotImplementedError()
 
     def _callstream(self, cmd, **args):
         """execute <cmd> on the server
 
         The command is expected to return a stream. Note that if the
         command doesn't return a stream, _callstream behaves
         differently for ssh and http peers.
 
         returns the server reply as a file like object.
         """
         raise NotImplementedError()
 
     def _callcompressable(self, cmd, **args):
         """execute <cmd> on the server
 
         The command is expected to return a stream.
 
         The stream may have been compressed in some implementations. This
         function takes care of the decompression. This is the only difference
         with _callstream.
 
         returns the server reply as a file like object.
         """
         raise NotImplementedError()
 
     def _callpush(self, cmd, fp, **args):
         """execute a <cmd> on server
 
         The command is expected to be related to a push. Push has a special
         return method.
 
         returns the server reply as a (ret, output) tuple. ret is either
         empty (error) or a stringified int.
         """
         raise NotImplementedError()
 
     def _calltwowaystream(self, cmd, fp, **args):
         """execute <cmd> on server
 
         The command will send a stream to the server and get a stream in reply.
         """
         raise NotImplementedError()
 
     def _abort(self, exception):
         """clearly abort the wire protocol connection and raise the exception
         """
         raise NotImplementedError()
 
 # server side
 
 # 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)
-    func, spec = commands[command]
+
+    transportversion = wireprototypes.TRANSPORTS[proto.name]['version']
+    commandtable = commandsv2 if transportversion == 2 else commands
+    func, spec = commandtable[command]
+
     args = proto.getargs(spec)
     return func(repo, proto, *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
 
 # Constants specifying which transports a wire protocol command should be
 # available on. For use with @wireprotocommand.
 POLICY_ALL = 'all'
 POLICY_V1_ONLY = 'v1-only'
 POLICY_V2_ONLY = 'v2-only'
 
+# For version 1 transports.
 commands = commanddict()
 
+# For version 2 transports.
+commandsv2 = commanddict()
+
 def wireprotocommand(name, args='', transportpolicy=POLICY_ALL,
                      permission='push'):
     """Decorator to declare a wire protocol command.
 
     ``name`` is the name of the wire protocol command being provided.
 
     ``args`` is a space-delimited list of named arguments that the command
     accepts. ``*`` is a special value that says to accept all arguments.
 
     ``transportpolicy`` is a POLICY_* constant denoting which transports
     this wire protocol command should be exposed to. By default, commands
     are exposed to all wire protocol transports.
 
     ``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.
     """
     if transportpolicy == POLICY_ALL:
         transports = set(wireprototypes.TRANSPORTS)
+        transportversions = {1, 2}
     elif transportpolicy == POLICY_V1_ONLY:
         transports = {k for k, v in wireprototypes.TRANSPORTS.items()
                       if v['version'] == 1}
+        transportversions = {1}
     elif transportpolicy == POLICY_V2_ONLY:
         transports = {k for k, v in wireprototypes.TRANSPORTS.items()
                       if v['version'] == 2}
+        transportversions = {2}
     else:
         raise error.ProgrammingError('invalid transport policy value: %s' %
                                      transportpolicy)
 
     # 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)
 
     def register(func):
-        commands[name] = commandentry(func, args=args, transports=transports,
+        if 1 in transportversions:
+            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)
+        if 2 in transportversions:
+            if name in commandsv2:
+                raise error.ProgrammingError('%s command already registered '
+                                             'for version 2' % name)
+            commandsv2[name] = 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',
                   transportpolicy=POLICY_V1_ONLY)
 def batch(repo, proto, cmds, others):
     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(escapearg(result))
 
     return wireprototypes.bytesresponse(';'.join(res))
 
 @wireprotocommand('between', 'pairs', transportpolicy=POLICY_V1_ONLY,
                   permission='pull')
 def between(repo, proto, pairs):
     pairs = [decodelist(p, '-') for p in pairs.split(" ")]
     r = []
     for b in repo.between(pairs):
         r.append(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 = encodelist(nodes)
         heads.append('%s %s' % (branchname, branchnodes))
 
     return wireprototypes.bytesresponse('\n'.join(heads))
 
 @wireprotocommand('branches', 'nodes', transportpolicy=POLICY_V1_ONLY,
                   permission='pull')
 def branches(repo, proto, nodes):
     nodes = decodelist(nodes)
     r = []
     for b in repo.branches(nodes):
         r.append(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):
     return wireprototypes.bytesresponse(' '.join(_capabilities(repo, proto)))
 
 @wireprotocommand('changegroup', 'roots', transportpolicy=POLICY_V1_ONLY,
                   permission='pull')
 def changegroup(repo, proto, roots):
     nodes = 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',
                   transportpolicy=POLICY_V1_ONLY,
                   permission='pull')
 def changegroupsubset(repo, proto, bases, heads):
     bases = decodelist(bases)
     heads = 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)))
 
 @wireprotocommand('getbundle', '*', permission='pull')
 def getbundle(repo, proto, others):
     opts = options('getbundle', gboptsmap.keys(), others)
     for k, v in opts.iteritems():
         keytype = gboptsmap[k]
         if keytype == 'nodes':
             opts[k] = 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:
         if repo.ui.configbool('server', 'disablefullbundle'):
             # Check to see if this is a full clone.
             clheads = set(repo.changelog.heads())
             changegroup = opts.get('cg', True)
             heads = set(opts.get('heads', set()))
             common = set(opts.get('common', set()))
             common.discard(nullid)
             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(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)
         c = repo[k]
         r = c.hex()
         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(decodelist(nodes)))
     return wireprototypes.bytesresponse(v)
 
 @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 = decodelist(heads)
 
     with proto.mayberedirectstdio() as output:
         try:
             exchange.check_heads(repo, their_heads, 'preparing changes')
 
             # write bundle data to temporary file because it can be big
             fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
             fp = os.fdopen(fd, r'wb+')
             r = 0
             try:
                 proto.forwardpayload(fp)
                 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:
                 fp.close()
                 os.unlink(tempname)
 
         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/wireprotoserver.py

@@ -1,1050 +1,1050 b''
 # 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
 import struct
 import sys
 import threading
 
 from .i18n import _
 from . import (
     encoding,
     error,
     hook,
     pycompat,
     util,
     wireproto,
     wireprotoframing,
     wireprototypes,
 )
 from .utils import (
     procutil,
 )
 
 stringio = util.stringio
 
 urlerr = util.urlerr
 urlreq = util.urlreq
 
 HTTP_OK = 200
 
 HGTYPE = 'application/mercurial-0.1'
 HGTYPE2 = 'application/mercurial-0.2'
 HGERRTYPE = 'application/hg-error'
 FRAMINGTYPE = b'application/mercurial-exp-framing-0003'
 
 HTTPV2 = wireprototypes.HTTPV2
 SSHV1 = wireprototypes.SSHV1
 SSHV2 = wireprototypes.SSHV2
 
 def decodevaluefromheaders(req, headerprefix):
     """Decode a long value from multiple HTTP request headers.
 
     Returns the value as a bytes, not a str.
     """
     chunks = []
     i = 1
     while True:
         v = req.headers.get(b'%s-%d' % (headerprefix, i))
         if v is None:
             break
         chunks.append(pycompat.bytesurl(v))
         i += 1
 
     return ''.join(chunks)
 
 class httpv1protocolhandler(wireprototypes.baseprotocolhandler):
     def __init__(self, req, ui, checkperm):
         self._req = req
         self._ui = ui
         self._checkperm = checkperm
 
     @property
     def name(self):
         return 'http-v1'
 
     def getargs(self, args):
         knownargs = self._args()
         data = {}
         keys = args.split()
         for k in keys:
             if k == '*':
                 star = {}
                 for key in knownargs.keys():
                     if key != 'cmd' and key not in keys:
                         star[key] = knownargs[key][0]
                 data['*'] = star
             else:
                 data[k] = knownargs[k][0]
         return [data[k] for k in keys]
 
     def _args(self):
         args = self._req.qsparams.asdictoflists()
         postlen = int(self._req.headers.get(b'X-HgArgs-Post', 0))
         if postlen:
             args.update(urlreq.parseqs(
                 self._req.bodyfh.read(postlen), keep_blank_values=True))
             return args
 
         argvalue = decodevaluefromheaders(self._req, b'X-HgArg')
         args.update(urlreq.parseqs(argvalue, keep_blank_values=True))
         return args
 
     def forwardpayload(self, fp):
         # Existing clients *always* send Content-Length.
         length = int(self._req.headers[b'Content-Length'])
 
         # If httppostargs is used, we need to read Content-Length
         # minus the amount that was consumed by args.
         length -= int(self._req.headers.get(b'X-HgArgs-Post', 0))
         for s in util.filechunkiter(self._req.bodyfh, limit=length):
             fp.write(s)
 
     @contextlib.contextmanager
     def mayberedirectstdio(self):
         oldout = self._ui.fout
         olderr = self._ui.ferr
 
         out = util.stringio()
 
         try:
             self._ui.fout = out
             self._ui.ferr = out
             yield out
         finally:
             self._ui.fout = oldout
             self._ui.ferr = olderr
 
     def client(self):
         return 'remote:%s:%s:%s' % (
             self._req.urlscheme,
             urlreq.quote(self._req.remotehost or ''),
             urlreq.quote(self._req.remoteuser or ''))
 
     def addcapabilities(self, repo, caps):
         caps.append(b'batch')
 
         caps.append('httpheader=%d' %
                     repo.ui.configint('server', 'maxhttpheaderlen'))
         if repo.ui.configbool('experimental', 'httppostargs'):
             caps.append('httppostargs')
 
         # FUTURE advertise 0.2rx once support is implemented
         # FUTURE advertise minrx and mintx after consulting config option
         caps.append('httpmediatype=0.1rx,0.1tx,0.2tx')
 
         compengines = wireproto.supportedcompengines(repo.ui, util.SERVERROLE)
         if compengines:
             comptypes = ','.join(urlreq.quote(e.wireprotosupport().name)
                                  for e in compengines)
             caps.append('compression=%s' % comptypes)
 
         return caps
 
     def checkperm(self, perm):
         return self._checkperm(perm)
 
 # This method exists mostly so that extensions like remotefilelog can
 # disable a kludgey legacy method only over http. As of early 2018,
 # there are no other known users, so with any luck we can discard this
 # hook if remotefilelog becomes a first-party extension.
 def iscmd(cmd):
     return cmd in wireproto.commands
 
 def handlewsgirequest(rctx, req, res, checkperm):
     """Possibly process a wire protocol request.
 
     If the current request is a wire protocol request, the request is
     processed by this function.
 
     ``req`` is a ``parsedrequest`` instance.
     ``res`` is a ``wsgiresponse`` instance.
 
     Returns a bool indicating if the request was serviced. If set, the caller
     should stop processing the request, as a response has already been issued.
     """
     # Avoid cycle involving hg module.
     from .hgweb import common as hgwebcommon
 
     repo = rctx.repo
 
     # HTTP version 1 wire protocol requests are denoted by a "cmd" query
     # string parameter. If it isn't present, this isn't a wire protocol
     # request.
     if 'cmd' not in req.qsparams:
         return False
 
     cmd = req.qsparams['cmd']
 
     # The "cmd" request parameter is used by both the wire protocol and hgweb.
     # While not all wire protocol commands are available for all transports,
     # if we see a "cmd" value that resembles a known wire protocol command, we
     # route it to a protocol handler. This is better than routing possible
     # wire protocol requests to hgweb because it prevents hgweb from using
     # known wire protocol commands and it is less confusing for machine
     # clients.
     if not iscmd(cmd):
         return False
 
     # The "cmd" query string argument is only valid on the root path of the
     # repo. e.g. ``/?cmd=foo``, ``/repo?cmd=foo``. URL paths within the repo
     # like ``/blah?cmd=foo`` are not allowed. So don't recognize the request
     # in this case. We send an HTTP 404 for backwards compatibility reasons.
     if req.dispatchpath:
         res.status = hgwebcommon.statusmessage(404)
         res.headers['Content-Type'] = HGTYPE
         # TODO This is not a good response to issue for this request. This
         # is mostly for BC for now.
         res.setbodybytes('0\n%s\n' % b'Not Found')
         return True
 
     proto = httpv1protocolhandler(req, repo.ui,
                                   lambda perm: checkperm(rctx, req, perm))
 
     # The permissions checker should be the only thing that can raise an
     # ErrorResponse. It is kind of a layer violation to catch an hgweb
     # exception here. So consider refactoring into a exception type that
     # is associated with the wire protocol.
     try:
         _callhttp(repo, req, res, proto, cmd)
     except hgwebcommon.ErrorResponse as e:
         for k, v in e.headers:
             res.headers[k] = v
         res.status = hgwebcommon.statusmessage(e.code, pycompat.bytestr(e))
         # TODO This response body assumes the failed command was
         # "unbundle." That assumption is not always valid.
         res.setbodybytes('0\n%s\n' % pycompat.bytestr(e))
 
     return True
 
 def handlewsgiapirequest(rctx, req, res, checkperm):
     """Handle requests to /api/*."""
     assert req.dispatchparts[0] == b'api'
 
     repo = rctx.repo
 
     # This whole URL space is experimental for now. But we want to
     # reserve the URL space. So, 404 all URLs if the feature isn't enabled.
     if not repo.ui.configbool('experimental', 'web.apiserver'):
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('Experimental API server endpoint not enabled'))
         return
 
     # The URL space is /api/<protocol>/*. The structure of URLs under varies
     # by <protocol>.
 
     # Registered APIs are made available via config options of the name of
     # the protocol.
     availableapis = set()
     for k, v in API_HANDLERS.items():
         section, option = v['config']
         if repo.ui.configbool(section, option):
             availableapis.add(k)
 
     # Requests to /api/ list available APIs.
     if req.dispatchparts == [b'api']:
         res.status = b'200 OK'
         res.headers[b'Content-Type'] = b'text/plain'
         lines = [_('APIs can be accessed at /api/<name>, where <name> can be '
                    'one of the following:\n')]
         if availableapis:
             lines.extend(sorted(availableapis))
         else:
             lines.append(_('(no available APIs)\n'))
         res.setbodybytes(b'\n'.join(lines))
         return
 
     proto = req.dispatchparts[1]
 
     if proto not in API_HANDLERS:
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('Unknown API: %s\nKnown APIs: %s') % (
             proto, b', '.join(sorted(availableapis))))
         return
 
     if proto not in availableapis:
         res.status = b'404 Not Found'
         res.headers[b'Content-Type'] = b'text/plain'
         res.setbodybytes(_('API %s not enabled\n') % proto)
         return
 
     API_HANDLERS[proto]['handler'](rctx, req, res, checkperm,
                                    req.dispatchparts[2:])
 
 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.commands and command not in extracommands:
+    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.commands.commandavailable(command, proto)
+    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.commands.commandavailable(command['command'], proto):
+        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.commands[command['command']]
+        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.bytesresponse):
         action, meta = reactor.onbytesresponseready(outstream,
                                                     command['requestid'],
                                                     rsp.data)
     else:
         action, meta = reactor.onapplicationerror(
             _('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)
 
 # Maps API name to metadata so custom API can be registered.
 API_HANDLERS = {
     HTTPV2: {
         'config': ('experimental', 'web.api.http-v2'),
         'handler': _handlehttpv2request,
     },
 }
 
 class httpv2protocolhandler(wireprototypes.baseprotocolhandler):
     def __init__(self, req, ui, args=None):
         self._req = req
         self._ui = ui
         self._args = args
 
     @property
     def name(self):
         return HTTPV2
 
     def getargs(self, args):
         data = {}
         for k in args.split():
             if k == '*':
                 raise NotImplementedError('do not support * args')
             else:
                 data[k] = self._args[k]
 
         return [data[k] for k in args.split()]
 
     def forwardpayload(self, fp):
         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 _httpresponsetype(ui, req, prefer_uncompressed):
     """Determine the appropriate response type and compression settings.
 
     Returns a tuple of (mediatype, compengine, engineopts).
     """
     # Determine the response media type and compression engine based
     # on the request parameters.
     protocaps = decodevaluefromheaders(req, 'X-HgProto').split(' ')
 
     if '0.2' in protocaps:
         # All clients are expected to support uncompressed data.
         if prefer_uncompressed:
             return HGTYPE2, util._noopengine(), {}
 
         # Default as defined by wire protocol spec.
         compformats = ['zlib', 'none']
         for cap in protocaps:
             if cap.startswith('comp='):
                 compformats = cap[5:].split(',')
                 break
 
         # Now find an agreed upon compression format.
         for engine in wireproto.supportedcompengines(ui, util.SERVERROLE):
             if engine.wireprotosupport().name in compformats:
                 opts = {}
                 level = ui.configint('server', '%slevel' % engine.name())
                 if level is not None:
                     opts['level'] = level
 
                 return HGTYPE2, engine, opts
 
         # No mutually supported compression format. Fall back to the
         # legacy protocol.
 
     # Don't allow untrusted settings because disabling compression or
     # setting a very high compression level could lead to flooding
     # the server's network or CPU.
     opts = {'level': ui.configint('server', 'zliblevel')}
     return HGTYPE, util.compengines['zlib'], opts
 
 def _callhttp(repo, req, res, proto, cmd):
     # Avoid cycle involving hg module.
     from .hgweb import common as hgwebcommon
 
     def genversion2(gen, engine, engineopts):
         # application/mercurial-0.2 always sends a payload header
         # identifying the compression engine.
         name = engine.wireprotosupport().name
         assert 0 < len(name) < 256
         yield struct.pack('B', len(name))
         yield name
 
         for chunk in gen:
             yield chunk
 
     def setresponse(code, contenttype, bodybytes=None, bodygen=None):
         if code == HTTP_OK:
             res.status = '200 Script output follows'
         else:
             res.status = hgwebcommon.statusmessage(code)
 
         res.headers['Content-Type'] = contenttype
 
         if bodybytes is not None:
             res.setbodybytes(bodybytes)
         if bodygen is not None:
             res.setbodygen(bodygen)
 
     if not wireproto.commands.commandavailable(cmd, proto):
         setresponse(HTTP_OK, HGERRTYPE,
                     _('requested wire protocol command is not available over '
                       'HTTP'))
         return
 
     proto.checkperm(wireproto.commands[cmd].permission)
 
     rsp = wireproto.dispatch(repo, proto, cmd)
 
     if isinstance(rsp, bytes):
         setresponse(HTTP_OK, HGTYPE, bodybytes=rsp)
     elif isinstance(rsp, wireprototypes.bytesresponse):
         setresponse(HTTP_OK, HGTYPE, bodybytes=rsp.data)
     elif isinstance(rsp, wireprototypes.streamreslegacy):
         setresponse(HTTP_OK, HGTYPE, bodygen=rsp.gen)
     elif isinstance(rsp, wireprototypes.streamres):
         gen = rsp.gen
 
         # This code for compression should not be streamres specific. It
         # is here because we only compress streamres at the moment.
         mediatype, engine, engineopts = _httpresponsetype(
             repo.ui, req, rsp.prefer_uncompressed)
         gen = engine.compressstream(gen, engineopts)
 
         if mediatype == HGTYPE2:
             gen = genversion2(gen, engine, engineopts)
 
         setresponse(HTTP_OK, mediatype, bodygen=gen)
     elif isinstance(rsp, wireprototypes.pushres):
         rsp = '%d\n%s' % (rsp.res, rsp.output)
         setresponse(HTTP_OK, HGTYPE, bodybytes=rsp)
     elif isinstance(rsp, wireprototypes.pusherr):
         rsp = '0\n%s\n' % rsp.res
         res.drain = True
         setresponse(HTTP_OK, HGTYPE, bodybytes=rsp)
     elif isinstance(rsp, wireprototypes.ooberror):
         setresponse(HTTP_OK, HGERRTYPE, bodybytes=rsp.message)
     else:
         raise error.ProgrammingError('hgweb.protocol internal failure', rsp)
 
 def _sshv1respondbytes(fout, value):
     """Send a bytes response for protocol version 1."""
     fout.write('%d\n' % len(value))
     fout.write(value)
     fout.flush()
 
 def _sshv1respondstream(fout, source):
     write = fout.write
     for chunk in source.gen:
         write(chunk)
     fout.flush()
 
 def _sshv1respondooberror(fout, ferr, rsp):
     ferr.write(b'%s\n-\n' % rsp)
     ferr.flush()
     fout.write(b'\n')
     fout.flush()
 
 class sshv1protocolhandler(wireprototypes.baseprotocolhandler):
     """Handler for requests services via version 1 of SSH protocol."""
     def __init__(self, ui, fin, fout):
         self._ui = ui
         self._fin = fin
         self._fout = fout
 
     @property
     def name(self):
         return wireprototypes.SSHV1
 
     def getargs(self, args):
         data = {}
         keys = args.split()
         for n in xrange(len(keys)):
             argline = self._fin.readline()[:-1]
             arg, l = argline.split()
             if arg not in keys:
                 raise error.Abort(_("unexpected parameter %r") % arg)
             if arg == '*':
                 star = {}
                 for k in xrange(int(l)):
                     argline = self._fin.readline()[:-1]
                     arg, l = argline.split()
                     val = self._fin.read(int(l))
                     star[arg] = val
                 data['*'] = star
             else:
                 val = self._fin.read(int(l))
                 data[arg] = val
         return [data[k] for k in keys]
 
     def forwardpayload(self, fpout):
         # We initially send an empty response. This tells the client it is
         # OK to start sending data. If a client sees any other response, it
         # interprets it as an error.
         _sshv1respondbytes(self._fout, b'')
 
         # The file is in the form:
         #
         # <chunk size>\n<chunk>
         # ...
         # 0\n
         count = int(self._fin.readline())
         while count:
             fpout.write(self._fin.read(count))
             count = int(self._fin.readline())
 
     @contextlib.contextmanager
     def mayberedirectstdio(self):
         yield None
 
     def client(self):
         client = encoding.environ.get('SSH_CLIENT', '').split(' ', 1)[0]
         return 'remote:ssh:' + client
 
     def addcapabilities(self, repo, caps):
         caps.append(b'batch')
         return caps
 
     def checkperm(self, perm):
         pass
 
 class sshv2protocolhandler(sshv1protocolhandler):
     """Protocol handler for version 2 of the SSH protocol."""
 
     @property
     def name(self):
         return wireprototypes.SSHV2
 
 def _runsshserver(ui, repo, fin, fout, ev):
     # This function operates like a state machine of sorts. The following
     # states are defined:
     #
     # protov1-serving
     #    Server is in protocol version 1 serving mode. Commands arrive on
     #    new lines. These commands are processed in this state, one command
     #    after the other.
     #
     # protov2-serving
     #    Server is in protocol version 2 serving mode.
     #
     # upgrade-initial
     #    The server is going to process an upgrade request.
     #
     # upgrade-v2-filter-legacy-handshake
     #    The protocol is being upgraded to version 2. The server is expecting
     #    the legacy handshake from version 1.
     #
     # upgrade-v2-finish
     #    The upgrade to version 2 of the protocol is imminent.
     #
     # shutdown
     #    The server is shutting down, possibly in reaction to a client event.
     #
     # And here are their transitions:
     #
     # protov1-serving -> shutdown
     #    When server receives an empty request or encounters another
     #    error.
     #
     # protov1-serving -> upgrade-initial
     #    An upgrade request line was seen.
     #
     # upgrade-initial -> upgrade-v2-filter-legacy-handshake
     #    Upgrade to version 2 in progress. Server is expecting to
     #    process a legacy handshake.
     #
     # upgrade-v2-filter-legacy-handshake -> shutdown
     #    Client did not fulfill upgrade handshake requirements.
     #
     # upgrade-v2-filter-legacy-handshake -> upgrade-v2-finish
     #    Client fulfilled version 2 upgrade requirements. Finishing that
     #    upgrade.
     #
     # upgrade-v2-finish -> protov2-serving
     #    Protocol upgrade to version 2 complete. Server can now speak protocol
     #    version 2.
     #
     # protov2-serving -> protov1-serving
     #    Ths happens by default since protocol version 2 is the same as
     #    version 1 except for the handshake.
 
     state = 'protov1-serving'
     proto = sshv1protocolhandler(ui, fin, fout)
     protoswitched = False
 
     while not ev.is_set():
         if state == 'protov1-serving':
             # Commands are issued on new lines.
             request = fin.readline()[:-1]
 
             # Empty lines signal to terminate the connection.
             if not request:
                 state = 'shutdown'
                 continue
 
             # It looks like a protocol upgrade request. Transition state to
             # handle it.
             if request.startswith(b'upgrade '):
                 if protoswitched:
                     _sshv1respondooberror(fout, ui.ferr,
                                           b'cannot upgrade protocols multiple '
                                           b'times')
                     state = 'shutdown'
                     continue
 
                 state = 'upgrade-initial'
                 continue
 
             available = wireproto.commands.commandavailable(request, proto)
 
             # This command isn't available. Send an empty response and go
             # back to waiting for a new command.
             if not available:
                 _sshv1respondbytes(fout, b'')
                 continue
 
             rsp = wireproto.dispatch(repo, proto, request)
 
             if isinstance(rsp, bytes):
                 _sshv1respondbytes(fout, rsp)
             elif isinstance(rsp, wireprototypes.bytesresponse):
                 _sshv1respondbytes(fout, rsp.data)
             elif isinstance(rsp, wireprototypes.streamres):
                 _sshv1respondstream(fout, rsp)
             elif isinstance(rsp, wireprototypes.streamreslegacy):
                 _sshv1respondstream(fout, rsp)
             elif isinstance(rsp, wireprototypes.pushres):
                 _sshv1respondbytes(fout, b'')
                 _sshv1respondbytes(fout, b'%d' % rsp.res)
             elif isinstance(rsp, wireprototypes.pusherr):
                 _sshv1respondbytes(fout, rsp.res)
             elif isinstance(rsp, wireprototypes.ooberror):
                 _sshv1respondooberror(fout, ui.ferr, rsp.message)
             else:
                 raise error.ProgrammingError('unhandled response type from '
                                              'wire protocol command: %s' % rsp)
 
         # For now, protocol version 2 serving just goes back to version 1.
         elif state == 'protov2-serving':
             state = 'protov1-serving'
             continue
 
         elif state == 'upgrade-initial':
             # We should never transition into this state if we've switched
             # protocols.
             assert not protoswitched
             assert proto.name == wireprototypes.SSHV1
 
             # Expected: upgrade <token> <capabilities>
             # If we get something else, the request is malformed. It could be
             # from a future client that has altered the upgrade line content.
             # We treat this as an unknown command.
             try:
                 token, caps = request.split(b' ')[1:]
             except ValueError:
                 _sshv1respondbytes(fout, b'')
                 state = 'protov1-serving'
                 continue
 
             # Send empty response if we don't support upgrading protocols.
             if not ui.configbool('experimental', 'sshserver.support-v2'):
                 _sshv1respondbytes(fout, b'')
                 state = 'protov1-serving'
                 continue
 
             try:
                 caps = urlreq.parseqs(caps)
             except ValueError:
                 _sshv1respondbytes(fout, b'')
                 state = 'protov1-serving'
                 continue
 
             # We don't see an upgrade request to protocol version 2. Ignore
             # the upgrade request.
             wantedprotos = caps.get(b'proto', [b''])[0]
             if SSHV2 not in wantedprotos:
                 _sshv1respondbytes(fout, b'')
                 state = 'protov1-serving'
                 continue
 
             # It looks like we can honor this upgrade request to protocol 2.
             # Filter the rest of the handshake protocol request lines.
             state = 'upgrade-v2-filter-legacy-handshake'
             continue
 
         elif state == 'upgrade-v2-filter-legacy-handshake':
             # Client should have sent legacy handshake after an ``upgrade``
             # request. Expected lines:
             #
             #    hello
             #    between
             #    pairs 81
             #    0000...-0000...
 
             ok = True
             for line in (b'hello', b'between', b'pairs 81'):
                 request = fin.readline()[:-1]
 
                 if request != line:
                     _sshv1respondooberror(fout, ui.ferr,
                                           b'malformed handshake protocol: '
                                           b'missing %s' % line)
                     ok = False
                     state = 'shutdown'
                     break
 
             if not ok:
                 continue
 
             request = fin.read(81)
             if request != b'%s-%s' % (b'0' * 40, b'0' * 40):
                 _sshv1respondooberror(fout, ui.ferr,
                                       b'malformed handshake protocol: '
                                       b'missing between argument value')
                 state = 'shutdown'
                 continue
 
             state = 'upgrade-v2-finish'
             continue
 
         elif state == 'upgrade-v2-finish':
             # Send the upgrade response.
             fout.write(b'upgraded %s %s\n' % (token, SSHV2))
             servercaps = wireproto.capabilities(repo, proto)
             rsp = b'capabilities: %s' % servercaps.data
             fout.write(b'%d\n%s\n' % (len(rsp), rsp))
             fout.flush()
 
             proto = sshv2protocolhandler(ui, fin, fout)
             protoswitched = True
 
             state = 'protov2-serving'
             continue
 
         elif state == 'shutdown':
             break
 
         else:
             raise error.ProgrammingError('unhandled ssh server state: %s' %
                                          state)
 
 class sshserver(object):
     def __init__(self, ui, repo, logfh=None):
         self._ui = ui
         self._repo = repo
         self._fin = ui.fin
         self._fout = ui.fout
 
         # Log write I/O to stdout and stderr if configured.
         if logfh:
             self._fout = util.makeloggingfileobject(
                 logfh, self._fout, 'o', logdata=True)
             ui.ferr = util.makeloggingfileobject(
                 logfh, ui.ferr, 'e', logdata=True)
 
         hook.redirect(True)
         ui.fout = repo.ui.fout = ui.ferr
 
         # Prevent insertion/deletion of CRs
         procutil.setbinary(self._fin)
         procutil.setbinary(self._fout)
 
     def serve_forever(self):
         self.serveuntil(threading.Event())
         sys.exit(0)
 
     def serveuntil(self, ev):
         """Serve until a threading.Event is set."""
         _runsshserver(self._ui, self._repo, self._fin, self._fout, ev)

tests/test-wireproto.py

@@ -1,95 +1,102 b''
 from __future__ import absolute_import, print_function
 
 from mercurial import (
     error,
     pycompat,
     ui as uimod,
     util,
     wireproto,
     wireprototypes,
 )
 stringio = util.stringio
 
 class proto(object):
     def __init__(self, args):
         self.args = args
+        self.name = 'dummyproto'
+
     def getargs(self, spec):
         args = self.args
         args.setdefault(b'*', {})
         names = spec.split()
         return [args[n] for n in names]
 
     def checkperm(self, perm):
         pass
 
+wireprototypes.TRANSPORTS['dummyproto'] = {
+    'transport': 'dummy',
+    'version': 1,
+}
+
 class clientpeer(wireproto.wirepeer):
     def __init__(self, serverrepo, ui):
         self.serverrepo = serverrepo
         self._ui = ui
 
     @property
     def ui(self):
         return self._ui
 
     def url(self):
         return b'test'
 
     def local(self):
         return None
 
     def peer(self):
         return self
 
     def canpush(self):
         return True
 
     def close(self):
         pass
 
     def capabilities(self):
         return [b'batch']
 
     def _call(self, cmd, **args):
         args = pycompat.byteskwargs(args)
         res = wireproto.dispatch(self.serverrepo, proto(args), cmd)
         if isinstance(res, wireprototypes.bytesresponse):
             return res.data
         elif isinstance(res, bytes):
             return res
         else:
             raise error.Abort('dummy client does not support response type')
 
     def _callstream(self, cmd, **args):
         return stringio(self._call(cmd, **args))
 
     @wireproto.batchable
     def greet(self, name):
         f = wireproto.future()
         yield {b'name': mangle(name)}, f
         yield unmangle(f.value)
 
 class serverrepo(object):
     def greet(self, name):
         return b"Hello, " + name
 
     def filtered(self, name):
         return self
 
 def mangle(s):
     return b''.join(pycompat.bytechr(ord(c) + 1) for c in pycompat.bytestr(s))
 def unmangle(s):
     return b''.join(pycompat.bytechr(ord(c) - 1) for c in pycompat.bytestr(s))
 
 def greet(repo, proto, name):
     return mangle(repo.greet(unmangle(name)))
 
 wireproto.commands[b'greet'] = (greet, b'name',)
 
 srv = serverrepo()
 clt = clientpeer(srv, uimod.ui())
 
 print(clt.greet(b"Foobar"))
 b = clt.iterbatch()
 list(map(b.greet, (b'Fo, =;:<o', b'Bar')))
 b.submit()
 print([r for r in b.results()])