mercurial/wireprotoframing.py
changeset 37051 40206e227412
child 37052 8c3c47362934
--- /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 <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
+
+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'<HH')
+
+def makeframe(frametype, frameflags, payload):
+    """Assemble a frame into a byte array."""
+    # TODO assert size of payload.
+    frame = bytearray(FRAME_HEADER_SIZE + len(payload))
+
+    l = struct.pack(r'<I', len(payload))
+    frame[0:3] = l[0:3]
+    frame[3] = (frametype << 4) | frameflags
+    frame[4:] = payload
+
+    return frame
+
+def makeframefromhumanstring(s):
+    """Given a string of the form: <type> <flags> <payload>, 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