tigrcorn-protocols 0.3.16.dev5__py3-none-any.whl

tigrcorn_protocols/websocket/extensions.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import zlib
+
+from tigrcorn_core.errors import ProtocolError
+from tigrcorn_core.utils.headers import get_headers
+
+_PERMESSAGE_DEFLATE = b"permessage-deflate"
+_SERVER_NO_CONTEXT_TAKEOVER = b"server_no_context_takeover"
+_CLIENT_NO_CONTEXT_TAKEOVER = b"client_no_context_takeover"
+_SERVER_MAX_WINDOW_BITS = b"server_max_window_bits"
+_CLIENT_MAX_WINDOW_BITS = b"client_max_window_bits"
+
+_VALID_PMD_PARAMETERS = {
+    _SERVER_NO_CONTEXT_TAKEOVER,
+    _CLIENT_NO_CONTEXT_TAKEOVER,
+    _SERVER_MAX_WINDOW_BITS,
+    _CLIENT_MAX_WINDOW_BITS,
+}
+
+
+def _split_quoted(value: bytes, delimiter: int) -> list[bytes]:
+    parts: list[bytes] = []
+    buf = bytearray()
+    in_quote = False
+    escape = False
+    for byte in value:
+        if escape:
+            buf.append(byte)
+            escape = False
+            continue
+        if in_quote:
+            if byte == 0x5C:  # backslash
+                escape = True
+                continue
+            if byte == 0x22:
+                in_quote = False
+            buf.append(byte)
+            continue
+        if byte == 0x22:
+            in_quote = True
+            buf.append(byte)
+            continue
+        if byte == delimiter:
+            part = bytes(buf).strip()
+            if part:
+                parts.append(part)
+            buf.clear()
+            continue
+        buf.append(byte)
+    if in_quote:
+        raise ProtocolError('malformed websocket extension header')
+    part = bytes(buf).strip()
+    if part:
+        parts.append(part)
+    return parts
+
+
+def _parse_token_value(value: bytes) -> bytes:
+    raw = value.strip()
+    if len(raw) >= 2 and raw[:1] == raw[-1:] == b'"':
+        inner = raw[1:-1]
+        if b'"' in inner:
+            raise ProtocolError('malformed websocket extension parameter value')
+        return inner
+    return raw
+
+
+def _parse_window_bits(value: bytes) -> int:
+    if not value or (len(value) > 1 and value.startswith(b'0')):
+        raise ProtocolError('invalid permessage-deflate window bits parameter')
+    try:
+        bits = int(value.decode('ascii', 'strict'))
+    except UnicodeDecodeError as exc:
+        raise ProtocolError('invalid permessage-deflate window bits parameter') from exc
+    except ValueError as exc:
+        raise ProtocolError('invalid permessage-deflate window bits parameter') from exc
+    if not 8 <= bits <= 15:
+        raise ProtocolError('invalid permessage-deflate window bits parameter')
+    return bits
+
+
+@dataclass(slots=True, frozen=True)
+class PerMessageDeflateOffer:
+    server_no_context_takeover: bool = False
+    client_no_context_takeover: bool = False
+    server_max_window_bits: int | None = None
+    client_max_window_bits_requested: bool = False
+    client_max_window_bits: int | None = None
+
+
+@dataclass(slots=True, frozen=True)
+class PerMessageDeflateAgreement:
+    server_no_context_takeover: bool = False
+    client_no_context_takeover: bool = False
+    server_max_window_bits: int | None = None
+    client_max_window_bits: int | None = None
+
+    def as_header_value(self) -> bytes:
+        params = [b'permessage-deflate']
+        if self.server_no_context_takeover:
+            params.append(_SERVER_NO_CONTEXT_TAKEOVER)
+        if self.client_no_context_takeover:
+            params.append(_CLIENT_NO_CONTEXT_TAKEOVER)
+        if self.server_max_window_bits is not None:
+            params.append(_SERVER_MAX_WINDOW_BITS + b'=' + str(self.server_max_window_bits).encode('ascii'))
+        if self.client_max_window_bits is not None:
+            params.append(_CLIENT_MAX_WINDOW_BITS + b'=' + str(self.client_max_window_bits).encode('ascii'))
+        return b'; '.join(params)
+
+
+class PerMessageDeflateRuntime:
+    def __init__(self, agreement: PerMessageDeflateAgreement) -> None:
+        self.agreement = agreement
+        self._compressor: zlib.compressobj | None = None
+        self._decompressor: zlib.decompressobj | None = None
+
+    @staticmethod
+    def _runtime_window_bits(bits: int | None) -> int:
+        if bits is None:
+            return 15
+        # Python's raw DEFLATE bindings reject 8 here; fall back to the smallest
+        # supported raw window while preserving interoperable decompression.
+        return max(bits, 9)
+
+    def _new_compressor(self) -> zlib.compressobj:
+        return zlib.compressobj(wbits=-self._runtime_window_bits(self.agreement.server_max_window_bits))
+
+    def _new_decompressor(self) -> zlib.decompressobj:
+        return zlib.decompressobj(wbits=-self._runtime_window_bits(self.agreement.client_max_window_bits))
+
+    def compress_message(self, payload: bytes) -> bytes:
+        if self._compressor is None or self.agreement.server_no_context_takeover:
+            self._compressor = self._new_compressor()
+        compressed = self._compressor.compress(payload) + self._compressor.flush(zlib.Z_SYNC_FLUSH)
+        if not compressed.endswith(b'\x00\x00\xff\xff'):
+            raise RuntimeError('unexpected permessage-deflate trailer')
+        if self.agreement.server_no_context_takeover:
+            self._compressor = None
+        return compressed[:-4]
+
+    def decompress_message(self, payload: bytes) -> bytes:
+        if self._decompressor is None or self.agreement.client_no_context_takeover:
+            self._decompressor = self._new_decompressor()
+        try:
+            data = self._decompressor.decompress(payload + b'\x00\x00\xff\xff')
+        except zlib.error as exc:
+            raise ProtocolError('invalid permessage-deflate payload') from exc
+        if self._decompressor.unconsumed_tail or self._decompressor.unused_data:
+            raise ProtocolError('invalid permessage-deflate payload')
+        if self.agreement.client_no_context_takeover:
+            self._decompressor = None
+        return data
+
+
+def _iter_extension_elements(values: list[bytes]) -> list[tuple[bytes, list[tuple[bytes, bytes | None]]]]:
+    joined = b', '.join(value.strip() for value in values if value.strip())
+    if not joined:
+        return []
+    elements: list[tuple[bytes, list[tuple[bytes, bytes | None]]]] = []
+    for item in _split_quoted(joined, 0x2C):  # comma
+        parts = _split_quoted(item, 0x3B)  # semicolon
+        if not parts:
+            continue
+        name = parts[0].strip().lower()
+        params: list[tuple[bytes, bytes | None]] = []
+        for raw_param in parts[1:]:
+            if not raw_param:
+                continue
+            if b'=' in raw_param:
+                raw_name, raw_value = raw_param.split(b'=', 1)
+                param_name = raw_name.strip().lower()
+                param_value = _parse_token_value(raw_value)
+            else:
+                param_name = raw_param.strip().lower()
+                param_value = None
+            if not param_name:
+                raise ProtocolError('malformed websocket extension parameter')
+            params.append((param_name, param_value))
+        elements.append((name, params))
+    return elements
+
+
+def parse_permessage_deflate_offers(headers: list[tuple[bytes, bytes]]) -> list[PerMessageDeflateOffer]:
+    offers: list[PerMessageDeflateOffer] = []
+    for name, params in _iter_extension_elements(get_headers(headers, b'sec-websocket-extensions')):
+        if name != _PERMESSAGE_DEFLATE:
+            continue
+        try:
+            offers.append(_parse_offer_parameters(params))
+        except ProtocolError:
+            continue
+    return offers
+
+
+
+
+def default_permessage_deflate_agreement(offers: list[PerMessageDeflateOffer]) -> PerMessageDeflateAgreement | None:
+    """Choose a default server agreement for a valid permessage-deflate offer set.
+
+    The server accepts the first valid offer and mirrors explicit window constraints so
+    the generated response header corresponds to the client offer across websocket
+    carriers, including HTTP/2 and HTTP/3 third-party clients that require explicit
+    parameter echoing.
+    """
+    if not offers:
+        return None
+    offer = offers[0]
+    return PerMessageDeflateAgreement(
+        server_no_context_takeover=offer.server_no_context_takeover,
+        client_no_context_takeover=False,
+        server_max_window_bits=offer.server_max_window_bits,
+        client_max_window_bits=offer.client_max_window_bits if offer.client_max_window_bits_requested else None,
+    )
+def negotiate_permessage_deflate(
+    *,
+    request_headers: list[tuple[bytes, bytes]],
+    response_headers: list[tuple[bytes, bytes]],
+) -> PerMessageDeflateAgreement | None:
+    response_values = get_headers(response_headers, b'sec-websocket-extensions')
+    if not response_values:
+        return None
+    response_elements = _iter_extension_elements(response_values)
+    if len(response_elements) != 1 or response_elements[0][0] != _PERMESSAGE_DEFLATE:
+        raise RuntimeError('unsupported websocket extension negotiation')
+    agreement = _parse_response_parameters(response_elements[0][1])
+    offers = parse_permessage_deflate_offers(request_headers)
+    if not offers:
+        raise RuntimeError('websocket extension not offered by the client')
+    for offer in offers:
+        if _agreement_matches_offer(agreement, offer):
+            return agreement
+    raise RuntimeError('websocket extension negotiation does not correspond to a client offer')
+
+
+def _parse_offer_parameters(params: list[tuple[bytes, bytes | None]]) -> PerMessageDeflateOffer:
+    server_no_context_takeover = False
+    client_no_context_takeover = False
+    server_max_window_bits: int | None = None
+    client_max_window_bits_requested = False
+    client_max_window_bits: int | None = None
+    seen: set[bytes] = set()
+    for name, value in params:
+        if name not in _VALID_PMD_PARAMETERS or name in seen:
+            raise ProtocolError('invalid permessage-deflate offer')
+        seen.add(name)
+        if name == _SERVER_NO_CONTEXT_TAKEOVER:
+            if value is not None:
+                raise ProtocolError('invalid permessage-deflate offer')
+            server_no_context_takeover = True
+            continue
+        if name == _CLIENT_NO_CONTEXT_TAKEOVER:
+            if value is not None:
+                raise ProtocolError('invalid permessage-deflate offer')
+            client_no_context_takeover = True
+            continue
+        if name == _SERVER_MAX_WINDOW_BITS:
+            if value is None:
+                raise ProtocolError('invalid permessage-deflate offer')
+            server_max_window_bits = _parse_window_bits(value)
+            continue
+        if name == _CLIENT_MAX_WINDOW_BITS:
+            client_max_window_bits_requested = True
+            if value is not None:
+                client_max_window_bits = _parse_window_bits(value)
+            continue
+    return PerMessageDeflateOffer(
+        server_no_context_takeover=server_no_context_takeover,
+        client_no_context_takeover=client_no_context_takeover,
+        server_max_window_bits=server_max_window_bits,
+        client_max_window_bits_requested=client_max_window_bits_requested,
+        client_max_window_bits=client_max_window_bits,
+    )
+
+
+def _parse_response_parameters(params: list[tuple[bytes, bytes | None]]) -> PerMessageDeflateAgreement:
+    server_no_context_takeover = False
+    client_no_context_takeover = False
+    server_max_window_bits: int | None = None
+    client_max_window_bits: int | None = None
+    seen: set[bytes] = set()
+    for name, value in params:
+        if name not in _VALID_PMD_PARAMETERS or name in seen:
+            raise RuntimeError('unsupported websocket extension negotiation')
+        seen.add(name)
+        if name == _SERVER_NO_CONTEXT_TAKEOVER:
+            if value is not None:
+                raise RuntimeError('unsupported websocket extension negotiation')
+            server_no_context_takeover = True
+            continue
+        if name == _CLIENT_NO_CONTEXT_TAKEOVER:
+            if value is not None:
+                raise RuntimeError('unsupported websocket extension negotiation')
+            client_no_context_takeover = True
+            continue
+        if name == _SERVER_MAX_WINDOW_BITS:
+            if value is None:
+                raise RuntimeError('unsupported websocket extension negotiation')
+            server_max_window_bits = _parse_window_bits(value)
+            continue
+        if value is None:
+            raise RuntimeError('unsupported websocket extension negotiation')
+        client_max_window_bits = _parse_window_bits(value)
+    return PerMessageDeflateAgreement(
+        server_no_context_takeover=server_no_context_takeover,
+        client_no_context_takeover=client_no_context_takeover,
+        server_max_window_bits=server_max_window_bits,
+        client_max_window_bits=client_max_window_bits,
+    )
+
+
+def _agreement_matches_offer(agreement: PerMessageDeflateAgreement, offer: PerMessageDeflateOffer) -> bool:
+    if offer.server_no_context_takeover and not agreement.server_no_context_takeover:
+        return False
+    if offer.server_max_window_bits is not None:
+        if agreement.server_max_window_bits is None or agreement.server_max_window_bits > offer.server_max_window_bits:
+            return False
+    if agreement.client_max_window_bits is not None:
+        if not offer.client_max_window_bits_requested:
+            return False
+        if offer.client_max_window_bits is not None and agreement.client_max_window_bits > offer.client_max_window_bits:
+            return False
+    return True

tigrcorn_protocols/websocket/frames.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass
+
+from tigrcorn_core.errors import ProtocolError
+from tigrcorn_core.types import StreamReaderLike
+
+OP_CONT = 0x0
+OP_TEXT = 0x1
+OP_BINARY = 0x2
+OP_CLOSE = 0x8
+OP_PING = 0x9
+OP_PONG = 0xA
+_CONTROL_OPCODES = {OP_CLOSE, OP_PING, OP_PONG}
+_DATA_OPCODES = {OP_CONT, OP_TEXT, OP_BINARY}
+_VALID_OPCODES = _CONTROL_OPCODES | _DATA_OPCODES
+_FORBIDDEN_CLOSE_CODES = {1004, 1005, 1006, 1015}
+
+
+@dataclass(slots=True)
+class Frame:
+    fin: bool
+    opcode: int
+    payload: bytes
+    rsv1: bool = False
+
+
+def _mask_payload(mask_key: bytes, payload: bytes) -> bytes:
+    return bytes(b ^ mask_key[i % 4] for i, b in enumerate(payload))
+
+
+def validate_close_code(code: int) -> None:
+    if code < 1000 or code >= 5000:
+        raise ProtocolError('invalid close code')
+    if code in _FORBIDDEN_CLOSE_CODES:
+        raise ProtocolError('invalid close code')
+    if 1016 <= code <= 2999:
+        raise ProtocolError('invalid close code')
+
+
+def _validate_frame_semantics(fin: bool, opcode: int, payload_length: int) -> None:
+    if opcode not in _VALID_OPCODES:
+        raise ProtocolError('unsupported websocket opcode')
+    if opcode in _CONTROL_OPCODES:
+        if not fin:
+            raise ProtocolError('control frames must not be fragmented')
+        if payload_length > 125:
+            raise ProtocolError('control frame payload too large')
+
+
+def parse_frame_bytes(data: bytes, *, expect_masked: bool = False, max_payload_size: int | None = None, allow_rsv1: bool = False) -> Frame:
+    if len(data) < 2:
+        raise ProtocolError('incomplete websocket frame')
+    pos = 0
+    b1, b2 = data[pos], data[pos + 1]
+    pos += 2
+    fin = bool(b1 & 0x80)
+    rsv1 = bool(b1 & 0x40)
+    rsv = b1 & 0x70
+    opcode = b1 & 0x0F
+    masked = bool(b2 & 0x80)
+    length = b2 & 0x7F
+    if rsv & 0x30:
+        raise ProtocolError('RSV2/RSV3 bits are not supported')
+    if rsv1 and not allow_rsv1:
+        raise ProtocolError('RSV1 is not negotiated')
+    if expect_masked and not masked:
+        raise ProtocolError('client websocket frames must be masked')
+    if length == 126:
+        if len(data) < pos + 2:
+            raise ProtocolError('incomplete websocket frame')
+        length = struct.unpack('!H', data[pos : pos + 2])[0]
+        pos += 2
+    elif length == 127:
+        if len(data) < pos + 8:
+            raise ProtocolError('incomplete websocket frame')
+        length = struct.unpack('!Q', data[pos : pos + 8])[0]
+        pos += 8
+    _validate_frame_semantics(fin, opcode, length)
+    if max_payload_size is not None and length > max_payload_size:
+        raise ProtocolError('websocket frame exceeds configured max payload size')
+    mask_key = b''
+    if masked:
+        if len(data) < pos + 4:
+            raise ProtocolError('incomplete websocket frame')
+        mask_key = data[pos : pos + 4]
+        pos += 4
+    if len(data) < pos + length:
+        raise ProtocolError('incomplete websocket frame')
+    payload = data[pos : pos + length]
+    if masked:
+        payload = _mask_payload(mask_key, payload)
+    return Frame(fin=fin, opcode=opcode, payload=payload, rsv1=rsv1)
+
+
+async def read_frame(reader: StreamReaderLike, *, max_payload_size: int, expect_masked: bool = True, allow_rsv1: bool = False) -> Frame:
+    header = await reader.readexactly(2)
+    b1, b2 = header[0], header[1]
+    fin = bool(b1 & 0x80)
+    rsv1 = bool(b1 & 0x40)
+    rsv = b1 & 0x70
+    opcode = b1 & 0x0F
+    masked = bool(b2 & 0x80)
+    length = b2 & 0x7F
+    if rsv & 0x30:
+        raise ProtocolError('RSV2/RSV3 bits are not supported')
+    if rsv1 and not allow_rsv1:
+        raise ProtocolError('RSV1 is not negotiated')
+    if expect_masked and not masked:
+        raise ProtocolError('client websocket frames must be masked')
+    if length == 126:
+        length = struct.unpack('!H', await reader.readexactly(2))[0]
+    elif length == 127:
+        length = struct.unpack('!Q', await reader.readexactly(8))[0]
+    _validate_frame_semantics(fin, opcode, length)
+    if length > max_payload_size:
+        raise ProtocolError('websocket frame exceeds configured max payload size')
+    if masked:
+        mask_key = await reader.readexactly(4)
+        payload = await reader.readexactly(length)
+        payload = _mask_payload(mask_key, payload)
+    else:
+        payload = await reader.readexactly(length)
+    return Frame(fin=fin, opcode=opcode, payload=payload, rsv1=rsv1)
+
+
+def serialize_frame(opcode: int, payload: bytes = b'', *, fin: bool = True, mask: bool = False, mask_key: bytes = b'\x00\x00\x00\x00', rsv1: bool = False) -> bytes:
+    _validate_frame_semantics(fin, opcode, len(payload))
+    first = opcode | (0x80 if fin else 0) | (0x40 if rsv1 else 0)
+    length = len(payload)
+    mask_bit = 0x80 if mask else 0
+    if length < 126:
+        head = bytes([first, mask_bit | length])
+    elif length <= 0xFFFF:
+        head = bytes([first, mask_bit | 126]) + struct.pack('!H', length)
+    else:
+        head = bytes([first, mask_bit | 127]) + struct.pack('!Q', length)
+    if not mask:
+        return head + payload
+    masked = _mask_payload(mask_key, payload)
+    return head + mask_key + masked
+
+
+def encode_frame(opcode: int, payload: bytes = b'', *, fin: bool = True, masked: bool = False, mask_key: bytes = b'\x00\x00\x00\x00', rsv1: bool = False) -> bytes:
+    return serialize_frame(opcode, payload, fin=fin, mask=masked, mask_key=mask_key, rsv1=rsv1)
+
+
+def decode_frame(data: bytes, *, expect_masked: bool = False, allow_rsv1: bool = False) -> Frame:
+    return parse_frame_bytes(data, expect_masked=expect_masked, allow_rsv1=allow_rsv1)
+
+
+def encode_close_payload(code: int, reason: str = '') -> bytes:
+    validate_close_code(code)
+    encoded = reason.encode('utf-8')
+    if len(encoded) > 123:
+        raise ProtocolError('close reason too long')
+    return struct.pack('!H', code) + encoded if encoded or code != 1005 else b''
+
+
+def decode_close_payload(payload: bytes) -> tuple[int, str]:
+    if not payload:
+        return 1005, ''
+    if len(payload) == 1:
+        raise ProtocolError('invalid close payload')
+    if len(payload) > 125:
+        raise ProtocolError('control frame payload too large')
+    code = struct.unpack('!H', payload[:2])[0]
+    validate_close_code(code)
+    try:
+        reason = payload[2:].decode('utf-8', 'strict')
+    except UnicodeDecodeError as exc:
+        raise ProtocolError('invalid close reason utf-8') from exc
+    return code, reason