wireproto: use maybecapturestdio() for push responses (API)

The "pushres" and "pusherr" response types currently call
proto.restore() in the HTTP protocol. This completes the pairing
with proto.redirect() that occurs in the @wireprotocommand
functions. (But since the SSH protocol has a no-op redirect(),
it doesn't bother calling restore() because it would also be
a no-op.)

Having the disconnect between these paired calls is very confusing.
Knowing that you must use proto.redirect() if returning a "pushres"
or a "pusherr" is even wonkier.

We replace this confusing code with our new context manager for
[maybe] capturing output.

The "pushres" and "pusherr" types have gained an "output" argument
to their constructor and an attribute to hold captured data. The
HTTP protocol now retrieves output from these objects.

.. api::

   ``wireproto.pushres`` and ``wireproto.pusherr`` now explicitly
   track stdio output.

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

hgext/largefiles/proto.py

@@ -1,190 +1,189
 # Copyright 2011 Fog Creek Software
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 from __future__ import absolute_import
 
 import os
 import re
 
 from mercurial.i18n import _
 
 from mercurial import (
     error,
     httppeer,
     util,
     wireproto,
 )
 
 from . import (
     lfutil,
 )
 
 urlerr = util.urlerr
 urlreq = util.urlreq
 
 LARGEFILES_REQUIRED_MSG = ('\nThis repository uses the largefiles extension.'
                            '\n\nPlease enable it in your Mercurial config '
                            'file.\n')
 
 # these will all be replaced by largefiles.uisetup
 ssholdcallstream = None
 httpoldcallstream = None
 
 def putlfile(repo, proto, sha):
     '''Server command for putting a largefile into a repository's local store
     and into the user cache.'''
-    proto.redirect()
-
+    with proto.mayberedirectstdio() as output:
     path = lfutil.storepath(repo, sha)
     util.makedirs(os.path.dirname(path))
     tmpfp = util.atomictempfile(path, createmode=repo.store.createmode)
 
     try:
         proto.getfile(tmpfp)
         tmpfp._fp.seek(0)
         if sha != lfutil.hexsha1(tmpfp._fp):
             raise IOError(0, _('largefile contents do not match hash'))
         tmpfp.close()
         lfutil.linktousercache(repo, sha)
     except IOError as e:
         repo.ui.warn(_('largefiles: failed to put %s into store: %s\n') %
                      (sha, e.strerror))
-        return wireproto.pushres(1)
+            return wireproto.pushres(1, output.getvalue() if output else '')
     finally:
         tmpfp.discard()
 
-    return wireproto.pushres(0)
+    return wireproto.pushres(0, output.getvalue() if output else '')
 
 def getlfile(repo, proto, sha):
     '''Server command for retrieving a largefile from the repository-local
     cache or user cache.'''
     filename = lfutil.findfile(repo, sha)
     if not filename:
         raise error.Abort(_('requested largefile %s not present in cache')
                           % sha)
     f = open(filename, 'rb')
     length = os.fstat(f.fileno())[6]
 
     # Since we can't set an HTTP content-length header here, and
     # Mercurial core provides no way to give the length of a streamres
     # (and reading the entire file into RAM would be ill-advised), we
     # just send the length on the first line of the response, like the
     # ssh proto does for string responses.
     def generator():
         yield '%d\n' % length
         for chunk in util.filechunkiter(f):
             yield chunk
     return wireproto.streamres_legacy(gen=generator())
 
 def statlfile(repo, proto, sha):
     '''Server command for checking if a largefile is present - returns '2\n' if
     the largefile is missing, '0\n' if it seems to be in good condition.
 
     The value 1 is reserved for mismatched checksum, but that is too expensive
     to be verified on every stat and must be caught be running 'hg verify'
     server side.'''
     filename = lfutil.findfile(repo, sha)
     if not filename:
         return '2\n'
     return '0\n'
 
 def wirereposetup(ui, repo):
     class lfileswirerepository(repo.__class__):
         def putlfile(self, sha, fd):
             # unfortunately, httprepository._callpush tries to convert its
             # input file-like into a bundle before sending it, so we can't use
             # it ...
             if issubclass(self.__class__, httppeer.httppeer):
                 res = self._call('putlfile', data=fd, sha=sha,
                     headers={'content-type':'application/mercurial-0.1'})
                 try:
                     d, output = res.split('\n', 1)
                     for l in output.splitlines(True):
                         self.ui.warn(_('remote: '), l) # assume l ends with \n
                     return int(d)
                 except ValueError:
                     self.ui.warn(_('unexpected putlfile response: %r\n') % res)
                     return 1
             # ... but we can't use sshrepository._call because the data=
             # argument won't get sent, and _callpush does exactly what we want
             # in this case: send the data straight through
             else:
                 try:
                     ret, output = self._callpush("putlfile", fd, sha=sha)
                     if ret == "":
                         raise error.ResponseError(_('putlfile failed:'),
                                 output)
                     return int(ret)
                 except IOError:
                     return 1
                 except ValueError:
                     raise error.ResponseError(
                         _('putlfile failed (unexpected response):'), ret)
 
         def getlfile(self, sha):
             """returns an iterable with the chunks of the file with sha sha"""
             stream = self._callstream("getlfile", sha=sha)
             length = stream.readline()
             try:
                 length = int(length)
             except ValueError:
                 self._abort(error.ResponseError(_("unexpected response:"),
                                                 length))
 
             # SSH streams will block if reading more than length
             for chunk in util.filechunkiter(stream, limit=length):
                 yield chunk
             # HTTP streams must hit the end to process the last empty
             # chunk of Chunked-Encoding so the connection can be reused.
             if issubclass(self.__class__, httppeer.httppeer):
                 chunk = stream.read(1)
                 if chunk:
                     self._abort(error.ResponseError(_("unexpected response:"),
                                                     chunk))
 
         @wireproto.batchable
         def statlfile(self, sha):
             f = wireproto.future()
             result = {'sha': sha}
             yield result, f
             try:
                 yield int(f.value)
             except (ValueError, urlerr.httperror):
                 # If the server returns anything but an integer followed by a
                 # newline, newline, it's not speaking our language; if we get
                 # an HTTP error, we can't be sure the largefile is present;
                 # either way, consider it missing.
                 yield 2
 
     repo.__class__ = lfileswirerepository
 
 # advertise the largefiles=serve capability
 def _capabilities(orig, repo, proto):
     '''announce largefile server capability'''
     caps = orig(repo, proto)
     caps.append('largefiles=serve')
     return caps
 
 def heads(repo, proto):
     '''Wrap server command - largefile capable clients will know to call
     lheads instead'''
     if lfutil.islfilesrepo(repo):
         return wireproto.ooberror(LARGEFILES_REQUIRED_MSG)
     return wireproto.heads(repo, proto)
 
 def sshrepocallstream(self, cmd, **args):
     if cmd == 'heads' and self.capable('largefiles'):
         cmd = 'lheads'
     if cmd == 'batch' and self.capable('largefiles'):
         args[r'cmds'] = args[r'cmds'].replace('heads ', 'lheads ')
     return ssholdcallstream(self, cmd, **args)
 
 headsre = re.compile(r'(^|;)heads\b')
 
 def httprepocallstream(self, cmd, **args):
     if cmd == 'heads' and self.capable('largefiles'):
         cmd = 'lheads'
     if cmd == 'batch' and self.capable('largefiles'):
         args[r'cmds'] = headsre.sub('lheads', args[r'cmds'])
     return httpoldcallstream(self, cmd, **args)

mercurial/wireproto.py

@@ -1,1093 +1,1096
 # 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,
 )
 
 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.
         """
         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.
 class streamres(object):
     """wireproto reply: binary stream
 
     The call was successful and the result is a stream.
 
     Accepts a generator containing chunks of data to be sent to the client.
 
     ``prefer_uncompressed`` indicates that the data is expected to be
     uncompressable and that the stream should therefore use the ``none``
     engine.
     """
     def __init__(self, gen=None, prefer_uncompressed=False):
         self.gen = gen
         self.prefer_uncompressed = prefer_uncompressed
 
 class streamres_legacy(object):
     """wireproto reply: uncompressed binary stream
 
     The call was successful and the result is a stream.
 
     Accepts a generator containing chunks of data to be sent to the client.
 
     Like ``streamres``, but sends an uncompressed data for "version 1" clients
     using the application/mercurial-0.1 media type.
     """
     def __init__(self, gen=None):
         self.gen = gen
 
 class pushres(object):
     """wireproto reply: success with simple integer return
 
     The call was successful and returned an integer contained in `self.res`.
     """
-    def __init__(self, res):
+    def __init__(self, res, output):
         self.res = res
+        self.output = output
 
 class pusherr(object):
     """wireproto reply: failure
 
     The call failed. The `self.res` attribute contains the error message.
     """
-    def __init__(self, res):
+    def __init__(self, res, output):
         self.res = res
+        self.output = output
 
 class ooberror(object):
     """wireproto reply: failure of a batch of operation
 
     Something failed during a batch call. The error message is stored in
     `self.message`.
     """
     def __init__(self, message):
         self.message = message
 
 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]
     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:
         util.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, proto, 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=''):
         self.func = func
         self.args = args
 
     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)
 
     # 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:
                 v = commandentry(v[0], v[1])
         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."""
         # For now, commands are available for all protocols. So do a simple
         # membership test.
         return command in self
 
 commands = commanddict()
 
 def wireprotocommand(name, args=''):
     """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.
     """
     def register(func):
         commands[name] = commandentry(func, args)
         return func
     return register
 
 @wireprotocommand('batch', 'cmds *')
 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]
         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, ooberror):
             return result
         res.append(escapearg(result))
     return ';'.join(res)
 
 @wireprotocommand('between', 'pairs')
 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 "".join(r)
 
 @wireprotocommand('branchmap')
 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 '\n'.join(heads)
 
 @wireprotocommand('branches', 'nodes')
 def branches(repo, proto, nodes):
     nodes = decodelist(nodes)
     r = []
     for b in repo.branches(nodes):
         r.append(encodelist(b) + "\n")
     return "".join(r)
 
 @wireprotocommand('clonebundles', '')
 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 repo.vfs.tryread('clonebundles.manifest')
 
 wireprotocaps = ['lookup', 'changegroupsubset', 'branchmap', 'pushkey',
                  'known', 'getbundle', 'unbundlehash', 'batch']
 
 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)
     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))
 
     if proto.name == 'http':
         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 = supportedcompengines(repo.ui, proto, util.SERVERROLE)
         if compengines:
             comptypes = ','.join(urlreq.quote(e.wireprotosupport().name)
                                  for e in compengines)
             caps.append('compression=%s' % comptypes)
 
     return caps
 
 # If you are writing an extension and consider wrapping this function. Wrap
 # `_capabilities` instead.
 @wireprotocommand('capabilities')
 def capabilities(repo, proto):
     return ' '.join(_capabilities(repo, proto))
 
 @wireprotocommand('changegroup', 'roots')
 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 streamres(gen=gen)
 
 @wireprotocommand('changegroupsubset', 'bases heads')
 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 streamres(gen=gen)
 
 @wireprotocommand('debugwireargs', 'one two *')
 def debugwireargs(repo, proto, one, two, others):
     # only accept optional args from the known set
     opts = options('debugwireargs', ['three', 'four'], others)
     return repo.debugwireargs(one, two, **pycompat.strkwargs(opts))
 
 @wireprotocommand('getbundle', '*')
 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':
                 return 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':
                 return ooberror(str(exc) + '\n')
             raise # cannot do better for bundle1 + ssh
         # bundle2 request expect a bundle2 reply
         bundler = bundle2.bundle20(repo.ui)
         manargs = [('message', str(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 streamres(gen=chunks, prefer_uncompressed=not prefercompressed)
 
 @wireprotocommand('heads')
 def heads(repo, proto):
     h = repo.heads()
     return encodelist(h) + "\n"
 
 @wireprotocommand('hello')
 def hello(repo, proto):
     '''the hello command returns a set of lines describing various
     interesting things about the server, in an RFC822-like format.
     Currently the only one defined is "capabilities", which
     consists of a line in the form:
 
     capabilities: space separated list of tokens
     '''
     return "capabilities: %s\n" % (capabilities(repo, proto))
 
 @wireprotocommand('listkeys', 'namespace')
 def listkeys(repo, proto, namespace):
     d = repo.listkeys(encoding.tolocal(namespace)).items()
     return pushkeymod.encodekeys(d)
 
 @wireprotocommand('lookup', 'key')
 def lookup(repo, proto, key):
     try:
         k = encoding.tolocal(key)
         c = repo[k]
         r = c.hex()
         success = 1
     except Exception as inst:
         r = str(inst)
         success = 0
     return "%d %s\n" % (success, r)
 
 @wireprotocommand('known', 'nodes *')
 def known(repo, proto, nodes, others):
     return ''.join(b and "1" or "0" for b in repo.known(decodelist(nodes)))
 
 @wireprotocommand('pushkey', 'namespace key old new')
 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 util.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 '%s\n%s' % (int(r), output)
 
 @wireprotocommand('stream_out')
 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 streamres_legacy(streamclone.generatev1wireproto(repo))
 
 @wireprotocommand('unbundle', 'heads')
 def unbundle(repo, proto, heads):
     their_heads = decodelist(heads)
 
+    with proto.mayberedirectstdio() as output:
     try:
-        proto.redirect()
-
         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, pycompat.sysstr('wb+'))
         r = 0
         try:
             proto.getfile(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':
                     # 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.
+                        # 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 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.
+                    # The return looks streamable, we are in the bundle2 case
+                    # and should return a stream.
                 return streamres_legacy(gen=r.getchunks())
-            return pushres(r)
+                return 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 util.stderr directly.
                 # We did not change it to minimise code change.
                 # This need to be moved to something proper.
                 # Feel free to do it.
                 util.stderr.write("abort: %s\n" % exc)
                 if exc.hint is not None:
                     util.stderr.write("(%s)\n" % exc.hint)
-                return pushres(0)
+                    return pushres(0, output.getvalue() if output else '')
             except error.PushRaced:
-                return pusherr(str(exc))
+                    return pusherr(str(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)
+                        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', str(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', str(exc))])
         return streamres_legacy(gen=bundler.getchunks())

mercurial/wireprotoserver.py

@@ -1,481 +1,479
 # 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 abc
 import cgi
 import contextlib
 import struct
 import sys
 
 from .i18n import _
 from . import (
     encoding,
     error,
     hook,
     pycompat,
     util,
     wireproto,
 )
 
 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'
 
 # Names of the SSH protocol implementations.
 SSHV1 = 'ssh-v1'
 # This is advertised over the wire. Incremental the counter at the end
 # to reflect BC breakages.
 SSHV2 = 'exp-ssh-v2-0001'
 
 class baseprotocolhandler(object):
     """Abstract base class for wire protocol handlers.
 
     A wire protocol handler serves as an interface between protocol command
     handlers and the wire protocol transport layer. Protocol handlers provide
     methods to read command arguments, redirect stdio for the duration of
     the request, handle response types, etc.
     """
 
     __metaclass__ = abc.ABCMeta
 
     @abc.abstractproperty
     def name(self):
         """The name of the protocol implementation.
 
         Used for uniquely identifying the transport type.
         """
 
     @abc.abstractmethod
     def getargs(self, args):
         """return the value for arguments in <args>
 
         returns a list of values (same order as <args>)"""
 
     @abc.abstractmethod
     def getfile(self, fp):
         """write the whole content of a file into a file like object
 
         The file is in the form::
 
             (<chunk-size>\n<chunk>)+0\n
 
         chunk size is the ascii version of the int.
         """
 
     @abc.abstractmethod
     def mayberedirectstdio(self):
         """Context manager to possibly redirect stdio.
 
         The context manager yields a file-object like object that receives
         stdout and stderr output when the context manager is active. Or it
         yields ``None`` if no I/O redirection occurs.
 
         The intent of this context manager is to capture stdio output
         so it may be sent in the response. Some transports support streaming
         stdio to the client in real time. For these transports, stdio output
         won't be captured.
         """
 
     @abc.abstractmethod
     def redirect(self):
         """may setup interception for stdout and stderr
 
         See also the `restore` method."""
 
     # If the `redirect` function does install interception, the `restore`
     # function MUST be defined. If interception is not used, this function
     # MUST NOT be defined.
     #
     # left commented here on purpose
     #
     #def restore(self):
     #    """reinstall previous stdout and stderr and return intercepted stdout
     #    """
     #    raise NotImplementedError()
 
 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
     prefix = headerprefix.upper().replace(r'-', r'_')
     while True:
         v = req.env.get(r'HTTP_%s_%d' % (prefix, i))
         if v is None:
             break
         chunks.append(pycompat.bytesurl(v))
         i += 1
 
     return ''.join(chunks)
 
 class webproto(baseprotocolhandler):
     def __init__(self, req, ui):
         self._req = req
         self._ui = ui
 
     @property
     def name(self):
         return 'http'
 
     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 = util.rapply(pycompat.bytesurl, self._req.form.copy())
         postlen = int(self._req.env.get(r'HTTP_X_HGARGS_POST', 0))
         if postlen:
             args.update(cgi.parse_qs(
                 self._req.read(postlen), keep_blank_values=True))
             return args
 
         argvalue = decodevaluefromheaders(self._req, r'X-HgArg')
         args.update(cgi.parse_qs(argvalue, keep_blank_values=True))
         return args
 
     def getfile(self, fp):
         length = int(self._req.env[r'CONTENT_LENGTH'])
         # If httppostargs is used, we need to read Content-Length
         # minus the amount that was consumed by args.
         length -= int(self._req.env.get(r'HTTP_X_HGARGS_POST', 0))
         for s in util.filechunkiter(self._req, 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 redirect(self):
         self._oldio = self._ui.fout, self._ui.ferr
         self._ui.ferr = self._ui.fout = stringio()
 
     def restore(self):
         val = self._ui.fout.getvalue()
         self._ui.ferr, self._ui.fout = self._oldio
         return val
 
     def _client(self):
         return 'remote:%s:%s:%s' % (
             self._req.env.get('wsgi.url_scheme') or 'http',
             urlreq.quote(self._req.env.get('REMOTE_HOST', '')),
             urlreq.quote(self._req.env.get('REMOTE_USER', '')))
 
     def responsetype(self, 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(self._req, r'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(self._ui, self,
                                                          util.SERVERROLE):
                 if engine.wireprotosupport().name in compformats:
                     opts = {}
                     level = self._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': self._ui.configint('server', 'zliblevel')}
         return HGTYPE, util.compengines['zlib'], opts
 
 def iscmd(cmd):
     return cmd in wireproto.commands
 
 def parsehttprequest(repo, req, query):
     """Parse the HTTP request for a wire protocol request.
 
     If the current request appears to be a wire protocol request, this
     function returns a dict with details about that request, including
     an ``abstractprotocolserver`` instance suitable for handling the
     request. Otherwise, ``None`` is returned.
 
     ``req`` is a ``wsgirequest`` instance.
     """
     # 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 r'cmd' not in req.form:
         return None
 
     cmd = pycompat.sysbytes(req.form[r'cmd'][0])
 
     # 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 cmd not in wireproto.commands:
         return None
 
     proto = webproto(req, repo.ui)
 
     return {
         'cmd': cmd,
         'proto': proto,
         'dispatch': lambda: _callhttp(repo, req, proto, cmd),
         'handleerror': lambda ex: _handlehttperror(ex, req, cmd),
     }
 
 def _callhttp(repo, req, proto, cmd):
     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
 
     rsp = wireproto.dispatch(repo, proto, cmd)
 
     if not wireproto.commands.commandavailable(cmd, proto):
         req.respond(HTTP_OK, HGERRTYPE,
                     body=_('requested wire protocol command is not available '
                            'over HTTP'))
         return []
 
     if isinstance(rsp, bytes):
         req.respond(HTTP_OK, HGTYPE, body=rsp)
         return []
     elif isinstance(rsp, wireproto.streamres_legacy):
         gen = rsp.gen
         req.respond(HTTP_OK, HGTYPE)
         return gen
     elif isinstance(rsp, wireproto.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 = proto.responsetype(
             rsp.prefer_uncompressed)
         gen = engine.compressstream(gen, engineopts)
 
         if mediatype == HGTYPE2:
             gen = genversion2(gen, engine, engineopts)
 
         req.respond(HTTP_OK, mediatype)
         return gen
     elif isinstance(rsp, wireproto.pushres):
-        val = proto.restore()
-        rsp = '%d\n%s' % (rsp.res, val)
+        rsp = '%d\n%s' % (rsp.res, rsp.output)
         req.respond(HTTP_OK, HGTYPE, body=rsp)
         return []
     elif isinstance(rsp, wireproto.pusherr):
         # This is the httplib workaround documented in _handlehttperror().
         req.drain()
 
-        proto.restore()
         rsp = '0\n%s\n' % rsp.res
         req.respond(HTTP_OK, HGTYPE, body=rsp)
         return []
     elif isinstance(rsp, wireproto.ooberror):
         rsp = rsp.message
         req.respond(HTTP_OK, HGERRTYPE, body=rsp)
         return []
     raise error.ProgrammingError('hgweb.protocol internal failure', rsp)
 
 def _handlehttperror(e, req, cmd):
     """Called when an ErrorResponse is raised during HTTP request processing."""
 
     # Clients using Python's httplib are stateful: the HTTP client
     # won't process an HTTP response until all request data is
     # sent to the server. The intent of this code is to ensure
     # we always read HTTP request data from the client, thus
     # ensuring httplib transitions to a state that allows it to read
     # the HTTP response. In other words, it helps prevent deadlocks
     # on clients using httplib.
 
     if (req.env[r'REQUEST_METHOD'] == r'POST' and
         # But not if Expect: 100-continue is being used.
         (req.env.get('HTTP_EXPECT',
                      '').lower() != '100-continue') or
         # Or the non-httplib HTTP library is being advertised by
         # the client.
         req.env.get('X-HgHttp2', '')):
         req.drain()
     else:
         req.headers.append((r'Connection', r'Close'))
 
     # TODO This response body assumes the failed command was
     # "unbundle." That assumption is not always valid.
     req.respond(e, HGTYPE, body='0\n%s\n' % e)
 
     return ''
 
 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(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 'ssh'
 
     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 getfile(self, fpout):
         _sshv1respondbytes(self._fout, b'')
         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 redirect(self):
         pass
 
     def _client(self):
         client = encoding.environ.get('SSH_CLIENT', '').split(' ', 1)[0]
         return 'remote:ssh:' + client
 
 class sshserver(object):
     def __init__(self, ui, repo):
         self._ui = ui
         self._repo = repo
         self._fin = ui.fin
         self._fout = ui.fout
 
         hook.redirect(True)
         ui.fout = repo.ui.fout = ui.ferr
 
         # Prevent insertion/deletion of CRs
         util.setbinary(self._fin)
         util.setbinary(self._fout)
 
         self._proto = sshv1protocolhandler(self._ui, self._fin, self._fout)
 
     def serve_forever(self):
         while self.serve_one():
             pass
         sys.exit(0)
 
     def serve_one(self):
         cmd = self._fin.readline()[:-1]
         if cmd and wireproto.commands.commandavailable(cmd, self._proto):
             rsp = wireproto.dispatch(self._repo, self._proto, cmd)
 
             if isinstance(rsp, bytes):
                 _sshv1respondbytes(self._fout, rsp)
             elif isinstance(rsp, wireproto.streamres):
                 _sshv1respondstream(self._fout, rsp)
             elif isinstance(rsp, wireproto.streamres_legacy):
                 _sshv1respondstream(self._fout, rsp)
             elif isinstance(rsp, wireproto.pushres):
                 _sshv1respondbytes(self._fout, b'')
                 _sshv1respondbytes(self._fout, bytes(rsp.res))
             elif isinstance(rsp, wireproto.pusherr):
                 _sshv1respondbytes(self._fout, rsp.res)
             elif isinstance(rsp, wireproto.ooberror):
                 _sshv1respondooberror(self._fout, self._ui.ferr, rsp.message)
             else:
                 raise error.ProgrammingError('unhandled response type from '
                                              'wire protocol command: %s' % rsp)
         elif cmd:
             _sshv1respondbytes(self._fout, b'')
         return cmd != ''