# 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 collections

import struct

from .i18n import _

from .thirdparty import (

    attr,

    cbor,

)

from . import (

    encoding,

    error,

    util,

)

from .utils import (

    stringutil,

)

FRAME_HEADER_SIZE = 8

DEFAULT_MAX_FRAME_SIZE = 32768

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,

}

FRAME_TYPE_COMMAND_REQUEST = 0x01

FRAME_TYPE_COMMAND_DATA = 0x03

FRAME_TYPE_BYTES_RESPONSE = 0x04

FRAME_TYPE_ERROR_RESPONSE = 0x05

FRAME_TYPE_TEXT_OUTPUT = 0x06

FRAME_TYPE_PROGRESS = 0x07

FRAME_TYPE_STREAM_SETTINGS = 0x08

FRAME_TYPES = {

    b'command-request': FRAME_TYPE_COMMAND_REQUEST,

    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,

    b'progress': FRAME_TYPE_PROGRESS,

    b'stream-settings': FRAME_TYPE_STREAM_SETTINGS,

}

FLAG_COMMAND_REQUEST_NEW = 0x01

FLAG_COMMAND_REQUEST_CONTINUATION = 0x02

FLAG_COMMAND_REQUEST_MORE_FRAMES = 0x04

FLAG_COMMAND_REQUEST_EXPECT_DATA = 0x08

FLAGS_COMMAND_REQUEST = {

    b'new': FLAG_COMMAND_REQUEST_NEW,

    b'continuation': FLAG_COMMAND_REQUEST_CONTINUATION,

    b'more': FLAG_COMMAND_REQUEST_MORE_FRAMES,

    b'have-data': FLAG_COMMAND_REQUEST_EXPECT_DATA,

}

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_REQUEST: FLAGS_COMMAND_REQUEST,

    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: {},

    FRAME_TYPE_PROGRESS: {},

    FRAME_TYPE_STREAM_SETTINGS: {},

}

ARGUMENT_RECORD_HEADER = struct.Struct(r'<HH')

def humanflags(mapping, value):

    """Convert a numeric flags value to a human value, using a mapping table."""

    namemap = {v: k for k, v in mapping.iteritems()}

    flags = []

    val = 1

    while value >= val:

        if value & val:

            flags.append(namemap.get(val, '<unknown 0x%02x>' % val))

        val <<= 1

    return b'|'.join(flags)

@attr.s(slots=True)

class frameheader(object):

    """Represents the data in a frame header."""

    length = attr.ib()

    requestid = attr.ib()

    streamid = attr.ib()

    streamflags = attr.ib()

    typeid = attr.ib()

    flags = attr.ib()

@attr.s(slots=True, repr=False)

class frame(object):

    """Represents a parsed frame."""

    requestid = attr.ib()

    streamid = attr.ib()

    streamflags = attr.ib()

    typeid = attr.ib()

    flags = attr.ib()

    payload = attr.ib()

    @encoding.strmethod

    def __repr__(self):

        typename = '<unknown 0x%02x>' % self.typeid

        for name, value in FRAME_TYPES.iteritems():

            if value == self.typeid:

                typename = name

                break

        return ('frame(size=%d; request=%d; stream=%d; streamflags=%s; '

                'type=%s; flags=%s)' % (

            len(self.payload), self.requestid, self.streamid,

            humanflags(STREAM_FLAGS, self.streamflags), typename,

            humanflags(FRAME_TYPE_FLAGS.get(self.typeid, {}), self.flags)))

def makeframe(requestid, streamid, streamflags, 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

    # 8 bits stream id

    # 8 bits stream flags

    # 4 bits type

    # 4 bits flags

    l = struct.pack(r'<I', len(payload))

    frame[0:3] = l[0:3]

    struct.pack_into(r'<HBB', frame, 3, requestid, streamid, streamflags)

    frame[7] = (typeid << 4) | flags

    frame[8:] = payload

    return frame

def makeframefromhumanstring(s):

    """Create a frame from a human readable string

    Strings have the form:

        <request-id> <stream-id> <stream-flags> <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 and stream IDs are integers.

    Stream flags, frame type, and flags can be specified by integer or

    named constant.

    Flags can be delimited by `|` to bitwise OR them together.

    If the payload begins with ``cbor:``, the following string will be

    evaluated as Python literal and the resulting object will be fed into

    a CBOR encoder. Otherwise, the payload is interpreted as a Python

    byte string literal.

    """

    fields = s.split(b' ', 5)

    requestid, streamid, streamflags, frametype, frameflags, payload = fields

    requestid = int(requestid)

    streamid = int(streamid)

    finalstreamflags = 0

    for flag in streamflags.split(b'|'):

        if flag in STREAM_FLAGS:

            finalstreamflags |= STREAM_FLAGS[flag]

        else:

            finalstreamflags |= int(flag)

    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)

    if payload.startswith(b'cbor:'):

        payload = cbor.dumps(stringutil.evalpythonliteral(payload[5:]),

                             canonical=True)

    else:

        payload = stringutil.unescapestr(payload)

    return makeframe(requestid=requestid, streamid=streamid,

                     streamflags=finalstreamflags, 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)

    # 16 bits request ID

    # 8 bits stream ID

    # 8 bits stream flags

    # 4 bits frame type

    # 4 bits frame flags

    # ... payload

    framelength = data[0] + 256 * data[1] + 16384 * data[2]

    requestid, streamid, streamflags = struct.unpack_from(r'<HBB', data, 3)

    typeflags = data[7]

    frametype = (typeflags & 0xf0) >> 4

    frameflags = typeflags & 0x0f

    return frameheader(framelength, requestid, streamid, streamflags,

                       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.streamid, h.streamflags, h.typeid, h.flags,

                 payload)

def createcommandframes(stream, requestid, cmd, args, datafh=None,

                        maxframesize=DEFAULT_MAX_FRAME_SIZE):

    """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.

    """

    data = {b'name': cmd}

    if args:

        data[b'args'] = args

    data = cbor.dumps(data, canonical=True)

    offset = 0

    while True:

        flags = 0

        # Must set new or continuation flag.

        if not offset:

            flags |= FLAG_COMMAND_REQUEST_NEW

        else:

            flags |= FLAG_COMMAND_REQUEST_CONTINUATION

        # Data frames is set on all frames.

        if datafh:

            flags |= FLAG_COMMAND_REQUEST_EXPECT_DATA

        payload = data[offset:offset + maxframesize]

        offset += len(payload)

        if len(payload) == maxframesize and offset < len(data):

            flags |= FLAG_COMMAND_REQUEST_MORE_FRAMES

        yield stream.makeframe(requestid=requestid,

                               typeid=FRAME_TYPE_COMMAND_REQUEST,

                               flags=flags,

                               payload=payload)

        if not (flags & FLAG_COMMAND_REQUEST_MORE_FRAMES):

            break

    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 stream.makeframe(requestid=requestid,

                                   typeid=FRAME_TYPE_COMMAND_DATA,

                                   flags=flags,

                                   payload=data)

            if done:

                break

                                       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:

        flags = FLAG_BYTES_RESPONSE_EOS

        yield stream.makeframe(requestid=requestid,

                               typeid=FRAME_TYPE_BYTES_RESPONSE,

                               flags=flags,

                               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 stream.makeframe(requestid=requestid,

                               typeid=FRAME_TYPE_BYTES_RESPONSE,

                               flags=flags,

                               payload=chunk)

        if done:

            break

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 stream.makeframe(requestid=requestid,

                           typeid=FRAME_TYPE_ERROR_RESPONSE,

                           flags=flags,

                           payload=msg)

def createtextoutputframe(stream, requestid, atoms,

                          maxframesize=DEFAULT_MAX_FRAME_SIZE):

    """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()``.

    """

    atomdicts = []

    for (formatting, args, labels) in atoms:

        # 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 ASCII.

        formatting = formatting.decode(r'ascii', r'replace').encode(r'ascii')

        # 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]

        atom = {b'msg': formatting}

        if args:

            atom[b'args'] = args

        if labels:

            atom[b'labels'] = labels

        atomdicts.append(atom)

    payload = cbor.dumps(atomdicts, canonical=True)

    if len(payload) > maxframesize:

        raise ValueError('cannot encode data in a single frame')

    yield stream.makeframe(requestid=requestid,

                           typeid=FRAME_TYPE_TEXT_OUTPUT,

                           flags=0,

                           payload=payload)

class stream(object):

    """Represents a logical unidirectional series of frames."""

    def __init__(self, streamid, active=False):

        self.streamid = streamid

        self._active = active

    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.

        """

        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)

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