diff -r 37d7a1d18b97 -r 40206e227412 mercurial/wireprotoframing.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/mercurial/wireprotoframing.py Mon Mar 19 16:49:53 2018 -0700 @@ -0,0 +1,156 @@ +# wireprotoframing.py - unified framing protocol for wire protocol +# +# Copyright 2018 Gregory Szorc +# +# 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 . import ( + util, +) + +FRAME_HEADER_SIZE = 4 +DEFAULT_MAX_FRAME_SIZE = 32768 + +FRAME_TYPE_COMMAND_NAME = 0x01 +FRAME_TYPE_COMMAND_ARGUMENT = 0x02 +FRAME_TYPE_COMMAND_DATA = 0x03 + +FRAME_TYPES = { + b'command-name': FRAME_TYPE_COMMAND_NAME, + b'command-argument': FRAME_TYPE_COMMAND_ARGUMENT, + b'command-data': FRAME_TYPE_COMMAND_DATA, +} + +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, +} + +# 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, +} + +ARGUMENT_FRAME_HEADER = struct.Struct(r' , creates a frame. + + This can be used by user-facing applications and tests for creating + frames easily without having to type out a bunch of constants. + + Frame type and flags can be specified by integer or named constant. + Flags can be delimited by `|` to bitwise OR them together. + """ + frametype, frameflags, payload = s.split(b' ', 2) + + 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 = util.unescapestr(payload) + + return makeframe(frametype, finalflags, payload) + +def createcommandframes(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(FRAME_TYPE_COMMAND_NAME, flags, 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(FRAME_TYPE_COMMAND_ARGUMENT, flags, 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(FRAME_TYPE_COMMAND_DATA, flags, data) + + if done: + break