This is an archive of the discontinued Mercurial Phabricator instance.

Differential D2899

wireproto: define attr-based classes for representing frames
ClosedPublic

Authored by indygreg on Mar 19 2018, 5:41 PM.

Details

Reviewers
durin42
Group Reviewers
hg-reviewers
Commits
rHG884a0c1604ad: wireproto: define attr-based classes for representing frames
Summary

When frames only had 3 attributes, it was reasonable to
represent them as a tuple. With them growing more attributes,
it will be easier to pass them around as a more formal type.
So let's define attr-based classes to represent frame headers and
full frames.

Diff Detail

Repository
rHG Mercurial
Lint
Automatic diff as part of commit; lint not applicable.
Unit
Automatic diff as part of commit; unit tests not applicable.

Event Timeline

indygreg created this revision.Mar 19 2018, 5:41 PM
durin42 accepted this revision.Mar 21 2018, 9:33 PM
This revision is now accepted and ready to land.Mar 21 2018, 9:33 PM
Closed by commit rHG884a0c1604ad: wireproto: define attr-based classes for representing frames (authored by indygreg). · Explain WhyMar 21 2018, 9:41 PM
This revision was automatically updated to reflect the committed changes.

Revision Contents
Changeset List

Diff 7237

mercurial/wireprotoframing.py

# wireprotoframing.py - unified framing protocol for wire protocol # wireprotoframing.py - unified framing protocol for wire protocol
# #
# Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com> # Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com>
# #
# This software may be used and distributed according to the terms of the # This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version. # GNU General Public License version 2 or any later version.
# This file contains functionality to support the unified frame-based wire # This file contains functionality to support the unified frame-based wire
# protocol. For details about the protocol, see # protocol. For details about the protocol, see
# `hg help internals.wireprotocol`. # `hg help internals.wireprotocol`.
from __future__ import absolute_import from __future__ import absolute_import
import struct import struct
from .i18n import _ from .i18n import _
from .thirdparty import (
attr,
)
from . import ( from . import (
error, error,
util, util,
) )
FRAME_HEADER_SIZE = 6 FRAME_HEADER_SIZE = 6
DEFAULT_MAX_FRAME_SIZE = 32768 DEFAULT_MAX_FRAME_SIZE = 32768
FRAME_TYPE_COMMAND_DATA: FLAGS_COMMAND_DATA, FRAME_TYPE_COMMAND_DATA: FLAGS_COMMAND_DATA,
FRAME_TYPE_BYTES_RESPONSE: FLAGS_BYTES_RESPONSE, FRAME_TYPE_BYTES_RESPONSE: FLAGS_BYTES_RESPONSE,
FRAME_TYPE_ERROR_RESPONSE: FLAGS_ERROR_RESPONSE, FRAME_TYPE_ERROR_RESPONSE: FLAGS_ERROR_RESPONSE,
FRAME_TYPE_TEXT_OUTPUT: {}, FRAME_TYPE_TEXT_OUTPUT: {},
} }
ARGUMENT_FRAME_HEADER = struct.Struct(r'<HH') 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, frametype, frameflags, payload): def makeframe(requestid, frametype, frameflags, payload):
"""Assemble a frame into a byte array.""" """Assemble a frame into a byte array."""
# TODO assert size of payload. # TODO assert size of payload.
frame = bytearray(FRAME_HEADER_SIZE + len(payload)) frame = bytearray(FRAME_HEADER_SIZE + len(payload))
# 24 bits length # 24 bits length
# 16 bits request id # 16 bits request id
# 4 bits type # 4 bits type
# ... payload # ... payload
framelength = data[0] + 256 * data[1] + 16384 * data[2] framelength = data[0] + 256 * data[1] + 16384 * data[2]
requestid = struct.unpack_from(r'<H', data, 3)[0] requestid = struct.unpack_from(r'<H', data, 3)[0]
typeflags = data[5] typeflags = data[5]
frametype = (typeflags & 0xf0) >> 4 frametype = (typeflags & 0xf0) >> 4
frameflags = typeflags & 0x0f frameflags = typeflags & 0x0f
return requestid, frametype, frameflags, framelength return frameheader(framelength, requestid, frametype, frameflags)
def readframe(fh): def readframe(fh):
"""Read a unified framing protocol frame from a file object. """Read a unified framing protocol frame from a file object.
Returns a 3-tuple of (type, flags, payload) for the decoded frame or 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 None if no frame is available. May raise if a malformed frame is
seen. seen.
""" """
header = bytearray(FRAME_HEADER_SIZE) header = bytearray(FRAME_HEADER_SIZE)
readcount = fh.readinto(header) readcount = fh.readinto(header)
if readcount == 0: if readcount == 0:
return None return None
if readcount != FRAME_HEADER_SIZE: if readcount != FRAME_HEADER_SIZE:
raise error.Abort(_('received incomplete frame: got %d bytes: %s') % raise error.Abort(_('received incomplete frame: got %d bytes: %s') %
(readcount, header)) (readcount, header))
requestid, frametype, frameflags, framelength = parseheader(header) h = parseheader(header)
payload = fh.read(framelength) payload = fh.read(h.length)
if len(payload) != framelength: if len(payload) != h.length:
raise error.Abort(_('frame length error: expected %d; got %d') % raise error.Abort(_('frame length error: expected %d; got %d') %
(framelength, len(payload))) (h.length, len(payload)))
return requestid, frametype, frameflags, payload return frame(h.requestid, h.typeid, h.flags, payload)
def createcommandframes(requestid, cmd, args, datafh=None): def createcommandframes(requestid, cmd, args, datafh=None):
"""Create frames necessary to transmit a request to run a command. """Create frames necessary to transmit a request to run a command.
This is a generator of bytearrays. Each item represents a frame This is a generator of bytearrays. Each item represents a frame
ready to be sent over the wire to a peer. ready to be sent over the wire to a peer.
""" """
flags = 0 flags = 0
sender cannot receive until all data has been transmitted. sender cannot receive until all data has been transmitted.
""" """
self._deferoutput = deferoutput self._deferoutput = deferoutput
self._state = 'idle' self._state = 'idle'
self._bufferedframegens = [] self._bufferedframegens = []
# request id -> dict of commands that are actively being received. # request id -> dict of commands that are actively being received.
self._receivingcommands = {} self._receivingcommands = {}
def onframerecv(self, requestid, frametype, frameflags, payload): def onframerecv(self, frame):
"""Process a frame that has been received off the wire. """Process a frame that has been received off the wire.
Returns a dict with an ``action`` key that details what action, Returns a dict with an ``action`` key that details what action,
if any, the consumer should take next. if any, the consumer should take next.
""" """
handlers = { handlers = {
'idle': self._onframeidle, 'idle': self._onframeidle,
'command-receiving': self._onframecommandreceiving, 'command-receiving': self._onframecommandreceiving,
'errored': self._onframeerrored, 'errored': self._onframeerrored,
} }
meth = handlers.get(self._state) meth = handlers.get(self._state)
if not meth: if not meth:
raise error.ProgrammingError('unhandled state: %s' % self._state) raise error.ProgrammingError('unhandled state: %s' % self._state)
return meth(requestid, frametype, frameflags, payload) return meth(frame)
def onbytesresponseready(self, requestid, data): def onbytesresponseready(self, requestid, data):
"""Signal that a bytes response is ready to be sent to the client. """Signal that a bytes response is ready to be sent to the client.
The raw bytes response is passed as an argument. The raw bytes response is passed as an argument.
""" """
framegen = createbytesresponseframesfrombytes(requestid, data) framegen = createbytesresponseframesfrombytes(requestid, data)
'data': entry['data'].getvalue() if entry['data'] else None, 'data': entry['data'].getvalue() if entry['data'] else None,
} }
def _makewantframeresult(self): def _makewantframeresult(self):
return 'wantframe', { return 'wantframe', {
'state': self._state, 'state': self._state,
} }
def _onframeidle(self, requestid, frametype, frameflags, payload): def _onframeidle(self, frame):
# The only frame type that should be received in this state is a # The only frame type that should be received in this state is a
# command request. # command request.
if frametype != FRAME_TYPE_COMMAND_NAME: if frame.typeid != FRAME_TYPE_COMMAND_NAME:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult( return self._makeerrorresult(
_('expected command frame; got %d') % frametype) _('expected command frame; got %d') % frame.typeid)
if requestid in self._receivingcommands: if frame.requestid in self._receivingcommands:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult( return self._makeerrorresult(
_('request with ID %d already received') % requestid) _('request with ID %d already received') % frame.requestid)
expectingargs = bool(frameflags & FLAG_COMMAND_NAME_HAVE_ARGS) expectingargs = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_ARGS)
expectingdata = bool(frameflags & FLAG_COMMAND_NAME_HAVE_DATA) expectingdata = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_DATA)
self._receivingcommands[requestid] = { self._receivingcommands[frame.requestid] = {
'command': payload, 'command': frame.payload,
'args': {}, 'args': {},
'data': None, 'data': None,
'expectingargs': expectingargs, 'expectingargs': expectingargs,
'expectingdata': expectingdata, 'expectingdata': expectingdata,
} }
if frameflags & FLAG_COMMAND_NAME_EOS: if frame.flags & FLAG_COMMAND_NAME_EOS:
return self._makeruncommandresult(requestid) return self._makeruncommandresult(frame.requestid)
if expectingargs or expectingdata: if expectingargs or expectingdata:
self._state = 'command-receiving' self._state = 'command-receiving'
return self._makewantframeresult() return self._makewantframeresult()
else: else:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult(_('missing frame flags on ' return self._makeerrorresult(_('missing frame flags on '
'command frame')) 'command frame'))
def _onframecommandreceiving(self, requestid, frametype, frameflags, def _onframecommandreceiving(self, frame):
payload):
# It could be a new command request. Process it as such. # It could be a new command request. Process it as such.
if frametype == FRAME_TYPE_COMMAND_NAME: if frame.typeid == FRAME_TYPE_COMMAND_NAME:
return self._onframeidle(requestid, frametype, frameflags, payload) return self._onframeidle(frame)
# All other frames should be related to a command that is currently # All other frames should be related to a command that is currently
# receiving. # receiving.
if requestid not in self._receivingcommands: if frame.requestid not in self._receivingcommands:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult( return self._makeerrorresult(
_('received frame for request that is not receiving: %d') % _('received frame for request that is not receiving: %d') %
requestid) frame.requestid)
entry = self._receivingcommands[requestid] entry = self._receivingcommands[frame.requestid]
if frametype == FRAME_TYPE_COMMAND_ARGUMENT: if frame.typeid == FRAME_TYPE_COMMAND_ARGUMENT:
if not entry['expectingargs']: if not entry['expectingargs']:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult(_( return self._makeerrorresult(_(
'received command argument frame for request that is not ' 'received command argument frame for request that is not '
'expecting arguments: %d') % requestid) 'expecting arguments: %d') % frame.requestid)
return self._handlecommandargsframe(requestid, entry, frametype, return self._handlecommandargsframe(frame, entry)
frameflags, payload)
elif frametype == FRAME_TYPE_COMMAND_DATA: elif frame.typeid == FRAME_TYPE_COMMAND_DATA:
if not entry['expectingdata']: if not entry['expectingdata']:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult(_( return self._makeerrorresult(_(
'received command data frame for request that is not ' 'received command data frame for request that is not '
'expecting data: %d') % requestid) 'expecting data: %d') % frame.requestid)
if entry['data'] is None: if entry['data'] is None:
entry['data'] = util.bytesio() entry['data'] = util.bytesio()
return self._handlecommanddataframe(requestid, entry, frametype, return self._handlecommanddataframe(frame, entry)
frameflags, payload)
def _handlecommandargsframe(self, requestid, entry, frametype, frameflags, def _handlecommandargsframe(self, frame, entry):
payload):
# The frame and state of command should have already been validated. # The frame and state of command should have already been validated.
assert frametype == FRAME_TYPE_COMMAND_ARGUMENT assert frame.typeid == FRAME_TYPE_COMMAND_ARGUMENT
offset = 0 offset = 0
namesize, valuesize = ARGUMENT_FRAME_HEADER.unpack_from(payload) namesize, valuesize = ARGUMENT_FRAME_HEADER.unpack_from(frame.payload)
offset += ARGUMENT_FRAME_HEADER.size offset += ARGUMENT_FRAME_HEADER.size
# The argument name MUST fit inside the frame. # The argument name MUST fit inside the frame.
argname = bytes(payload[offset:offset + namesize]) argname = bytes(frame.payload[offset:offset + namesize])
offset += namesize offset += namesize
if len(argname) != namesize: if len(argname) != namesize:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult(_('malformed argument frame: ' return self._makeerrorresult(_('malformed argument frame: '
'partial argument name')) 'partial argument name'))
argvalue = bytes(payload[offset:]) argvalue = bytes(frame.payload[offset:])
# Argument value spans multiple frames. Record our active state # Argument value spans multiple frames. Record our active state
# and wait for the next frame. # and wait for the next frame.
if frameflags & FLAG_COMMAND_ARGUMENT_CONTINUATION: if frame.flags & FLAG_COMMAND_ARGUMENT_CONTINUATION:
raise error.ProgrammingError('not yet implemented') raise error.ProgrammingError('not yet implemented')
# Common case: the argument value is completely contained in this # Common case: the argument value is completely contained in this
# frame. # frame.
if len(argvalue) != valuesize: if len(argvalue) != valuesize:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult(_('malformed argument frame: ' return self._makeerrorresult(_('malformed argument frame: '
'partial argument value')) 'partial argument value'))
entry['args'][argname] = argvalue entry['args'][argname] = argvalue
if frameflags & FLAG_COMMAND_ARGUMENT_EOA: if frame.flags & FLAG_COMMAND_ARGUMENT_EOA:
if entry['expectingdata']: if entry['expectingdata']:
# TODO signal request to run a command once we don't # TODO signal request to run a command once we don't
# buffer data frames. # buffer data frames.
return self._makewantframeresult() return self._makewantframeresult()
else: else:
return self._makeruncommandresult(requestid) return self._makeruncommandresult(frame.requestid)
else: else:
return self._makewantframeresult() return self._makewantframeresult()
def _handlecommanddataframe(self, requestid, entry, frametype, frameflags, def _handlecommanddataframe(self, frame, entry):
payload): assert frame.typeid == FRAME_TYPE_COMMAND_DATA
assert frametype == FRAME_TYPE_COMMAND_DATA
# TODO support streaming data instead of buffering it. # TODO support streaming data instead of buffering it.
entry['data'].write(payload) entry['data'].write(frame.payload)
if frameflags & FLAG_COMMAND_DATA_CONTINUATION: if frame.flags & FLAG_COMMAND_DATA_CONTINUATION:
return self._makewantframeresult() return self._makewantframeresult()
elif frameflags & FLAG_COMMAND_DATA_EOS: elif frame.flags & FLAG_COMMAND_DATA_EOS:
entry['data'].seek(0) entry['data'].seek(0)
return self._makeruncommandresult(requestid) return self._makeruncommandresult(frame.requestid)
else: else:
self._state = 'errored' self._state = 'errored'
return self._makeerrorresult(_('command data frame without ' return self._makeerrorresult(_('command data frame without '
'flags')) 'flags'))
def _onframeerrored(self, requestid, frametype, frameflags, payload): def _onframeerrored(self, frame):
return self._makeerrorresult(_('server already errored')) return self._makeerrorresult(_('server already errored'))

mercurial/wireprotoserver.py

while True: while True:
frame = wireprotoframing.readframe(req.bodyfh) frame = wireprotoframing.readframe(req.bodyfh)
if not frame: if not frame:
states.append(b'received: <no frame>') states.append(b'received: <no frame>')
break break
requestid, frametype, frameflags, payload = frame states.append(b'received: %d %d %d %s' % (frame.typeid, frame.flags,
states.append(b'received: %d %d %d %s' % (frametype, frameflags, frame.requestid,
requestid, payload)) frame.payload))
action, meta = reactor.onframerecv(requestid, frametype, frameflags, action, meta = reactor.onframerecv(frame)
payload)
states.append(json.dumps((action, meta), sort_keys=True, states.append(json.dumps((action, meta), sort_keys=True,
separators=(', ', ': '))) separators=(', ', ': ')))
action, meta = reactor.oninputeof() action, meta = reactor.oninputeof()
meta['action'] = action meta['action'] = action
states.append(json.dumps(meta, sort_keys=True, separators=(', ',': '))) states.append(json.dumps(meta, sort_keys=True, separators=(', ',': ')))
res.status = b'200 OK' res.status = b'200 OK'
reactor = wireprotoframing.serverreactor(deferoutput=True) reactor = wireprotoframing.serverreactor(deferoutput=True)
seencommand = False seencommand = False
while True: while True:
frame = wireprotoframing.readframe(req.bodyfh) frame = wireprotoframing.readframe(req.bodyfh)
if not frame: if not frame:
break break
action, meta = reactor.onframerecv(*frame) action, meta = reactor.onframerecv(frame)
if action == 'wantframe': if action == 'wantframe':
# Need more data before we can do anything. # Need more data before we can do anything.
continue continue
elif action == 'runcommand': elif action == 'runcommand':
sentoutput = _httpv2runcommand(ui, repo, req, res, authedperm, sentoutput = _httpv2runcommand(ui, repo, req, res, authedperm,
reqcommand, reactor, meta, reqcommand, reactor, meta,
issubsequent=seencommand) issubsequent=seencommand)

tests/test-wireproto-serverreactor.py

return framing.serverreactor(deferoutput=deferoutput) return framing.serverreactor(deferoutput=deferoutput)
def sendframes(reactor, gen): def sendframes(reactor, gen):
"""Send a generator of frame bytearray to a reactor. """Send a generator of frame bytearray to a reactor.
Emits a generator of results from ``onframerecv()`` calls. Emits a generator of results from ``onframerecv()`` calls.
""" """
for frame in gen: for frame in gen:
rid, frametype, frameflags, framelength = framing.parseheader(frame) header = framing.parseheader(frame)
payload = frame[framing.FRAME_HEADER_SIZE:] payload = frame[framing.FRAME_HEADER_SIZE:]
assert len(payload) == framelength assert len(payload) == header.length
yield reactor.onframerecv(rid, frametype, frameflags, payload) yield reactor.onframerecv(framing.frame(header.requestid,
header.typeid,
header.flags,
payload))
def sendcommandframes(reactor, rid, cmd, args, datafh=None): def sendcommandframes(reactor, rid, cmd, args, datafh=None):
"""Generate frames to run a command and send them to a reactor.""" """Generate frames to run a command and send them to a reactor."""
return sendframes(reactor, return sendframes(reactor,
framing.createcommandframes(rid, cmd, args, datafh)) framing.createcommandframes(rid, cmd, args, datafh))
class FrameTests(unittest.TestCase): class FrameTests(unittest.TestCase):
def testdataexactframesize(self): def testdataexactframesize(self):