wireproto: start to associate frame generation with a stream

An upcoming commit will introduce "streams" into the frame-based wire
protocol. In preparation for this invasive change, we introduce a basic
"stream" class and have all operations that create frames also operate
alongside a stream instance.

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

mercurial/wireprotoframing.py

@@ -1,720 +1,732
 # wireprotoframing.py - unified framing protocol for wire protocol
 #
 # Copyright 2018 Gregory Szorc <[e-mail]>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 # This file contains functionality to support the unified frame-based wire
 # protocol. For details about the protocol, see
 # `hg help internals.wireprotocol`.
 
 from __future__ import absolute_import
 
 import struct
 
 from .i18n import _
 from .thirdparty import (
     attr,
 )
 from . import (
     error,
     util,
 )
 from .utils import (
     stringutil,
 )
 
 FRAME_HEADER_SIZE = 6
 DEFAULT_MAX_FRAME_SIZE = 32768
 
 FRAME_TYPE_COMMAND_NAME = 0x01
 FRAME_TYPE_COMMAND_ARGUMENT = 0x02
 FRAME_TYPE_COMMAND_DATA = 0x03
 FRAME_TYPE_BYTES_RESPONSE = 0x04
 FRAME_TYPE_ERROR_RESPONSE = 0x05
 FRAME_TYPE_TEXT_OUTPUT = 0x06
 
 FRAME_TYPES = {
     b'command-name': FRAME_TYPE_COMMAND_NAME,
     b'command-argument': FRAME_TYPE_COMMAND_ARGUMENT,
     b'command-data': FRAME_TYPE_COMMAND_DATA,
     b'bytes-response': FRAME_TYPE_BYTES_RESPONSE,
     b'error-response': FRAME_TYPE_ERROR_RESPONSE,
     b'text-output': FRAME_TYPE_TEXT_OUTPUT,
 }
 
 FLAG_COMMAND_NAME_EOS = 0x01
 FLAG_COMMAND_NAME_HAVE_ARGS = 0x02
 FLAG_COMMAND_NAME_HAVE_DATA = 0x04
 
 FLAGS_COMMAND = {
     b'eos': FLAG_COMMAND_NAME_EOS,
     b'have-args': FLAG_COMMAND_NAME_HAVE_ARGS,
     b'have-data': FLAG_COMMAND_NAME_HAVE_DATA,
 }
 
 FLAG_COMMAND_ARGUMENT_CONTINUATION = 0x01
 FLAG_COMMAND_ARGUMENT_EOA = 0x02
 
 FLAGS_COMMAND_ARGUMENT = {
     b'continuation': FLAG_COMMAND_ARGUMENT_CONTINUATION,
     b'eoa': FLAG_COMMAND_ARGUMENT_EOA,
 }
 
 FLAG_COMMAND_DATA_CONTINUATION = 0x01
 FLAG_COMMAND_DATA_EOS = 0x02
 
 FLAGS_COMMAND_DATA = {
     b'continuation': FLAG_COMMAND_DATA_CONTINUATION,
     b'eos': FLAG_COMMAND_DATA_EOS,
 }
 
 FLAG_BYTES_RESPONSE_CONTINUATION = 0x01
 FLAG_BYTES_RESPONSE_EOS = 0x02
 
 FLAGS_BYTES_RESPONSE = {
     b'continuation': FLAG_BYTES_RESPONSE_CONTINUATION,
     b'eos': FLAG_BYTES_RESPONSE_EOS,
 }
 
 FLAG_ERROR_RESPONSE_PROTOCOL = 0x01
 FLAG_ERROR_RESPONSE_APPLICATION = 0x02
 
 FLAGS_ERROR_RESPONSE = {
     b'protocol': FLAG_ERROR_RESPONSE_PROTOCOL,
     b'application': FLAG_ERROR_RESPONSE_APPLICATION,
 }
 
 # Maps frame types to their available flags.
 FRAME_TYPE_FLAGS = {
     FRAME_TYPE_COMMAND_NAME: FLAGS_COMMAND,
     FRAME_TYPE_COMMAND_ARGUMENT: FLAGS_COMMAND_ARGUMENT,
     FRAME_TYPE_COMMAND_DATA: FLAGS_COMMAND_DATA,
     FRAME_TYPE_BYTES_RESPONSE: FLAGS_BYTES_RESPONSE,
     FRAME_TYPE_ERROR_RESPONSE: FLAGS_ERROR_RESPONSE,
     FRAME_TYPE_TEXT_OUTPUT: {},
 }
 
 ARGUMENT_FRAME_HEADER = struct.Struct(r'<HH')
 
 @attr.s(slots=True)
 class frameheader(object):
     """Represents the data in a frame header."""
 
     length = attr.ib()
     requestid = attr.ib()
     typeid = attr.ib()
     flags = attr.ib()
 
 @attr.s(slots=True)
 class frame(object):
     """Represents a parsed frame."""
 
     requestid = attr.ib()
     typeid = attr.ib()
     flags = attr.ib()
     payload = attr.ib()
 
 def makeframe(requestid, typeid, flags, payload):
     """Assemble a frame into a byte array."""
     # TODO assert size of payload.
     frame = bytearray(FRAME_HEADER_SIZE + len(payload))
 
     # 24 bits length
     # 16 bits request id
     # 4 bits type
     # 4 bits flags
 
     l = struct.pack(r'<I', len(payload))
     frame[0:3] = l[0:3]
     struct.pack_into(r'<H', frame, 3, requestid)
     frame[5] = (typeid << 4) | flags
     frame[6:] = payload
 
     return frame
 
 def makeframefromhumanstring(s):
     """Create a frame from a human readable string
 
     Strings have the form:
 
         <request-id> <type> <flags> <payload>
 
     This can be used by user-facing applications and tests for creating
     frames easily without having to type out a bunch of constants.
 
     Request ID is an integer.
 
     Frame type and flags can be specified by integer or named constant.
 
     Flags can be delimited by `|` to bitwise OR them together.
     """
     requestid, frametype, frameflags, payload = s.split(b' ', 3)
 
     requestid = int(requestid)
 
     if frametype in FRAME_TYPES:
         frametype = FRAME_TYPES[frametype]
     else:
         frametype = int(frametype)
 
     finalflags = 0
     validflags = FRAME_TYPE_FLAGS[frametype]
     for flag in frameflags.split(b'|'):
         if flag in validflags:
             finalflags |= validflags[flag]
         else:
             finalflags |= int(flag)
 
     payload = stringutil.unescapestr(payload)
 
     return makeframe(requestid=requestid, typeid=frametype,
                      flags=finalflags, payload=payload)
 
 def parseheader(data):
     """Parse a unified framing protocol frame header from a buffer.
 
     The header is expected to be in the buffer at offset 0 and the
     buffer is expected to be large enough to hold a full header.
     """
     # 24 bits payload length (little endian)
     # 4 bits frame type
     # 4 bits frame flags
     # ... payload
     framelength = data[0] + 256 * data[1] + 16384 * data[2]
     requestid = struct.unpack_from(r'<H', data, 3)[0]
     typeflags = data[5]
 
     frametype = (typeflags & 0xf0) >> 4
     frameflags = typeflags & 0x0f
 
     return frameheader(framelength, requestid, frametype, frameflags)
 
 def readframe(fh):
     """Read a unified framing protocol frame from a file object.
 
     Returns a 3-tuple of (type, flags, payload) for the decoded frame or
     None if no frame is available. May raise if a malformed frame is
     seen.
     """
     header = bytearray(FRAME_HEADER_SIZE)
 
     readcount = fh.readinto(header)
 
     if readcount == 0:
         return None
 
     if readcount != FRAME_HEADER_SIZE:
         raise error.Abort(_('received incomplete frame: got %d bytes: %s') %
                           (readcount, header))
 
     h = parseheader(header)
 
     payload = fh.read(h.length)
     if len(payload) != h.length:
         raise error.Abort(_('frame length error: expected %d; got %d') %
                           (h.length, len(payload)))
 
     return frame(h.requestid, h.typeid, h.flags, payload)
 
-def createcommandframes(requestid, cmd, args, datafh=None):
+def createcommandframes(stream, requestid, cmd, args, datafh=None):
     """Create frames necessary to transmit a request to run a command.
 
     This is a generator of bytearrays. Each item represents a frame
     ready to be sent over the wire to a peer.
     """
     flags = 0
     if args:
         flags |= FLAG_COMMAND_NAME_HAVE_ARGS
     if datafh:
         flags |= FLAG_COMMAND_NAME_HAVE_DATA
 
     if not flags:
         flags |= FLAG_COMMAND_NAME_EOS
 
-    yield makeframe(requestid=requestid, typeid=FRAME_TYPE_COMMAND_NAME,
-                    flags=flags, payload=cmd)
+    yield stream.makeframe(requestid=requestid, typeid=FRAME_TYPE_COMMAND_NAME,
+                           flags=flags, payload=cmd)
 
     for i, k in enumerate(sorted(args)):
         v = args[k]
         last = i == len(args) - 1
 
         # TODO handle splitting of argument values across frames.
         payload = bytearray(ARGUMENT_FRAME_HEADER.size + len(k) + len(v))
         offset = 0
         ARGUMENT_FRAME_HEADER.pack_into(payload, offset, len(k), len(v))
         offset += ARGUMENT_FRAME_HEADER.size
         payload[offset:offset + len(k)] = k
         offset += len(k)
         payload[offset:offset + len(v)] = v
 
         flags = FLAG_COMMAND_ARGUMENT_EOA if last else 0
-        yield makeframe(requestid=requestid,
-                        typeid=FRAME_TYPE_COMMAND_ARGUMENT,
-                        flags=flags,
-                        payload=payload)
+        yield stream.makeframe(requestid=requestid,
+                               typeid=FRAME_TYPE_COMMAND_ARGUMENT,
+                               flags=flags,
+                               payload=payload)
 
     if datafh:
         while True:
             data = datafh.read(DEFAULT_MAX_FRAME_SIZE)
 
             done = False
             if len(data) == DEFAULT_MAX_FRAME_SIZE:
                 flags = FLAG_COMMAND_DATA_CONTINUATION
             else:
                 flags = FLAG_COMMAND_DATA_EOS
                 assert datafh.read(1) == b''
                 done = True
 
-            yield makeframe(requestid=requestid,
-                            typeid=FRAME_TYPE_COMMAND_DATA,
-                            flags=flags,
-                            payload=data)
+            yield stream.makeframe(requestid=requestid,
+                                   typeid=FRAME_TYPE_COMMAND_DATA,
+                                   flags=flags,
+                                   payload=data)
 
             if done:
                 break
 
-def createbytesresponseframesfrombytes(requestid, data,
+def createbytesresponseframesfrombytes(stream, requestid, data,
                                        maxframesize=DEFAULT_MAX_FRAME_SIZE):
     """Create a raw frame to send a bytes response from static bytes input.
 
     Returns a generator of bytearrays.
     """
 
     # Simple case of a single frame.
     if len(data) <= maxframesize:
-        yield makeframe(requestid=requestid,
-                        typeid=FRAME_TYPE_BYTES_RESPONSE,
-                        flags=FLAG_BYTES_RESPONSE_EOS,
-                        payload=data)
+        yield stream.makeframe(requestid=requestid,
+                               typeid=FRAME_TYPE_BYTES_RESPONSE,
+                               flags=FLAG_BYTES_RESPONSE_EOS,
+                               payload=data)
         return
 
     offset = 0
     while True:
         chunk = data[offset:offset + maxframesize]
         offset += len(chunk)
         done = offset == len(data)
 
         if done:
             flags = FLAG_BYTES_RESPONSE_EOS
         else:
             flags = FLAG_BYTES_RESPONSE_CONTINUATION
 
-        yield makeframe(requestid=requestid,
-                        typeid=FRAME_TYPE_BYTES_RESPONSE,
-                        flags=flags,
-                        payload=chunk)
+        yield stream.makeframe(requestid=requestid,
+                               typeid=FRAME_TYPE_BYTES_RESPONSE,
+                               flags=flags,
+                               payload=chunk)
 
         if done:
             break
 
-def createerrorframe(requestid, msg, protocol=False, application=False):
+def createerrorframe(stream, requestid, msg, protocol=False, application=False):
     # TODO properly handle frame size limits.
     assert len(msg) <= DEFAULT_MAX_FRAME_SIZE
 
     flags = 0
     if protocol:
         flags |= FLAG_ERROR_RESPONSE_PROTOCOL
     if application:
         flags |= FLAG_ERROR_RESPONSE_APPLICATION
 
-    yield makeframe(requestid=requestid,
-                    typeid=FRAME_TYPE_ERROR_RESPONSE,
-                    flags=flags,
-                    payload=msg)
+    yield stream.makeframe(requestid=requestid,
+                           typeid=FRAME_TYPE_ERROR_RESPONSE,
+                           flags=flags,
+                           payload=msg)
 
-def createtextoutputframe(requestid, atoms):
+def createtextoutputframe(stream, requestid, atoms):
     """Create a text output frame to render text to people.
 
     ``atoms`` is a 3-tuple of (formatting string, args, labels).
 
     The formatting string contains ``%s`` tokens to be replaced by the
     corresponding indexed entry in ``args``. ``labels`` is an iterable of
     formatters to be applied at rendering time. In terms of the ``ui``
     class, each atom corresponds to a ``ui.write()``.
     """
     bytesleft = DEFAULT_MAX_FRAME_SIZE
     atomchunks = []
 
     for (formatting, args, labels) in atoms:
         if len(args) > 255:
             raise ValueError('cannot use more than 255 formatting arguments')
         if len(labels) > 255:
             raise ValueError('cannot use more than 255 labels')
 
         # TODO look for localstr, other types here?
 
         if not isinstance(formatting, bytes):
             raise ValueError('must use bytes formatting strings')
         for arg in args:
             if not isinstance(arg, bytes):
                 raise ValueError('must use bytes for arguments')
         for label in labels:
             if not isinstance(label, bytes):
                 raise ValueError('must use bytes for labels')
 
         # Formatting string must be UTF-8.
         formatting = formatting.decode(r'utf-8', r'replace').encode(r'utf-8')
 
         # Arguments must be UTF-8.
         args = [a.decode(r'utf-8', r'replace').encode(r'utf-8') for a in args]
 
         # Labels must be ASCII.
         labels = [l.decode(r'ascii', r'strict').encode(r'ascii')
                   for l in labels]
 
         if len(formatting) > 65535:
             raise ValueError('formatting string cannot be longer than 64k')
 
         if any(len(a) > 65535 for a in args):
             raise ValueError('argument string cannot be longer than 64k')
 
         if any(len(l) > 255 for l in labels):
             raise ValueError('label string cannot be longer than 255 bytes')
 
         chunks = [
             struct.pack(r'<H', len(formatting)),
             struct.pack(r'<BB', len(labels), len(args)),
             struct.pack(r'<' + r'B' * len(labels), *map(len, labels)),
             struct.pack(r'<' + r'H' * len(args), *map(len, args)),
         ]
         chunks.append(formatting)
         chunks.extend(labels)
         chunks.extend(args)
 
         atom = b''.join(chunks)
         atomchunks.append(atom)
         bytesleft -= len(atom)
 
     if bytesleft < 0:
         raise ValueError('cannot encode data in a single frame')
 
-    yield makeframe(requestid=requestid,
-                    typeid=FRAME_TYPE_TEXT_OUTPUT,
-                    flags=0,
-                    payload=b''.join(atomchunks))
+    yield stream.makeframe(requestid=requestid,
+                           typeid=FRAME_TYPE_TEXT_OUTPUT,
+                           flags=0,
+                           payload=b''.join(atomchunks))
+
+class stream(object):
+    """Represents a logical unidirectional series of frames."""
+
+    def makeframe(self, requestid, typeid, flags, payload):
+        """Create a frame to be sent out over this stream.
+
+        Only returns the frame instance. Does not actually send it.
+        """
+        return makeframe(requestid, typeid, flags, payload)
 
 class serverreactor(object):
     """Holds state of a server handling frame-based protocol requests.
 
     This class is the "brain" of the unified frame-based protocol server
     component. While the protocol is stateless from the perspective of
     requests/commands, something needs to track which frames have been
     received, what frames to expect, etc. This class is that thing.
 
     Instances are modeled as a state machine of sorts. Instances are also
     reactionary to external events. The point of this class is to encapsulate
     the state of the connection and the exchange of frames, not to perform
     work. Instead, callers tell this class when something occurs, like a
     frame arriving. If that activity is worthy of a follow-up action (say
     *run a command*), the return value of that handler will say so.
 
     I/O and CPU intensive operations are purposefully delegated outside of
     this class.
 
     Consumers are expected to tell instances when events occur. They do so by
     calling the various ``on*`` methods. These methods return a 2-tuple
     describing any follow-up action(s) to take. The first element is the
     name of an action to perform. The second is a data structure (usually
     a dict) specific to that action that contains more information. e.g.
     if the server wants to send frames back to the client, the data structure
     will contain a reference to those frames.
 
     Valid actions that consumers can be instructed to take are:
 
     sendframes
        Indicates that frames should be sent to the client. The ``framegen``
        key contains a generator of frames that should be sent. The server
        assumes that all frames are sent to the client.
 
     error
        Indicates that an error occurred. Consumer should probably abort.
 
     runcommand
        Indicates that the consumer should run a wire protocol command. Details
        of the command to run are given in the data structure.
 
     wantframe
        Indicates that nothing of interest happened and the server is waiting on
        more frames from the client before anything interesting can be done.
 
     noop
        Indicates no additional action is required.
 
     Known Issues
     ------------
 
     There are no limits to the number of partially received commands or their
     size. A malicious client could stream command request data and exhaust the
     server's memory.
 
     Partially received commands are not acted upon when end of input is
     reached. Should the server error if it receives a partial request?
     Should the client send a message to abort a partially transmitted request
     to facilitate graceful shutdown?
 
     Active requests that haven't been responded to aren't tracked. This means
     that if we receive a command and instruct its dispatch, another command
     with its request ID can come in over the wire and there will be a race
     between who responds to what.
     """
 
     def __init__(self, deferoutput=False):
         """Construct a new server reactor.
 
         ``deferoutput`` can be used to indicate that no output frames should be
         instructed to be sent until input has been exhausted. In this mode,
         events that would normally generate output frames (such as a command
         response being ready) will instead defer instructing the consumer to
         send those frames. This is useful for half-duplex transports where the
         sender cannot receive until all data has been transmitted.
         """
         self._deferoutput = deferoutput
         self._state = 'idle'
         self._bufferedframegens = []
         # request id -> dict of commands that are actively being received.
         self._receivingcommands = {}
         # Request IDs that have been received and are actively being processed.
         # Once all output for a request has been sent, it is removed from this
         # set.
         self._activecommands = set()
 
     def onframerecv(self, frame):
         """Process a frame that has been received off the wire.
 
         Returns a dict with an ``action`` key that details what action,
         if any, the consumer should take next.
         """
         handlers = {
             'idle': self._onframeidle,
             'command-receiving': self._onframecommandreceiving,
             'errored': self._onframeerrored,
         }
 
         meth = handlers.get(self._state)
         if not meth:
             raise error.ProgrammingError('unhandled state: %s' % self._state)
 
         return meth(frame)
 
-    def onbytesresponseready(self, requestid, data):
+    def onbytesresponseready(self, stream, requestid, data):
         """Signal that a bytes response is ready to be sent to the client.
 
         The raw bytes response is passed as an argument.
         """
         def sendframes():
-            for frame in createbytesresponseframesfrombytes(requestid, data):
+            for frame in createbytesresponseframesfrombytes(stream, requestid,
+                                                            data):
                 yield frame
 
             self._activecommands.remove(requestid)
 
         result = sendframes()
 
         if self._deferoutput:
             self._bufferedframegens.append(result)
             return 'noop', {}
         else:
             return 'sendframes', {
                 'framegen': result,
             }
 
     def oninputeof(self):
         """Signals that end of input has been received.
 
         No more frames will be received. All pending activity should be
         completed.
         """
         # TODO should we do anything about in-flight commands?
 
         if not self._deferoutput or not self._bufferedframegens:
             return 'noop', {}
 
         # If we buffered all our responses, emit those.
         def makegen():
             for gen in self._bufferedframegens:
                 for frame in gen:
                     yield frame
 
         return 'sendframes', {
             'framegen': makegen(),
         }
 
-    def onapplicationerror(self, requestid, msg):
+    def onapplicationerror(self, stream, requestid, msg):
         return 'sendframes', {
-            'framegen': createerrorframe(requestid, msg, application=True),
+            'framegen': createerrorframe(stream, requestid, msg,
+                                         application=True),
         }
 
     def _makeerrorresult(self, msg):
         return 'error', {
             'message': msg,
         }
 
     def _makeruncommandresult(self, requestid):
         entry = self._receivingcommands[requestid]
         del self._receivingcommands[requestid]
 
         if self._receivingcommands:
             self._state = 'command-receiving'
         else:
             self._state = 'idle'
 
         assert requestid not in self._activecommands
         self._activecommands.add(requestid)
 
         return 'runcommand', {
             'requestid': requestid,
             'command': entry['command'],
             'args': entry['args'],
             'data': entry['data'].getvalue() if entry['data'] else None,
         }
 
     def _makewantframeresult(self):
         return 'wantframe', {
             'state': self._state,
         }
 
     def _onframeidle(self, frame):
         # The only frame type that should be received in this state is a
         # command request.
         if frame.typeid != FRAME_TYPE_COMMAND_NAME:
             self._state = 'errored'
             return self._makeerrorresult(
                 _('expected command frame; got %d') % frame.typeid)
 
         if frame.requestid in self._receivingcommands:
             self._state = 'errored'
             return self._makeerrorresult(
                 _('request with ID %d already received') % frame.requestid)
 
         if frame.requestid in self._activecommands:
             self._state = 'errored'
             return self._makeerrorresult((
                 _('request with ID %d is already active') % frame.requestid))
 
         expectingargs = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_ARGS)
         expectingdata = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_DATA)
 
         self._receivingcommands[frame.requestid] = {
             'command': frame.payload,
             'args': {},
             'data': None,
             'expectingargs': expectingargs,
             'expectingdata': expectingdata,
         }
 
         if frame.flags & FLAG_COMMAND_NAME_EOS:
             return self._makeruncommandresult(frame.requestid)
 
         if expectingargs or expectingdata:
             self._state = 'command-receiving'
             return self._makewantframeresult()
         else:
             self._state = 'errored'
             return self._makeerrorresult(_('missing frame flags on '
                                            'command frame'))
 
     def _onframecommandreceiving(self, frame):
         # It could be a new command request. Process it as such.
         if frame.typeid == FRAME_TYPE_COMMAND_NAME:
             return self._onframeidle(frame)
 
         # All other frames should be related to a command that is currently
         # receiving but is not active.
         if frame.requestid in self._activecommands:
             self._state = 'errored'
             return self._makeerrorresult(
                 _('received frame for request that is still active: %d') %
                 frame.requestid)
 
         if frame.requestid not in self._receivingcommands:
             self._state = 'errored'
             return self._makeerrorresult(
                 _('received frame for request that is not receiving: %d') %
                   frame.requestid)
 
         entry = self._receivingcommands[frame.requestid]
 
         if frame.typeid == FRAME_TYPE_COMMAND_ARGUMENT:
             if not entry['expectingargs']:
                 self._state = 'errored'
                 return self._makeerrorresult(_(
                     'received command argument frame for request that is not '
                     'expecting arguments: %d') % frame.requestid)
 
             return self._handlecommandargsframe(frame, entry)
 
         elif frame.typeid == FRAME_TYPE_COMMAND_DATA:
             if not entry['expectingdata']:
                 self._state = 'errored'
                 return self._makeerrorresult(_(
                     'received command data frame for request that is not '
                     'expecting data: %d') % frame.requestid)
 
             if entry['data'] is None:
                 entry['data'] = util.bytesio()
 
             return self._handlecommanddataframe(frame, entry)
 
     def _handlecommandargsframe(self, frame, entry):
         # The frame and state of command should have already been validated.
         assert frame.typeid == FRAME_TYPE_COMMAND_ARGUMENT
 
         offset = 0
         namesize, valuesize = ARGUMENT_FRAME_HEADER.unpack_from(frame.payload)
         offset += ARGUMENT_FRAME_HEADER.size
 
         # The argument name MUST fit inside the frame.
         argname = bytes(frame.payload[offset:offset + namesize])
         offset += namesize
 
         if len(argname) != namesize:
             self._state = 'errored'
             return self._makeerrorresult(_('malformed argument frame: '
                                            'partial argument name'))
 
         argvalue = bytes(frame.payload[offset:])
 
         # Argument value spans multiple frames. Record our active state
         # and wait for the next frame.
         if frame.flags & FLAG_COMMAND_ARGUMENT_CONTINUATION:
             raise error.ProgrammingError('not yet implemented')
 
         # Common case: the argument value is completely contained in this
         # frame.
 
         if len(argvalue) != valuesize:
             self._state = 'errored'
             return self._makeerrorresult(_('malformed argument frame: '
                                            'partial argument value'))
 
         entry['args'][argname] = argvalue
 
         if frame.flags & FLAG_COMMAND_ARGUMENT_EOA:
             if entry['expectingdata']:
                 # TODO signal request to run a command once we don't
                 # buffer data frames.
                 return self._makewantframeresult()
             else:
                 return self._makeruncommandresult(frame.requestid)
         else:
             return self._makewantframeresult()
 
     def _handlecommanddataframe(self, frame, entry):
         assert frame.typeid == FRAME_TYPE_COMMAND_DATA
 
         # TODO support streaming data instead of buffering it.
         entry['data'].write(frame.payload)
 
         if frame.flags & FLAG_COMMAND_DATA_CONTINUATION:
             return self._makewantframeresult()
         elif frame.flags & FLAG_COMMAND_DATA_EOS:
             entry['data'].seek(0)
             return self._makeruncommandresult(frame.requestid)
         else:
             self._state = 'errored'
             return self._makeerrorresult(_('command data frame without '
                                            'flags'))
 
     def _onframeerrored(self, frame):
         return self._makeerrorresult(_('server already errored'))

mercurial/wireprotoserver.py

@@ -1,1047 +1,1049
 # 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-0002'
 
 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:
         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)
         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
 
     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, 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,
                       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):
             # 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']]
         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
+    stream = wireprotoframing.stream()
 
     if isinstance(rsp, wireprototypes.bytesresponse):
-        action, meta = reactor.onbytesresponseready(command['requestid'],
+        action, meta = reactor.onbytesresponseready(stream,
+                                                    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-serverreactor.py

@@ -1,639 +1,678
 from __future__ import absolute_import, print_function
 
 import unittest
 
 from mercurial import (
     util,
     wireprotoframing as framing,
 )
 
 ffs = framing.makeframefromhumanstring
 
 def makereactor(deferoutput=False):
     return framing.serverreactor(deferoutput=deferoutput)
 
 def sendframes(reactor, gen):
     """Send a generator of frame bytearray to a reactor.
 
     Emits a generator of results from ``onframerecv()`` calls.
     """
     for frame in gen:
         header = framing.parseheader(frame)
         payload = frame[framing.FRAME_HEADER_SIZE:]
         assert len(payload) == header.length
 
         yield reactor.onframerecv(framing.frame(header.requestid,
                                                 header.typeid,
                                                 header.flags,
                                                 payload))
 
-def sendcommandframes(reactor, rid, cmd, args, datafh=None):
+def sendcommandframes(reactor, stream, rid, cmd, args, datafh=None):
     """Generate frames to run a command and send them to a reactor."""
     return sendframes(reactor,
-                      framing.createcommandframes(rid, cmd, args, datafh))
+                      framing.createcommandframes(stream, rid, cmd, args,
+                                                  datafh))
 
 class FrameTests(unittest.TestCase):
     def testdataexactframesize(self):
         data = util.bytesio(b'x' * framing.DEFAULT_MAX_FRAME_SIZE)
 
-        frames = list(framing.createcommandframes(1, b'command', {}, data))
+        stream = framing.stream()
+        frames = list(framing.createcommandframes(stream, 1, b'command',
+                                                  {}, data))
         self.assertEqual(frames, [
             ffs(b'1 command-name have-data command'),
             ffs(b'1 command-data continuation %s' % data.getvalue()),
             ffs(b'1 command-data eos ')
         ])
 
     def testdatamultipleframes(self):
         data = util.bytesio(b'x' * (framing.DEFAULT_MAX_FRAME_SIZE + 1))
-        frames = list(framing.createcommandframes(1, b'command', {}, data))
+
+        stream = framing.stream()
+        frames = list(framing.createcommandframes(stream, 1, b'command', {},
+                                                  data))
         self.assertEqual(frames, [
             ffs(b'1 command-name have-data command'),
             ffs(b'1 command-data continuation %s' % (
                 b'x' * framing.DEFAULT_MAX_FRAME_SIZE)),
             ffs(b'1 command-data eos x'),
         ])
 
     def testargsanddata(self):
         data = util.bytesio(b'x' * 100)
 
-        frames = list(framing.createcommandframes(1, b'command', {
+        stream = framing.stream()
+        frames = list(framing.createcommandframes(stream, 1, b'command', {
             b'key1': b'key1value',
             b'key2': b'key2value',
             b'key3': b'key3value',
         }, data))
 
         self.assertEqual(frames, [
             ffs(b'1 command-name have-args|have-data command'),
             ffs(br'1 command-argument 0 \x04\x00\x09\x00key1key1value'),
             ffs(br'1 command-argument 0 \x04\x00\x09\x00key2key2value'),
             ffs(br'1 command-argument eoa \x04\x00\x09\x00key3key3value'),
             ffs(b'1 command-data eos %s' % data.getvalue()),
         ])
 
     def testtextoutputexcessiveargs(self):
         """At most 255 formatting arguments are allowed."""
         with self.assertRaisesRegexp(ValueError,
                                      'cannot use more than 255 formatting'):
             args = [b'x' for i in range(256)]
-            list(framing.createtextoutputframe(1, [(b'bleh', args, [])]))
+            list(framing.createtextoutputframe(None, 1,
+                                               [(b'bleh', args, [])]))
 
     def testtextoutputexcessivelabels(self):
         """At most 255 labels are allowed."""
         with self.assertRaisesRegexp(ValueError,
                                      'cannot use more than 255 labels'):
             labels = [b'l' for i in range(256)]
-            list(framing.createtextoutputframe(1, [(b'bleh', [], labels)]))
+            list(framing.createtextoutputframe(None, 1,
+                                               [(b'bleh', [], labels)]))
 
     def testtextoutputformattingstringtype(self):
         """Formatting string must be bytes."""
         with self.assertRaisesRegexp(ValueError, 'must use bytes formatting '):
-            list(framing.createtextoutputframe(1, [
+            list(framing.createtextoutputframe(None, 1, [
                 (b'foo'.decode('ascii'), [], [])]))
 
     def testtextoutputargumentbytes(self):
         with self.assertRaisesRegexp(ValueError, 'must use bytes for argument'):
-            list(framing.createtextoutputframe(1, [
+            list(framing.createtextoutputframe(None, 1, [
                 (b'foo', [b'foo'.decode('ascii')], [])]))
 
     def testtextoutputlabelbytes(self):
         with self.assertRaisesRegexp(ValueError, 'must use bytes for labels'):
-            list(framing.createtextoutputframe(1, [
+            list(framing.createtextoutputframe(None, 1, [
                 (b'foo', [], [b'foo'.decode('ascii')])]))
 
     def testtextoutputtoolongformatstring(self):
         with self.assertRaisesRegexp(ValueError,
                                      'formatting string cannot be longer than'):
-            list(framing.createtextoutputframe(1, [
+            list(framing.createtextoutputframe(None, 1, [
                 (b'x' * 65536, [], [])]))
 
     def testtextoutputtoolongargumentstring(self):
         with self.assertRaisesRegexp(ValueError,
                                      'argument string cannot be longer than'):
-            list(framing.createtextoutputframe(1, [
+            list(framing.createtextoutputframe(None, 1, [
                 (b'bleh', [b'x' * 65536], [])]))
 
     def testtextoutputtoolonglabelstring(self):
         with self.assertRaisesRegexp(ValueError,
                                      'label string cannot be longer than'):
-            list(framing.createtextoutputframe(1, [
+            list(framing.createtextoutputframe(None, 1, [
                 (b'bleh', [], [b'x' * 65536])]))
 
     def testtextoutput1simpleatom(self):
-        val = list(framing.createtextoutputframe(1, [
+        stream = framing.stream()
+        val = list(framing.createtextoutputframe(stream, 1, [
             (b'foo', [], [])]))
 
         self.assertEqual(val, [
             ffs(br'1 text-output 0 \x03\x00\x00\x00foo'),
         ])
 
     def testtextoutput2simpleatoms(self):
-        val = list(framing.createtextoutputframe(1, [
+        stream = framing.stream()
+        val = list(framing.createtextoutputframe(stream, 1, [
             (b'foo', [], []),
             (b'bar', [], []),
         ]))
 
         self.assertEqual(val, [
             ffs(br'1 text-output 0 \x03\x00\x00\x00foo\x03\x00\x00\x00bar'),
         ])
 
     def testtextoutput1arg(self):
-        val = list(framing.createtextoutputframe(1, [
+        stream = framing.stream()
+        val = list(framing.createtextoutputframe(stream, 1, [
             (b'foo %s', [b'val1'], []),
         ]))
 
         self.assertEqual(val, [
             ffs(br'1 text-output 0 \x06\x00\x00\x01\x04\x00foo %sval1'),
         ])
 
     def testtextoutput2arg(self):
-        val = list(framing.createtextoutputframe(1, [
+        stream = framing.stream()
+        val = list(framing.createtextoutputframe(stream, 1, [
             (b'foo %s %s', [b'val', b'value'], []),
         ]))
 
         self.assertEqual(val, [
             ffs(br'1 text-output 0 \x09\x00\x00\x02\x03\x00\x05\x00'
                 br'foo %s %svalvalue'),
         ])
 
     def testtextoutput1label(self):
-        val = list(framing.createtextoutputframe(1, [
+        stream = framing.stream()
+        val = list(framing.createtextoutputframe(stream, 1, [
             (b'foo', [], [b'label']),
         ]))
 
         self.assertEqual(val, [
             ffs(br'1 text-output 0 \x03\x00\x01\x00\x05foolabel'),
         ])
 
     def testargandlabel(self):
-        val = list(framing.createtextoutputframe(1, [
+        stream = framing.stream()
+        val = list(framing.createtextoutputframe(stream, 1, [
             (b'foo %s', [b'arg'], [b'label']),
         ]))
 
         self.assertEqual(val, [
             ffs(br'1 text-output 0 \x06\x00\x01\x01\x05\x03\x00foo %slabelarg'),
         ])
 
 class ServerReactorTests(unittest.TestCase):
     def _sendsingleframe(self, reactor, f):
         results = list(sendframes(reactor, [f]))
         self.assertEqual(len(results), 1)
 
         return results[0]
 
     def assertaction(self, res, expected):
         self.assertIsInstance(res, tuple)
         self.assertEqual(len(res), 2)
         self.assertIsInstance(res[1], dict)
         self.assertEqual(res[0], expected)
 
     def assertframesequal(self, frames, framestrings):
         expected = [ffs(s) for s in framestrings]
         self.assertEqual(list(frames), expected)
 
     def test1framecommand(self):
         """Receiving a command in a single frame yields request to run it."""
         reactor = makereactor()
-        results = list(sendcommandframes(reactor, 1, b'mycommand', {}))
+        stream = framing.stream()
+        results = list(sendcommandframes(reactor, stream, 1, b'mycommand', {}))
         self.assertEqual(len(results), 1)
         self.assertaction(results[0], 'runcommand')
         self.assertEqual(results[0][1], {
             'requestid': 1,
             'command': b'mycommand',
             'args': {},
             'data': None,
         })
 
         result = reactor.oninputeof()
         self.assertaction(result, 'noop')
 
     def test1argument(self):
         reactor = makereactor()
-        results = list(sendcommandframes(reactor, 41, b'mycommand',
+        stream = framing.stream()
+        results = list(sendcommandframes(reactor, stream, 41, b'mycommand',
                                          {b'foo': b'bar'}))
         self.assertEqual(len(results), 2)
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'runcommand')
         self.assertEqual(results[1][1], {
             'requestid': 41,
             'command': b'mycommand',
             'args': {b'foo': b'bar'},
             'data': None,
         })
 
     def testmultiarguments(self):
         reactor = makereactor()
-        results = list(sendcommandframes(reactor, 1, b'mycommand',
+        stream = framing.stream()
+        results = list(sendcommandframes(reactor, stream, 1, b'mycommand',
                                          {b'foo': b'bar', b'biz': b'baz'}))
         self.assertEqual(len(results), 3)
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'wantframe')
         self.assertaction(results[2], 'runcommand')
         self.assertEqual(results[2][1], {
             'requestid': 1,
             'command': b'mycommand',
             'args': {b'foo': b'bar', b'biz': b'baz'},
             'data': None,
         })
 
     def testsimplecommanddata(self):
         reactor = makereactor()
-        results = list(sendcommandframes(reactor, 1, b'mycommand', {},
+        stream = framing.stream()
+        results = list(sendcommandframes(reactor, stream, 1, b'mycommand', {},
                                          util.bytesio(b'data!')))
         self.assertEqual(len(results), 2)
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'runcommand')
         self.assertEqual(results[1][1], {
             'requestid': 1,
             'command': b'mycommand',
             'args': {},
             'data': b'data!',
         })
 
     def testmultipledataframes(self):
         frames = [
             ffs(b'1 command-name have-data mycommand'),
             ffs(b'1 command-data continuation data1'),
             ffs(b'1 command-data continuation data2'),
             ffs(b'1 command-data eos data3'),
         ]
 
         reactor = makereactor()
         results = list(sendframes(reactor, frames))
         self.assertEqual(len(results), 4)
         for i in range(3):
             self.assertaction(results[i], 'wantframe')
         self.assertaction(results[3], 'runcommand')
         self.assertEqual(results[3][1], {
             'requestid': 1,
             'command': b'mycommand',
             'args': {},
             'data': b'data1data2data3',
         })
 
     def testargumentanddata(self):
         frames = [
             ffs(b'1 command-name have-args|have-data command'),
             ffs(br'1 command-argument 0 \x03\x00\x03\x00keyval'),
             ffs(br'1 command-argument eoa \x03\x00\x03\x00foobar'),
             ffs(b'1 command-data continuation value1'),
             ffs(b'1 command-data eos value2'),
         ]
 
         reactor = makereactor()
         results = list(sendframes(reactor, frames))
 
         self.assertaction(results[-1], 'runcommand')
         self.assertEqual(results[-1][1], {
             'requestid': 1,
             'command': b'command',
             'args': {
                 b'key': b'val',
                 b'foo': b'bar',
             },
             'data': b'value1value2',
         })
 
     def testunexpectedcommandargument(self):
         """Command argument frame when not running a command is an error."""
         result = self._sendsingleframe(makereactor(),
                                        ffs(b'1 command-argument 0 ignored'))
         self.assertaction(result, 'error')
         self.assertEqual(result[1], {
             'message': b'expected command frame; got 2',
         })
 
     def testunexpectedcommandargumentreceiving(self):
         """Same as above but the command is receiving."""
         results = list(sendframes(makereactor(), [
             ffs(b'1 command-name have-data command'),
             ffs(b'1 command-argument eoa ignored'),
         ]))
 
         self.assertaction(results[1], 'error')
         self.assertEqual(results[1][1], {
             'message': b'received command argument frame for request that is '
                        b'not expecting arguments: 1',
         })
 
     def testunexpectedcommanddata(self):
         """Command argument frame when not running a command is an error."""
         result = self._sendsingleframe(makereactor(),
                                        ffs(b'1 command-data 0 ignored'))
         self.assertaction(result, 'error')
         self.assertEqual(result[1], {
             'message': b'expected command frame; got 3',
         })
 
     def testunexpectedcommanddatareceiving(self):
         """Same as above except the command is receiving."""
         results = list(sendframes(makereactor(), [
             ffs(b'1 command-name have-args command'),
             ffs(b'1 command-data eos ignored'),
         ]))
 
         self.assertaction(results[1], 'error')
         self.assertEqual(results[1][1], {
             'message': b'received command data frame for request that is not '
                        b'expecting data: 1',
         })
 
     def testmissingcommandframeflags(self):
         """Command name frame must have flags set."""
         result = self._sendsingleframe(makereactor(),
                                        ffs(b'1 command-name 0 command'))
         self.assertaction(result, 'error')
         self.assertEqual(result[1], {
             'message': b'missing frame flags on command frame',
         })
 
     def testconflictingrequestidallowed(self):
         """Multiple fully serviced commands with same request ID is allowed."""
         reactor = makereactor()
         results = []
+        outstream = framing.stream()
         results.append(self._sendsingleframe(
             reactor, ffs(b'1 command-name eos command')))
-        result = reactor.onbytesresponseready(1, b'response1')
+        result = reactor.onbytesresponseready(outstream, 1, b'response1')
         self.assertaction(result, 'sendframes')
         list(result[1]['framegen'])
         results.append(self._sendsingleframe(
             reactor, ffs(b'1 command-name eos command')))
-        result = reactor.onbytesresponseready(1, b'response2')
+        result = reactor.onbytesresponseready(outstream, 1, b'response2')
         self.assertaction(result, 'sendframes')
         list(result[1]['framegen'])
         results.append(self._sendsingleframe(
             reactor, ffs(b'1 command-name eos command')))
-        result = reactor.onbytesresponseready(1, b'response3')
+        result = reactor.onbytesresponseready(outstream, 1, b'response3')
         self.assertaction(result, 'sendframes')
         list(result[1]['framegen'])
 
         for i in range(3):
             self.assertaction(results[i], 'runcommand')
             self.assertEqual(results[i][1], {
                 'requestid': 1,
                 'command': b'command',
                 'args': {},
                 'data': None,
             })
 
     def testconflictingrequestid(self):
         """Request ID for new command matching in-flight command is illegal."""
         results = list(sendframes(makereactor(), [
             ffs(b'1 command-name have-args command'),
             ffs(b'1 command-name eos command'),
         ]))
 
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'error')
         self.assertEqual(results[1][1], {
             'message': b'request with ID 1 already received',
         })
 
     def testinterleavedcommands(self):
         results = list(sendframes(makereactor(), [
             ffs(b'1 command-name have-args command1'),
             ffs(b'3 command-name have-args command3'),
             ffs(br'1 command-argument 0 \x03\x00\x03\x00foobar'),
             ffs(br'3 command-argument 0 \x03\x00\x03\x00bizbaz'),
             ffs(br'3 command-argument eoa \x03\x00\x03\x00keyval'),
             ffs(br'1 command-argument eoa \x04\x00\x03\x00key1val'),
         ]))
 
         self.assertEqual([t[0] for t in results], [
             'wantframe',
             'wantframe',
             'wantframe',
             'wantframe',
             'runcommand',
             'runcommand',
         ])
 
         self.assertEqual(results[4][1], {
             'requestid': 3,
             'command': 'command3',
             'args': {b'biz': b'baz', b'key': b'val'},
             'data': None,
         })
         self.assertEqual(results[5][1], {
             'requestid': 1,
             'command': 'command1',
             'args': {b'foo': b'bar', b'key1': b'val'},
             'data': None,
         })
 
     def testmissingargumentframe(self):
         # This test attempts to test behavior when reactor has an incomplete
         # command request waiting on argument data. But it doesn't handle that
         # scenario yet. So this test does nothing of value.
         frames = [
             ffs(b'1 command-name have-args command'),
         ]
 
         results = list(sendframes(makereactor(), frames))
         self.assertaction(results[0], 'wantframe')
 
     def testincompleteargumentname(self):
         """Argument frame with incomplete name."""
         frames = [
             ffs(b'1 command-name have-args command1'),
             ffs(br'1 command-argument eoa \x04\x00\xde\xadfoo'),
         ]
 
         results = list(sendframes(makereactor(), frames))
         self.assertEqual(len(results), 2)
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'error')
         self.assertEqual(results[1][1], {
             'message': b'malformed argument frame: partial argument name',
         })
 
     def testincompleteargumentvalue(self):
         """Argument frame with incomplete value."""
         frames = [
             ffs(b'1 command-name have-args command'),
             ffs(br'1 command-argument eoa \x03\x00\xaa\xaafoopartialvalue'),
         ]
 
         results = list(sendframes(makereactor(), frames))
         self.assertEqual(len(results), 2)
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'error')
         self.assertEqual(results[1][1], {
             'message': b'malformed argument frame: partial argument value',
         })
 
     def testmissingcommanddataframe(self):
         # The reactor doesn't currently handle partially received commands.
         # So this test is failing to do anything with request 1.
         frames = [
             ffs(b'1 command-name have-data command1'),
             ffs(b'3 command-name eos command2'),
         ]
         results = list(sendframes(makereactor(), frames))
         self.assertEqual(len(results), 2)
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'runcommand')
 
     def testmissingcommanddataframeflags(self):
         frames = [
             ffs(b'1 command-name have-data command1'),
             ffs(b'1 command-data 0 data'),
         ]
         results = list(sendframes(makereactor(), frames))
         self.assertEqual(len(results), 2)
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'error')
         self.assertEqual(results[1][1], {
             'message': b'command data frame without flags',
         })
 
     def testframefornonreceivingrequest(self):
         """Receiving a frame for a command that is not receiving is illegal."""
         results = list(sendframes(makereactor(), [
             ffs(b'1 command-name eos command1'),
             ffs(b'3 command-name have-data command3'),
             ffs(b'5 command-argument eoa ignored'),
         ]))
         self.assertaction(results[2], 'error')
         self.assertEqual(results[2][1], {
             'message': b'received frame for request that is not receiving: 5',
         })
 
     def testsimpleresponse(self):
         """Bytes response to command sends result frames."""
         reactor = makereactor()
-        list(sendcommandframes(reactor, 1, b'mycommand', {}))
+        instream = framing.stream()
+        list(sendcommandframes(reactor, instream, 1, b'mycommand', {}))
 
-        result = reactor.onbytesresponseready(1, b'response')
+        outstream = framing.stream()
+        result = reactor.onbytesresponseready(outstream, 1, b'response')
         self.assertaction(result, 'sendframes')
         self.assertframesequal(result[1]['framegen'], [
             b'1 bytes-response eos response',
         ])
 
     def testmultiframeresponse(self):
         """Bytes response spanning multiple frames is handled."""
         first = b'x' * framing.DEFAULT_MAX_FRAME_SIZE
         second = b'y' * 100
 
         reactor = makereactor()
-        list(sendcommandframes(reactor, 1, b'mycommand', {}))
+        instream = framing.stream()
+        list(sendcommandframes(reactor, instream, 1, b'mycommand', {}))
 
-        result = reactor.onbytesresponseready(1, first + second)
+        outstream = framing.stream()
+        result = reactor.onbytesresponseready(outstream, 1, first + second)
         self.assertaction(result, 'sendframes')
         self.assertframesequal(result[1]['framegen'], [
             b'1 bytes-response continuation %s' % first,
             b'1 bytes-response eos %s' % second,
         ])
 
     def testapplicationerror(self):
         reactor = makereactor()
-        list(sendcommandframes(reactor, 1, b'mycommand', {}))
+        instream = framing.stream()
+        list(sendcommandframes(reactor, instream, 1, b'mycommand', {}))
 
-        result = reactor.onapplicationerror(1, b'some message')
+        outstream = framing.stream()
+        result = reactor.onapplicationerror(outstream, 1, b'some message')
         self.assertaction(result, 'sendframes')
         self.assertframesequal(result[1]['framegen'], [
             b'1 error-response application some message',
         ])
 
     def test1commanddeferresponse(self):
         """Responses when in deferred output mode are delayed until EOF."""
         reactor = makereactor(deferoutput=True)
-        results = list(sendcommandframes(reactor, 1, b'mycommand', {}))
+        instream = framing.stream()
+        results = list(sendcommandframes(reactor, instream, 1, b'mycommand',
+                                         {}))
         self.assertEqual(len(results), 1)
         self.assertaction(results[0], 'runcommand')
 
-        result = reactor.onbytesresponseready(1, b'response')
+        outstream = framing.stream()
+        result = reactor.onbytesresponseready(outstream, 1, b'response')
         self.assertaction(result, 'noop')
         result = reactor.oninputeof()
         self.assertaction(result, 'sendframes')
         self.assertframesequal(result[1]['framegen'], [
             b'1 bytes-response eos response',
         ])
 
     def testmultiplecommanddeferresponse(self):
         reactor = makereactor(deferoutput=True)
-        list(sendcommandframes(reactor, 1, b'command1', {}))
-        list(sendcommandframes(reactor, 3, b'command2', {}))
+        instream = framing.stream()
+        list(sendcommandframes(reactor, instream, 1, b'command1', {}))
+        list(sendcommandframes(reactor, instream, 3, b'command2', {}))
 
-        result = reactor.onbytesresponseready(1, b'response1')
+        outstream = framing.stream()
+        result = reactor.onbytesresponseready(outstream, 1, b'response1')
         self.assertaction(result, 'noop')
-        result = reactor.onbytesresponseready(3, b'response2')
+        result = reactor.onbytesresponseready(outstream, 3, b'response2')
         self.assertaction(result, 'noop')
         result = reactor.oninputeof()
         self.assertaction(result, 'sendframes')
         self.assertframesequal(result[1]['framegen'], [
             b'1 bytes-response eos response1',
             b'3 bytes-response eos response2'
         ])
 
     def testrequestidtracking(self):
         reactor = makereactor(deferoutput=True)
-        list(sendcommandframes(reactor, 1, b'command1', {}))
-        list(sendcommandframes(reactor, 3, b'command2', {}))
-        list(sendcommandframes(reactor, 5, b'command3', {}))
+        instream = framing.stream()
+        list(sendcommandframes(reactor, instream, 1, b'command1', {}))
+        list(sendcommandframes(reactor, instream, 3, b'command2', {}))
+        list(sendcommandframes(reactor, instream, 5, b'command3', {}))
 
         # Register results for commands out of order.
-        reactor.onbytesresponseready(3, b'response3')
-        reactor.onbytesresponseready(1, b'response1')
-        reactor.onbytesresponseready(5, b'response5')
+        outstream = framing.stream()
+        reactor.onbytesresponseready(outstream, 3, b'response3')
+        reactor.onbytesresponseready(outstream, 1, b'response1')
+        reactor.onbytesresponseready(outstream, 5, b'response5')
 
         result = reactor.oninputeof()
         self.assertaction(result, 'sendframes')
         self.assertframesequal(result[1]['framegen'], [
             b'3 bytes-response eos response3',
             b'1 bytes-response eos response1',
             b'5 bytes-response eos response5',
         ])
 
     def testduplicaterequestonactivecommand(self):
         """Receiving a request ID that matches a request that isn't finished."""
         reactor = makereactor()
-        list(sendcommandframes(reactor, 1, b'command1', {}))
-        results = list(sendcommandframes(reactor, 1, b'command1', {}))
+        stream = framing.stream()
+        list(sendcommandframes(reactor, stream, 1, b'command1', {}))
+        results = list(sendcommandframes(reactor, stream, 1, b'command1', {}))
 
         self.assertaction(results[0], 'error')
         self.assertEqual(results[0][1], {
             'message': b'request with ID 1 is already active',
         })
 
     def testduplicaterequestonactivecommandnosend(self):
         """Same as above but we've registered a response but haven't sent it."""
         reactor = makereactor()
-        list(sendcommandframes(reactor, 1, b'command1', {}))
-        reactor.onbytesresponseready(1, b'response')
+        instream = framing.stream()
+        list(sendcommandframes(reactor, instream, 1, b'command1', {}))
+        outstream = framing.stream()
+        reactor.onbytesresponseready(outstream, 1, b'response')
 
         # We've registered the response but haven't sent it. From the
         # perspective of the reactor, the command is still active.
 
-        results = list(sendcommandframes(reactor, 1, b'command1', {}))
+        results = list(sendcommandframes(reactor, instream, 1, b'command1', {}))
         self.assertaction(results[0], 'error')
         self.assertEqual(results[0][1], {
             'message': b'request with ID 1 is already active',
         })
 
     def testduplicaterequestargumentframe(self):
         """Variant on above except we sent an argument frame instead of name."""
         reactor = makereactor()
-        list(sendcommandframes(reactor, 1, b'command', {}))
+        stream = framing.stream()
+        list(sendcommandframes(reactor, stream, 1, b'command', {}))
         results = list(sendframes(reactor, [
             ffs(b'3 command-name have-args command'),
             ffs(b'1 command-argument 0 ignored'),
         ]))
         self.assertaction(results[0], 'wantframe')
         self.assertaction(results[1], 'error')
         self.assertEqual(results[1][1], {
             'message': 'received frame for request that is still active: 1',
         })
 
     def testduplicaterequestaftersend(self):
         """We can use a duplicate request ID after we've sent the response."""
         reactor = makereactor()
-        list(sendcommandframes(reactor, 1, b'command1', {}))
-        res = reactor.onbytesresponseready(1, b'response')
+        instream = framing.stream()
+        list(sendcommandframes(reactor, instream, 1, b'command1', {}))
+        outstream = framing.stream()
+        res = reactor.onbytesresponseready(outstream, 1, b'response')
         list(res[1]['framegen'])
 
-        results = list(sendcommandframes(reactor, 1, b'command1', {}))
+        results = list(sendcommandframes(reactor, instream, 1, b'command1', {}))
         self.assertaction(results[0], 'runcommand')
 
 if __name__ == '__main__':
     import silenttestrunner
     silenttestrunner.main(__name__)