diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fe46003..e34c53e 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,7 @@ Release History - Require ``event_hint`` when constructing ``RemoteProtocolError``. This is an API-breaking change. +- Speed up parsing and serialization of short frames. 1.3.2 (2025-11-20) ------------------ diff --git a/src/wsproto/frame_protocol.py b/src/wsproto/frame_protocol.py index 288276d..273d703 100644 --- a/src/wsproto/frame_protocol.py +++ b/src/wsproto/frame_protocol.py @@ -101,6 +101,10 @@ def iscontrol(self) -> bool: return bool(self & 0x08) +# Invalid opcodes still go through the ordinary parser and its error handling. +_OPCODES = tuple(next((op for op in Opcode if op.value == n), None) for n in range(16)) + + class CloseReason(IntEnum): """ RFC 6455, Section 7.4.1 - Defined Status Codes @@ -224,6 +228,18 @@ class RsvBits(NamedTuple): rsv3: bool +# All possible first header bytes. These immutable fields do not depend on the +# connection, extensions, payload, or how the input was split into chunks. +_HEADER_FIELDS = tuple( + ( + bool(first & FIN_MASK), + RsvBits(bool(first & RSV1_MASK), bool(first & RSV2_MASK), bool(first & RSV3_MASK)), + first & OPCODE_MASK, + ) + for first in range(256) +) + + class Header(NamedTuple): fin: bool rsv: RsvBits @@ -400,18 +416,46 @@ def process_buffer(self) -> Frame | None: return frame def parse_header(self) -> bool: + # Decode complete ordinary short headers without temporary header slices. + # The original path handles extensions, partial headers and every error. + if ( + type(self.buffer) is Buffer + and not self.buffer.bytes_used + and type(self.buffer.buffer) is bytearray + and len(self.buffer.buffer) >= 2 + and (self.buffer.buffer[1] & PAYLOAD_LEN_MASK) <= MAX_PAYLOAD_NORMAL + and type(self) is FrameDecoder + and type(self.extensions) is list + and not self.extensions + and type(self.client) is bool + ): + buffered = self.buffer.buffer + first, second = buffered[0], buffered[1] + short_opcode = _OPCODES[first & OPCODE_MASK] + masked = bool(second & MASK_MASK) + size = 6 if masked else 2 + length = second & PAYLOAD_LEN_MASK + if ( + short_opcode is not None + and not first & (RSV1_MASK | RSV2_MASK | RSV3_MASK) + and (first & FIN_MASK or not short_opcode.iscontrol()) + and masked != self.client + and len(buffered) >= size + ): + fin, rsv, _ = _HEADER_FIELDS[first] + self.masker = XorMaskerSimple(buffered[2:6]) if masked else XorMaskerNull() + del buffered[:size] + self.header = Header(fin, rsv, short_opcode, length, None) + self.effective_opcode = short_opcode + self.payload_required = length if short_opcode.iscontrol() else 0 + self.payload_consumed = 0 + return True data = self.buffer.consume_exactly(2) if data is None: self.buffer.rollback() return False - fin = bool(data[0] & FIN_MASK) - rsv = RsvBits( - bool(data[0] & RSV1_MASK), - bool(data[0] & RSV2_MASK), - bool(data[0] & RSV3_MASK), - ) - opcode = data[0] & OPCODE_MASK + fin, rsv, opcode = _HEADER_FIELDS[data[0]] try: opcode = Opcode(opcode) except ValueError: @@ -520,6 +564,57 @@ def __init__(self, client: bool, extensions: list[Extension]) -> None: self._outbound_opcode: Opcode | None = None + def _parse_complete_data_frame(self) -> Frame | None: + """ + Decode a complete ordinary message, or leave it to the staged parser. + + The caller has checked that the buffer starts with a FIN text or binary + header without reserved bits, that its payload length fits in 7 bits, + and that the whole frame is buffered. + """ + decoder = self._frame_decoder + message = self._message_decoder + if ( + type(decoder) is not FrameDecoder + or type(message) is not MessageDecoder + or decoder.header is not None + or message.opcode is not None + or message.decoder is not None + or type(decoder.extensions) is not list + or decoder.extensions + or type(decoder.client) is not bool + or type(decoder.buffer) is not Buffer + or decoder.buffer.bytes_used + ): + return None + data = decoder.buffer.buffer + masked = bool(data[1] & MASK_MASK) + if masked == decoder.client: + return None + length = data[1] & PAYLOAD_LEN_MASK + start = 2 + payload_start = start + (4 if masked else 0) + end = payload_start + length + payload = data[payload_start:end] + if masked: + payload = XorMaskerSimple(data[start:payload_start]).process(payload) + opcode = Opcode.TEXT if data[0] == 0x81 else Opcode.BINARY + content: bytes | str + if opcode is Opcode.TEXT: + try: + content = payload.decode("utf-8") + except UnicodeDecodeError: + # Let the staged decoder report the error and retain its state. + return None + else: + content = bytes(payload) + del data[:end] + decoder.payload_required = 0 + decoder.payload_consumed = length + decoder.effective_opcode = None + decoder.masker = None + return Frame(opcode, content, True, True) + def _process_close(self, frame: Frame) -> Frame: data = frame.payload assert isinstance(data, (bytes, bytearray)) @@ -563,8 +658,45 @@ def _parse_more_gen(self) -> Generator[Frame | None, None, None]: # disabled extensions in the first place... self.extensions = [ext for ext in self.extensions if ext.enabled()] closed = False + fast_path = ( + type(self._frame_decoder) is FrameDecoder + and type(self._message_decoder) is MessageDecoder + and type(self._frame_decoder.extensions) is list + and not self._frame_decoder.extensions + and type(self._frame_decoder.client) is bool + ) while not closed: + # There is no header or payload to process in an empty ordinary buffer. + if ( + fast_path + and type(self._frame_decoder) is FrameDecoder + and self._frame_decoder.header is None + and type(self._frame_decoder.buffer) is Buffer + and not self._frame_decoder.buffer.buffer + ): + self._frame_decoder.buffer.rollback() + # Match the staged path's release of the previous payload. + frame = None + yield None + continue + if ( + fast_path + and self._frame_decoder.header is None + and self._message_decoder.opcode is None + and len(self._frame_decoder.buffer) >= 2 + and self._frame_decoder.buffer.buffer[0] in (0x81, 0x82) + ): + buffered = self._frame_decoder.buffer.buffer + payload_len_short = buffered[1] & PAYLOAD_LEN_MASK + minimum_size = 2 + payload_len_short + if buffered[1] & MASK_MASK: + minimum_size += 4 + if payload_len_short <= MAX_PAYLOAD_NORMAL and len(buffered) >= minimum_size: + complete = self._parse_complete_data_frame() + if complete is not None: + yield complete + continue frame = self._frame_decoder.process_buffer() if frame is not None: @@ -693,6 +825,13 @@ def _serialize_frame( # -- https://tools.ietf.org/html/rfc6455#section-5.3 masking_key = bytearray(os.urandom(4)) masker = XorMaskerSimple(masking_key) - return bytearray(header + masking_key + masker.process(bytearray(payload))) - + if payload_length > MAX_PAYLOAD_NORMAL: + return bytearray(header + masking_key + masker.process(bytearray(payload))) + header += masking_key + header += masker.process(bytearray(payload)) + return header + + if payload_length <= MAX_PAYLOAD_NORMAL and type(payload) in (bytes, bytearray): + header += payload + return header return bytearray(header + payload) diff --git a/tests/test_frame_protocol.py b/tests/test_frame_protocol.py index 5443348..15d0dcd 100644 --- a/tests/test_frame_protocol.py +++ b/tests/test_frame_protocol.py @@ -1,7 +1,9 @@ from __future__ import annotations +import gc import itertools import struct +import weakref from binascii import unhexlify from codecs import getincrementaldecoder from typing import Dict, Optional, Tuple, Union @@ -1208,3 +1210,192 @@ def test_xor_mask_simple() -> None: assert masker.process(b"some very long data for masking by websocket") == ( b"B]^Q\x11DVFH\x12_[_U\x13PPFR\x14W]A\x14\\S@_X\\T\x14SK\x13CTP@[RYV@" ) + + +@pytest.mark.parametrize("first", range(256)) +def test_first_header_byte_fields(first: int) -> None: + class ObserveHeader(wpext.Extension): + def offer(self) -> bool: + return True + + def frame_inbound_header( + self, proto: fp.FrameDecoder | fp.FrameProtocol, opcode: fp.Opcode, + rsv: fp.RsvBits, payload_length: int, + ) -> fp.RsvBits: + assert type(rsv) is fp.RsvBits + assert rsv == fp.RsvBits(bool(first & 64), bool(first & 32), bool(first & 16)) + return fp.RsvBits(True, True, True) + + decoder = fp.FrameDecoder(client=True, extensions=[ObserveHeader()]) + decoder.receive_bytes(bytes([first, 0])) + if first & 15 not in [0, 1, 2, 8, 9, 10]: + with pytest.raises(fp.ParseFailed, match="Invalid opcode"): + decoder.parse_header() + elif first & 8 and not first & 128: + with pytest.raises(fp.ParseFailed, match="Invalid attempt to fragment control frame"): + decoder.parse_header() + else: + assert decoder.parse_header() + assert decoder.header.fin is bool(first & 128) + assert decoder.header.opcode is fp.Opcode(first & 15) + + +@pytest.mark.parametrize("client", [False, True]) +@pytest.mark.parametrize("size", [0, 1, 125, 126, 65535, 65536]) +def test_serialized_frame_does_not_share_payload(client: bool, size: int) -> None: + payload = bytearray(b"x" * size) + sender = fp.FrameProtocol(client=client, extensions=[]) + wire = sender.send_data(payload) + assert type(wire) is bytearray + payload[:] = b"y" * size + receiver = fp.FrameProtocol(client=not client, extensions=[]) + receiver.receive_bytes(wire) + frames = list(receiver.received_frames()) + assert frames[0].payload == b"x" * size + assert type(frames[0].payload) is bytes + + +@pytest.mark.parametrize("client", [False, True]) +def test_coalesced_complete_fragmented_and_control_frames(client: bool) -> None: + sender = fp.FrameProtocol(client=not client, extensions=[]) + wire = ( + sender.send_data("first") + + sender.send_data("é", fin=False) + + sender.ping(b"heartbeat") + + sender.send_data("last") + + sender.close(1000, "done") + ) + receiver = fp.FrameProtocol(client=client, extensions=[]) + receiver.receive_bytes(wire) + assert list(receiver.received_frames()) == [ + fp.Frame(fp.Opcode.TEXT, "first", True, True), + fp.Frame(fp.Opcode.TEXT, "é", True, False), + fp.Frame(fp.Opcode.PING, b"heartbeat", True, True), + fp.Frame(fp.Opcode.TEXT, "last", True, True), + fp.Frame(fp.Opcode.CLOSE, (fp.CloseReason.NORMAL_CLOSURE, "done"), True, True), + ] + + +@pytest.mark.parametrize("client", [False, True]) +def test_bad_complete_text_preserves_parser_failure(client: bool) -> None: + sender = fp.FrameProtocol(client=not client, extensions=[]) + wire = sender._serialize_frame(fp.Opcode.TEXT, b"\xff") + sender.send_data("later") + receiver = fp.FrameProtocol(client=client, extensions=[]) + receiver.receive_bytes(wire) + with pytest.raises(fp.ParseFailed) as error: + list(receiver.received_frames()) + assert error.value.code is fp.CloseReason.INVALID_FRAME_PAYLOAD_DATA + assert list(receiver.received_frames()) == [] + + +@pytest.mark.parametrize("client", [True, 1, 2, "client", [1]]) +def test_client_role_preserves_mask_validation(client) -> None: + sender = fp.FrameProtocol(client=True, extensions=[]) + receiver = fp.FrameProtocol(client=client, extensions=[]) + receiver.receive_bytes(sender.send_data("hello")) + with pytest.raises(fp.ParseFailed, match="client received unexpected masked frame"): + list(receiver.received_frames()) + + +def test_extensions_added_between_frames_are_observed() -> None: + calls = [] + + class Observe(wpext.Extension): + def offer(self) -> bool: + return True + + def frame_inbound_header(self, proto, opcode, rsv, payload_length): + calls.append("header") + return fp.RsvBits(False, False, False) + + def frame_inbound_payload_data(self, proto, data): + calls.append("payload") + return data + + def frame_inbound_complete(self, proto, fin): + calls.append("complete") + return None + + sender = fp.FrameProtocol(client=False, extensions=[]) + receiver = fp.FrameProtocol(client=True, extensions=[]) + receiver.receive_bytes(sender.send_data("first") + sender.send_data("second")) + frames = receiver.received_frames() + assert next(frames).payload == "first" + receiver._frame_decoder.extensions.append(Observe()) + assert next(frames).payload == "second" + assert calls == ["header", "payload", "complete"] + assert list(frames) == [] + + +@pytest.mark.parametrize("client", [False, True]) +def test_control_frame_waits_for_every_payload_byte(client: bool) -> None: + sender = fp.FrameProtocol(client=not client, extensions=[]) + receiver = fp.FrameProtocol(client=client, extensions=[]) + wire = sender.ping(b"heartbeat") + for part in wire[:-1]: + receiver.receive_bytes(bytes([part])) + assert list(receiver.received_frames()) == [] + receiver.receive_bytes(wire[-1:]) + assert list(receiver.received_frames()) == [ + fp.Frame(fp.Opcode.PING, b"heartbeat", True, True), + ] + assert list(receiver.received_frames()) == [] + receiver.receive_bytes(sender.pong()) + assert list(receiver.received_frames()) == [ + fp.Frame(fp.Opcode.PONG, b"", True, True), + ] + + +def test_replaced_frame_decoder_is_called_with_empty_input() -> None: + calls: list[bool] = [] + + class Observe(fp.FrameDecoder): + def process_buffer(self) -> fp.Frame | None: + calls.append(True) + return super().process_buffer() + + receiver = fp.FrameProtocol(client=True, extensions=[]) + receiver.receive_bytes(b"\x81\x01x") + frames = receiver.received_frames() + assert next(frames).payload == "x" + receiver._frame_decoder = Observe(client=True) + assert list(frames) == [] + assert calls == [True] + + +def test_buffer_subclass_is_used_when_empty() -> None: + calls: list[int] = [] + + class Observe(fp.Buffer): + def consume_exactly(self, nbytes: int) -> bytearray | None: + calls.append(nbytes) + return super().consume_exactly(nbytes) + + receiver = fp.FrameProtocol(client=True, extensions=[]) + receiver._frame_decoder.buffer = Observe() + assert list(receiver.received_frames()) == [] + assert calls == [2] + + +def test_empty_buffer_releases_staged_payload() -> None: + class Payload(str): + pass + + refs: list[weakref.ReferenceType[Payload]] = [] + receiver = fp.FrameProtocol(client=True, extensions=[]) + original = receiver._message_decoder.process_frame + + def observe(frame: fp.Frame) -> fp.Frame: + result = original(frame) + payload = Payload(result.payload) + refs.append(weakref.ref(payload)) + return result._replace(payload=payload) + + receiver._message_decoder.process_frame = observe + sender = fp.FrameProtocol(client=False, extensions=[]) + receiver.receive_bytes(sender.send_data("x" * 126)) + frames = list(receiver.received_frames()) + assert frames[0].payload == "x" * 126 + del frames + gc.collect() + assert refs[0]() is None