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

FLAG_BYTES_RESPONSE_CBOR = 0x04

FLAGS_BYTES_RESPONSE = {

    b'continuation': FLAG_BYTES_RESPONSE_CONTINUATION,

    b'eos': FLAG_BYTES_RESPONSE_EOS,

    b'cbor': FLAG_BYTES_RESPONSE_CBOR,

}

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

def createbytesresponseframesfrombytes(stream, requestid, data, iscbor=False,

                                       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

        if iscbor:

            flags |= FLAG_BYTES_RESPONSE_CBOR

        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

        if iscbor:

            flags |= FLAG_BYTES_RESPONSE_CBOR

        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 = False

    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