wireprotoframing.py
813 lines
| 27.5 KiB
| text/x-python
|
PythonLexer
/ mercurial / wireprotoframing.py
|
r37069 | # wireprotoframing.py - unified framing protocol for wire protocol | ||
# | ||||
# Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com> | ||||
# | ||||
# 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 | ||||
|
r37070 | from .i18n import _ | ||
|
r37079 | from .thirdparty import ( | ||
attr, | ||||
) | ||||
|
r37069 | from . import ( | ||
|
r37070 | error, | ||
|
r37069 | util, | ||
) | ||||
|
r37102 | from .utils import ( | ||
stringutil, | ||||
) | ||||
|
r37069 | |||
|
r37304 | FRAME_HEADER_SIZE = 8 | ||
|
r37069 | DEFAULT_MAX_FRAME_SIZE = 32768 | ||
|
r37304 | STREAM_FLAG_BEGIN_STREAM = 0x01 | ||
STREAM_FLAG_END_STREAM = 0x02 | ||||
STREAM_FLAG_ENCODING_APPLIED = 0x04 | ||||
STREAM_FLAGS = { | ||||
b'stream-begin': STREAM_FLAG_BEGIN_STREAM, | ||||
b'stream-end': STREAM_FLAG_END_STREAM, | ||||
b'encoded': STREAM_FLAG_ENCODING_APPLIED, | ||||
} | ||||
|
r37069 | FRAME_TYPE_COMMAND_NAME = 0x01 | ||
FRAME_TYPE_COMMAND_ARGUMENT = 0x02 | ||||
FRAME_TYPE_COMMAND_DATA = 0x03 | ||||
|
r37073 | FRAME_TYPE_BYTES_RESPONSE = 0x04 | ||
FRAME_TYPE_ERROR_RESPONSE = 0x05 | ||||
|
r37078 | FRAME_TYPE_TEXT_OUTPUT = 0x06 | ||
|
r37304 | FRAME_TYPE_STREAM_SETTINGS = 0x08 | ||
|
r37069 | |||
FRAME_TYPES = { | ||||
b'command-name': FRAME_TYPE_COMMAND_NAME, | ||||
b'command-argument': FRAME_TYPE_COMMAND_ARGUMENT, | ||||
b'command-data': FRAME_TYPE_COMMAND_DATA, | ||||
|
r37073 | b'bytes-response': FRAME_TYPE_BYTES_RESPONSE, | ||
b'error-response': FRAME_TYPE_ERROR_RESPONSE, | ||||
|
r37078 | b'text-output': FRAME_TYPE_TEXT_OUTPUT, | ||
|
r37304 | b'stream-settings': FRAME_TYPE_STREAM_SETTINGS, | ||
|
r37069 | } | ||
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, | ||||
} | ||||
|
r37073 | 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, | ||||
} | ||||
|
r37069 | # 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, | ||||
|
r37073 | FRAME_TYPE_BYTES_RESPONSE: FLAGS_BYTES_RESPONSE, | ||
FRAME_TYPE_ERROR_RESPONSE: FLAGS_ERROR_RESPONSE, | ||||
|
r37078 | FRAME_TYPE_TEXT_OUTPUT: {}, | ||
|
r37304 | FRAME_TYPE_STREAM_SETTINGS: {}, | ||
|
r37069 | } | ||
ARGUMENT_FRAME_HEADER = struct.Struct(r'<HH') | ||||
|
r37079 | @attr.s(slots=True) | ||
class frameheader(object): | ||||
"""Represents the data in a frame header.""" | ||||
length = attr.ib() | ||||
requestid = attr.ib() | ||||
|
r37304 | streamid = attr.ib() | ||
streamflags = attr.ib() | ||||
|
r37079 | typeid = attr.ib() | ||
flags = attr.ib() | ||||
@attr.s(slots=True) | ||||
class frame(object): | ||||
"""Represents a parsed frame.""" | ||||
requestid = attr.ib() | ||||
|
r37304 | streamid = attr.ib() | ||
streamflags = attr.ib() | ||||
|
r37079 | typeid = attr.ib() | ||
flags = attr.ib() | ||||
payload = attr.ib() | ||||
|
r37304 | def makeframe(requestid, streamid, streamflags, typeid, flags, payload): | ||
|
r37069 | """Assemble a frame into a byte array.""" | ||
# TODO assert size of payload. | ||||
frame = bytearray(FRAME_HEADER_SIZE + len(payload)) | ||||
|
r37075 | # 24 bits length | ||
# 16 bits request id | ||||
|
r37304 | # 8 bits stream id | ||
# 8 bits stream flags | ||||
|
r37075 | # 4 bits type | ||
# 4 bits flags | ||||
|
r37069 | l = struct.pack(r'<I', len(payload)) | ||
frame[0:3] = l[0:3] | ||||
|
r37304 | struct.pack_into(r'<HBB', frame, 3, requestid, streamid, streamflags) | ||
frame[7] = (typeid << 4) | flags | ||||
frame[8:] = payload | ||||
|
r37069 | |||
return frame | ||||
def makeframefromhumanstring(s): | ||||
|
r37075 | """Create a frame from a human readable string | ||
Strings have the form: | ||||
|
r37304 | <request-id> <stream-id> <stream-flags> <type> <flags> <payload> | ||
|
r37069 | |||
This can be used by user-facing applications and tests for creating | ||||
frames easily without having to type out a bunch of constants. | ||||
|
r37304 | Request ID and stream IDs are integers. | ||
|
r37075 | |||
|
r37304 | Stream flags, frame type, and flags can be specified by integer or | ||
named constant. | ||||
|
r37075 | |||
|
r37069 | Flags can be delimited by `|` to bitwise OR them together. | ||
""" | ||||
|
r37304 | fields = s.split(b' ', 5) | ||
requestid, streamid, streamflags, frametype, frameflags, payload = fields | ||||
|
r37075 | |||
requestid = int(requestid) | ||||
|
r37304 | streamid = int(streamid) | ||
finalstreamflags = 0 | ||||
for flag in streamflags.split(b'|'): | ||||
if flag in STREAM_FLAGS: | ||||
finalstreamflags |= STREAM_FLAGS[flag] | ||||
else: | ||||
finalstreamflags |= int(flag) | ||||
|
r37069 | |||
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) | ||||
|
r37102 | payload = stringutil.unescapestr(payload) | ||
|
r37069 | |||
|
r37304 | return makeframe(requestid=requestid, streamid=streamid, | ||
streamflags=finalstreamflags, typeid=frametype, | ||||
|
r37080 | flags=finalflags, payload=payload) | ||
|
r37069 | |||
|
r37070 | 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) | ||||
|
r37304 | # 16 bits request ID | ||
# 8 bits stream ID | ||||
# 8 bits stream flags | ||||
|
r37070 | # 4 bits frame type | ||
# 4 bits frame flags | ||||
# ... payload | ||||
framelength = data[0] + 256 * data[1] + 16384 * data[2] | ||||
|
r37304 | requestid, streamid, streamflags = struct.unpack_from(r'<HBB', data, 3) | ||
typeflags = data[7] | ||||
|
r37070 | |||
frametype = (typeflags & 0xf0) >> 4 | ||||
frameflags = typeflags & 0x0f | ||||
|
r37304 | return frameheader(framelength, requestid, streamid, streamflags, | ||
frametype, frameflags) | ||||
|
r37070 | |||
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)) | ||||
|
r37079 | h = parseheader(header) | ||
|
r37070 | |||
|
r37079 | payload = fh.read(h.length) | ||
if len(payload) != h.length: | ||||
|
r37070 | raise error.Abort(_('frame length error: expected %d; got %d') % | ||
|
r37079 | (h.length, len(payload))) | ||
|
r37070 | |||
|
r37304 | return frame(h.requestid, h.streamid, h.streamflags, h.typeid, h.flags, | ||
payload) | ||||
|
r37070 | |||
|
r37303 | def createcommandframes(stream, requestid, cmd, args, datafh=None): | ||
|
r37069 | """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 | ||||
|
r37303 | yield stream.makeframe(requestid=requestid, typeid=FRAME_TYPE_COMMAND_NAME, | ||
flags=flags, payload=cmd) | ||||
|
r37069 | |||
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 | ||||
|
r37303 | yield stream.makeframe(requestid=requestid, | ||
typeid=FRAME_TYPE_COMMAND_ARGUMENT, | ||||
flags=flags, | ||||
payload=payload) | ||||
|
r37069 | |||
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 | ||||
|
r37303 | yield stream.makeframe(requestid=requestid, | ||
typeid=FRAME_TYPE_COMMAND_DATA, | ||||
flags=flags, | ||||
payload=data) | ||||
|
r37069 | |||
if done: | ||||
break | ||||
|
r37070 | |||
|
r37303 | def createbytesresponseframesfrombytes(stream, requestid, data, | ||
|
r37073 | 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: | ||||
|
r37303 | yield stream.makeframe(requestid=requestid, | ||
typeid=FRAME_TYPE_BYTES_RESPONSE, | ||||
flags=FLAG_BYTES_RESPONSE_EOS, | ||||
payload=data) | ||||
|
r37073 | 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 | ||||
|
r37303 | yield stream.makeframe(requestid=requestid, | ||
typeid=FRAME_TYPE_BYTES_RESPONSE, | ||||
flags=flags, | ||||
payload=chunk) | ||||
|
r37073 | |||
if done: | ||||
break | ||||
|
r37303 | def createerrorframe(stream, requestid, msg, protocol=False, application=False): | ||
|
r37073 | # 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 | ||||
|
r37303 | yield stream.makeframe(requestid=requestid, | ||
typeid=FRAME_TYPE_ERROR_RESPONSE, | ||||
flags=flags, | ||||
payload=msg) | ||||
|
r37073 | |||
|
r37303 | def createtextoutputframe(stream, requestid, atoms): | ||
|
r37078 | """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') | ||||
|
r37303 | 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.""" | ||||
|
r37304 | def __init__(self, streamid, active=False): | ||
self.streamid = streamid | ||||
self._active = False | ||||
|
r37303 | 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. | ||||
""" | ||||
|
r37304 | streamflags = 0 | ||
if not self._active: | ||||
streamflags |= STREAM_FLAG_BEGIN_STREAM | ||||
self._active = True | ||||
return makeframe(requestid, self.streamid, streamflags, typeid, flags, | ||||
payload) | ||||
def ensureserverstream(stream): | ||||
if stream.streamid % 2: | ||||
raise error.ProgrammingError('server should only write to even ' | ||||
'numbered streams; %d is not even' % | ||||
stream.streamid) | ||||
|
r37078 | |||
|
r37070 | 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: | ||||
|
r37073 | 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. | ||||
|
r37070 | 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. | ||||
|
r37074 | |||
noop | ||||
Indicates no additional action is required. | ||||
|
r37076 | |||
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. | ||||
|
r37070 | """ | ||
|
r37074 | 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 | ||||
|
r37070 | self._state = 'idle' | ||
|
r37074 | self._bufferedframegens = [] | ||
|
r37304 | # stream id -> stream instance for all active streams from the client. | ||
self._incomingstreams = {} | ||||
|
r37076 | # request id -> dict of commands that are actively being received. | ||
self._receivingcommands = {} | ||||
|
r37081 | # 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() | ||||
|
r37070 | |||
|
r37079 | def onframerecv(self, frame): | ||
|
r37070 | """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. | ||||
""" | ||||
|
r37304 | if not frame.streamid % 2: | ||
self._state = 'errored' | ||||
return self._makeerrorresult( | ||||
_('received frame with even numbered stream ID: %d') % | ||||
frame.streamid) | ||||
if frame.streamid not in self._incomingstreams: | ||||
if not frame.streamflags & STREAM_FLAG_BEGIN_STREAM: | ||||
self._state = 'errored' | ||||
return self._makeerrorresult( | ||||
_('received frame on unknown inactive stream without ' | ||||
'beginning of stream flag set')) | ||||
self._incomingstreams[frame.streamid] = stream(frame.streamid) | ||||
if frame.streamflags & STREAM_FLAG_ENCODING_APPLIED: | ||||
# TODO handle decoding frames | ||||
self._state = 'errored' | ||||
raise error.ProgrammingError('support for decoding stream payloads ' | ||||
'not yet implemented') | ||||
if frame.streamflags & STREAM_FLAG_END_STREAM: | ||||
del self._incomingstreams[frame.streamid] | ||||
|
r37070 | handlers = { | ||
'idle': self._onframeidle, | ||||
|
r37076 | 'command-receiving': self._onframecommandreceiving, | ||
|
r37070 | 'errored': self._onframeerrored, | ||
} | ||||
meth = handlers.get(self._state) | ||||
if not meth: | ||||
raise error.ProgrammingError('unhandled state: %s' % self._state) | ||||
|
r37079 | return meth(frame) | ||
|
r37070 | |||
|
r37303 | def onbytesresponseready(self, stream, requestid, data): | ||
|
r37073 | """Signal that a bytes response is ready to be sent to the client. | ||
The raw bytes response is passed as an argument. | ||||
""" | ||||
|
r37304 | ensureserverstream(stream) | ||
|
r37081 | def sendframes(): | ||
|
r37303 | for frame in createbytesresponseframesfrombytes(stream, requestid, | ||
data): | ||||
|
r37081 | yield frame | ||
self._activecommands.remove(requestid) | ||||
result = sendframes() | ||||
|
r37074 | |||
if self._deferoutput: | ||||
|
r37081 | self._bufferedframegens.append(result) | ||
|
r37074 | return 'noop', {} | ||
else: | ||||
return 'sendframes', { | ||||
|
r37081 | 'framegen': result, | ||
|
r37074 | } | ||
def oninputeof(self): | ||||
"""Signals that end of input has been received. | ||||
No more frames will be received. All pending activity should be | ||||
completed. | ||||
""" | ||||
|
r37076 | # TODO should we do anything about in-flight commands? | ||
|
r37074 | 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 | ||||
|
r37073 | return 'sendframes', { | ||
|
r37074 | 'framegen': makegen(), | ||
|
r37073 | } | ||
|
r37303 | def onapplicationerror(self, stream, requestid, msg): | ||
|
r37304 | ensureserverstream(stream) | ||
|
r37073 | return 'sendframes', { | ||
|
r37303 | 'framegen': createerrorframe(stream, requestid, msg, | ||
application=True), | ||||
|
r37073 | } | ||
|
r37070 | def _makeerrorresult(self, msg): | ||
return 'error', { | ||||
'message': msg, | ||||
} | ||||
|
r37076 | def _makeruncommandresult(self, requestid): | ||
entry = self._receivingcommands[requestid] | ||||
del self._receivingcommands[requestid] | ||||
if self._receivingcommands: | ||||
self._state = 'command-receiving' | ||||
else: | ||||
self._state = 'idle' | ||||
|
r37081 | assert requestid not in self._activecommands | ||
self._activecommands.add(requestid) | ||||
|
r37070 | return 'runcommand', { | ||
|
r37076 | 'requestid': requestid, | ||
'command': entry['command'], | ||||
'args': entry['args'], | ||||
'data': entry['data'].getvalue() if entry['data'] else None, | ||||
|
r37070 | } | ||
def _makewantframeresult(self): | ||||
return 'wantframe', { | ||||
'state': self._state, | ||||
} | ||||
|
r37079 | def _onframeidle(self, frame): | ||
|
r37070 | # The only frame type that should be received in this state is a | ||
# command request. | ||||
|
r37079 | if frame.typeid != FRAME_TYPE_COMMAND_NAME: | ||
|
r37070 | self._state = 'errored' | ||
return self._makeerrorresult( | ||||
|
r37079 | _('expected command frame; got %d') % frame.typeid) | ||
|
r37070 | |||
|
r37079 | if frame.requestid in self._receivingcommands: | ||
|
r37076 | self._state = 'errored' | ||
return self._makeerrorresult( | ||||
|
r37079 | _('request with ID %d already received') % frame.requestid) | ||
|
r37076 | |||
|
r37081 | if frame.requestid in self._activecommands: | ||
self._state = 'errored' | ||||
return self._makeerrorresult(( | ||||
_('request with ID %d is already active') % frame.requestid)) | ||||
|
r37079 | expectingargs = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_ARGS) | ||
expectingdata = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_DATA) | ||||
|
r37076 | |||
|
r37079 | self._receivingcommands[frame.requestid] = { | ||
'command': frame.payload, | ||||
|
r37076 | 'args': {}, | ||
'data': None, | ||||
'expectingargs': expectingargs, | ||||
'expectingdata': expectingdata, | ||||
} | ||||
|
r37070 | |||
|
r37079 | if frame.flags & FLAG_COMMAND_NAME_EOS: | ||
return self._makeruncommandresult(frame.requestid) | ||||
|
r37070 | |||
|
r37076 | if expectingargs or expectingdata: | ||
self._state = 'command-receiving' | ||||
|
r37070 | return self._makewantframeresult() | ||
else: | ||||
self._state = 'errored' | ||||
return self._makeerrorresult(_('missing frame flags on ' | ||||
'command frame')) | ||||
|
r37079 | def _onframecommandreceiving(self, frame): | ||
|
r37076 | # It could be a new command request. Process it as such. | ||
|
r37079 | if frame.typeid == FRAME_TYPE_COMMAND_NAME: | ||
return self._onframeidle(frame) | ||||
|
r37076 | |||
# All other frames should be related to a command that is currently | ||||
|
r37081 | # 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) | ||||
|
r37079 | if frame.requestid not in self._receivingcommands: | ||
|
r37070 | self._state = 'errored' | ||
|
r37076 | return self._makeerrorresult( | ||
_('received frame for request that is not receiving: %d') % | ||||
|
r37079 | frame.requestid) | ||
|
r37076 | |||
|
r37079 | entry = self._receivingcommands[frame.requestid] | ||
|
r37076 | |||
|
r37079 | if frame.typeid == FRAME_TYPE_COMMAND_ARGUMENT: | ||
|
r37076 | if not entry['expectingargs']: | ||
self._state = 'errored' | ||||
return self._makeerrorresult(_( | ||||
'received command argument frame for request that is not ' | ||||
|
r37079 | 'expecting arguments: %d') % frame.requestid) | ||
|
r37076 | |||
|
r37079 | return self._handlecommandargsframe(frame, entry) | ||
|
r37076 | |||
|
r37079 | elif frame.typeid == FRAME_TYPE_COMMAND_DATA: | ||
|
r37076 | if not entry['expectingdata']: | ||
self._state = 'errored' | ||||
return self._makeerrorresult(_( | ||||
'received command data frame for request that is not ' | ||||
|
r37079 | 'expecting data: %d') % frame.requestid) | ||
|
r37076 | |||
if entry['data'] is None: | ||||
entry['data'] = util.bytesio() | ||||
|
r37079 | return self._handlecommanddataframe(frame, entry) | ||
|
r37076 | |||
|
r37079 | def _handlecommandargsframe(self, frame, entry): | ||
|
r37076 | # The frame and state of command should have already been validated. | ||
|
r37079 | assert frame.typeid == FRAME_TYPE_COMMAND_ARGUMENT | ||
|
r37070 | |||
offset = 0 | ||||
|
r37079 | namesize, valuesize = ARGUMENT_FRAME_HEADER.unpack_from(frame.payload) | ||
|
r37070 | offset += ARGUMENT_FRAME_HEADER.size | ||
# The argument name MUST fit inside the frame. | ||||
|
r37079 | argname = bytes(frame.payload[offset:offset + namesize]) | ||
|
r37070 | offset += namesize | ||
if len(argname) != namesize: | ||||
self._state = 'errored' | ||||
return self._makeerrorresult(_('malformed argument frame: ' | ||||
'partial argument name')) | ||||
|
r37079 | argvalue = bytes(frame.payload[offset:]) | ||
|
r37070 | |||
# Argument value spans multiple frames. Record our active state | ||||
# and wait for the next frame. | ||||
|
r37079 | if frame.flags & FLAG_COMMAND_ARGUMENT_CONTINUATION: | ||
|
r37070 | 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')) | ||||
|
r37076 | entry['args'][argname] = argvalue | ||
|
r37070 | |||
|
r37079 | if frame.flags & FLAG_COMMAND_ARGUMENT_EOA: | ||
|
r37076 | if entry['expectingdata']: | ||
|
r37070 | # TODO signal request to run a command once we don't | ||
# buffer data frames. | ||||
return self._makewantframeresult() | ||||
else: | ||||
|
r37079 | return self._makeruncommandresult(frame.requestid) | ||
|
r37070 | else: | ||
return self._makewantframeresult() | ||||
|
r37079 | def _handlecommanddataframe(self, frame, entry): | ||
assert frame.typeid == FRAME_TYPE_COMMAND_DATA | ||||
|
r37070 | |||
# TODO support streaming data instead of buffering it. | ||||
|
r37079 | entry['data'].write(frame.payload) | ||
|
r37070 | |||
|
r37079 | if frame.flags & FLAG_COMMAND_DATA_CONTINUATION: | ||
|
r37070 | return self._makewantframeresult() | ||
|
r37079 | elif frame.flags & FLAG_COMMAND_DATA_EOS: | ||
|
r37076 | entry['data'].seek(0) | ||
|
r37079 | return self._makeruncommandresult(frame.requestid) | ||
|
r37070 | else: | ||
self._state = 'errored' | ||||
return self._makeerrorresult(_('command data frame without ' | ||||
'flags')) | ||||
|
r37079 | def _onframeerrored(self, frame): | ||
|
r37070 | return self._makeerrorresult(_('server already errored')) | ||