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

)

from . import (

    encoding,

    error,

    pycompat,

    util,

    wireprototypes,

)

from .utils import (

    cborutil,

    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 = 0x02

FRAME_TYPE_COMMAND_RESPONSE = 0x03

FRAME_TYPE_ERROR_RESPONSE = 0x05

FRAME_TYPE_TEXT_OUTPUT = 0x06

FRAME_TYPE_PROGRESS = 0x07

FRAME_TYPES = {

    b'command-request': FRAME_TYPE_COMMAND_REQUEST,

    b'command-data': FRAME_TYPE_COMMAND_DATA,

    b'command-response': FRAME_TYPE_COMMAND_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_COMMAND_RESPONSE_CONTINUATION = 0x01

FLAG_COMMAND_RESPONSE_EOS = 0x02

FLAGS_COMMAND_RESPONSE = {

    b'continuation': FLAG_COMMAND_RESPONSE_CONTINUATION,

    b'eos': FLAG_COMMAND_RESPONSE_EOS,

}

# 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_COMMAND_RESPONSE: FLAGS_COMMAND_RESPONSE,

    FRAME_TYPE_ERROR_RESPONSE: {},

    FRAME_TYPE_TEXT_OUTPUT: {},

    FRAME_TYPE_PROGRESS: {},

}

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 = b''.join(cborutil.streamencode(

            stringutil.evalpythonliteral(payload[5:])))

    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,

                        redirect=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.

    """

    data = {b'name': cmd}

    if args:

        data[b'args'] = args

    if redirect:

        data[b'redirect'] = redirect

    data = b''.join(cborutil.streamencode(data))

    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 createcommandresponseframesfrombytes(stream, requestid, data,

                                         maxframesize=DEFAULT_MAX_FRAME_SIZE):

    """Create a raw frame to send a bytes response from static bytes input.

    Returns a generator of bytearrays.

    """

    # Automatically send the overall CBOR response map.

    overall = b''.join(cborutil.streamencode({b'status': b'ok'}))

    if len(overall) > maxframesize:

        raise error.ProgrammingError('not yet implemented')

    # Simple case where we can fit the full response in a single frame.

    if len(overall) + len(data) <= maxframesize:

        flags = FLAG_COMMAND_RESPONSE_EOS

        yield stream.makeframe(requestid=requestid,

                               typeid=FRAME_TYPE_COMMAND_RESPONSE,

                               flags=flags,

                               payload=overall + data)

        return

    # It's easier to send the overall CBOR map in its own frame than to track

    # offsets.

    yield stream.makeframe(requestid=requestid,

                           typeid=FRAME_TYPE_COMMAND_RESPONSE,

                           flags=FLAG_COMMAND_RESPONSE_CONTINUATION,

                           payload=overall)

    offset = 0

    while True:

        chunk = data[offset:offset + maxframesize]

        offset += len(chunk)

        done = offset == len(data)

        if done:

            flags = FLAG_COMMAND_RESPONSE_EOS

        else:

            flags = FLAG_COMMAND_RESPONSE_CONTINUATION

        yield stream.makeframe(requestid=requestid,

                               typeid=FRAME_TYPE_COMMAND_RESPONSE,

                               flags=flags,

                               payload=chunk)

        if done:

            break

def createbytesresponseframesfromgen(stream, requestid, gen,

                                     maxframesize=DEFAULT_MAX_FRAME_SIZE):

    """Generator of frames from a generator of byte chunks.

    This assumes that another frame will follow whatever this emits. i.e.

    this always emits the continuation flag and never emits the end-of-stream

    flag.

    """

    cb = util.chunkbuffer(gen)

    flags = FLAG_COMMAND_RESPONSE_CONTINUATION

    while True:

        chunk = cb.read(maxframesize)

        if not chunk:

            break

        yield stream.makeframe(requestid=requestid,

                               typeid=FRAME_TYPE_COMMAND_RESPONSE,

                               flags=flags,

                               payload=chunk)

        flags |= FLAG_COMMAND_RESPONSE_CONTINUATION

def createcommandresponseokframe(stream, requestid):

    overall = b''.join(cborutil.streamencode({b'status': b'ok'}))

    return stream.makeframe(requestid=requestid,

                            typeid=FRAME_TYPE_COMMAND_RESPONSE,

                            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,

                            payload=overall)

def createcommandresponseeosframe(stream, requestid):

    """Create an empty payload frame representing command end-of-stream."""

    return stream.makeframe(requestid=requestid,

                            typeid=FRAME_TYPE_COMMAND_RESPONSE,

                            flags=FLAG_COMMAND_RESPONSE_EOS,

                            payload=b'')

def createalternatelocationresponseframe(stream, requestid, location):

    data = {

        b'status': b'redirect',

        b'location': {

            b'url': location.url,

            b'mediatype': location.mediatype,

        }

    }

    for a in (r'size', r'fullhashes', r'fullhashseed', r'serverdercerts',

              r'servercadercerts'):

        value = getattr(location, a)

        if value is not None:

            data[b'location'][pycompat.bytestr(a)] = value

    return stream.makeframe(requestid=requestid,

                            typeid=FRAME_TYPE_COMMAND_RESPONSE,

                            flags=FLAG_COMMAND_RESPONSE_CONTINUATION,

                            payload=b''.join(cborutil.streamencode(data)))

def createcommanderrorresponse(stream, requestid, message, args=None):

    # TODO should this be using a list of {'msg': ..., 'args': {}} so atom

    # formatting works consistently?

    m = {

        b'status': b'error',

        b'error': {

            b'message': message,

        }

    }

    if args:

        m[b'error'][b'args'] = args

    overall = b''.join(cborutil.streamencode(m))

    yield stream.makeframe(requestid=requestid,

                           typeid=FRAME_TYPE_COMMAND_RESPONSE,

                           flags=FLAG_COMMAND_RESPONSE_EOS,

                           payload=overall)