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

import abc

import collections

import struct

import typing

from typing import (

    Protocol,

    Type,

)

from .i18n import _

from .thirdparty import attr

# Force pytype to use the non-vendored package

if typing.TYPE_CHECKING:

    # noinspection PyPackageRequirements

    import attr

from . import (

    encoding,

    error,

    pycompat,

    util,

    wireprototypes,

)

from .utils import (

    cborutil,

    stringutil,

)

if typing.TYPE_CHECKING:

    from typing import (

        Iterator,

    )

    HandleSendFramesReturnT = tuple[bytes, dict[bytes, Iterator[bytearray]]]

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: dict[bytes, int] = {

    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_TYPE_SENDER_PROTOCOL_SETTINGS = 0x08

FRAME_TYPE_STREAM_SETTINGS = 0x09

FRAME_TYPES: dict[bytes, int] = {

    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'sender-protocol-settings': FRAME_TYPE_SENDER_PROTOCOL_SETTINGS,

    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: dict[bytes, int] = {

    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: dict[bytes, int] = {

    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: dict[bytes, int] = {

    b'continuation': FLAG_COMMAND_RESPONSE_CONTINUATION,

    b'eos': FLAG_COMMAND_RESPONSE_EOS,

}

FLAG_SENDER_PROTOCOL_SETTINGS_CONTINUATION = 0x01

FLAG_SENDER_PROTOCOL_SETTINGS_EOS = 0x02

FLAGS_SENDER_PROTOCOL_SETTINGS: dict[bytes, int] = {

    b'continuation': FLAG_SENDER_PROTOCOL_SETTINGS_CONTINUATION,

    b'eos': FLAG_SENDER_PROTOCOL_SETTINGS_EOS,

}

FLAG_STREAM_ENCODING_SETTINGS_CONTINUATION = 0x01

FLAG_STREAM_ENCODING_SETTINGS_EOS = 0x02

FLAGS_STREAM_ENCODING_SETTINGS: dict[bytes, int] = {

    b'continuation': FLAG_STREAM_ENCODING_SETTINGS_CONTINUATION,

    b'eos': FLAG_STREAM_ENCODING_SETTINGS_EOS,

}

# Maps frame types to their available flags.

FRAME_TYPE_FLAGS: dict[int, dict[bytes, int]] = {

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

    FRAME_TYPE_SENDER_PROTOCOL_SETTINGS: FLAGS_SENDER_PROTOCOL_SETTINGS,

    FRAME_TYPE_STREAM_SETTINGS: FLAGS_STREAM_ENCODING_SETTINGS,

}

ARGUMENT_RECORD_HEADER = struct.Struct('<HH')

def humanflags(mapping: dict[bytes, int], value: int) -> bytes:

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

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

    flags = []

    val = 1

    while value >= val:

        if value & val:

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

        val <<= 1

    return b'|'.join(flags)

@attr.s(slots=True)

class frameheader:

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

    length = attr.ib(type=int)

    requestid = attr.ib(type=int)

    streamid = attr.ib(type=int)

    streamflags = attr.ib(type=int)

    typeid = attr.ib(type=int)

    flags = attr.ib(type=int)

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

class frame:

    """Represents a parsed frame."""

    requestid = attr.ib(type=int)

    streamid = attr.ib(type=int)

    streamflags = attr.ib(type=int)

    typeid = attr.ib(type=int)

    flags = attr.ib(type=int)

    payload = attr.ib(type=bytes)

    @encoding.strmethod

    def __repr__(self):

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

        for name, value in FRAME_TYPES.items():

            if value == self.typeid:

                typename = name

                break

        return (

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

            b'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: int,

    streamid: int,

    streamflags: int,

    typeid: int,

    flags: int,

    payload: bytes,

) -> bytearray:

    """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('<I', len(payload))

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

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

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

    frame[8:] = payload

    return frame

def makeframefromhumanstring(s: bytes) -> bytearray:

    """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: bytes) -> frameheader:

    """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('<HBB', data, 3)

    typeflags = data[7]

    frametype = (typeflags & 0xF0) >> 4

    frameflags = typeflags & 0x0F

    return frameheader(

        framelength, requestid, streamid, streamflags, frametype, frameflags

    )

def readframe(fh) -> frame | None:

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

            _(b'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(

            _(b'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: stream,

    requestid: int,

    cmd,

    args,

    datafh=None,

    maxframesize: int = DEFAULT_MAX_FRAME_SIZE,

    redirect=None,

) -> Iterator[bytearray]:

    """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 createcommandresponseokframe(

    stream: outputstream, requestid: int

) -> bytearray | None:

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

    if stream.streamsettingssent:

        overall = stream.encode(overall)

        encoded = True

        if not overall:

            return None

    else:

        encoded = False

    return stream.makeframe(

        requestid=requestid,

        typeid=FRAME_TYPE_COMMAND_RESPONSE,

        flags=FLAG_COMMAND_RESPONSE_CONTINUATION,

        payload=overall,

        encoded=encoded,

    )

def createcommandresponseeosframes(

    stream: outputstream,

    requestid: int,

    maxframesize: int = DEFAULT_MAX_FRAME_SIZE,

) -> Iterator[bytearray]:

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

    payload = stream.flush()

    offset = 0

    while True:

        chunk = payload[offset : offset + maxframesize]

        offset += len(chunk)

        done = offset == len(payload)

        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,

            encoded=payload != b'',

        )

        if done:

            break

def createalternatelocationresponseframe(

    stream: outputstream, requestid: int, location

) -> bytearray: