From c80d3fb7923c1957210eb963cf380bd86e5276de Mon Sep 17 00:00:00 2001 From: Dickson Date: Fri, 6 Nov 2026 02:58:49 -0500 Subject: [PATCH 1/3] Reduce frame header and output buffer allocations FrameDecoder.parse_header() rebuilt the same FIN/RSV/opcode fields, including a new RsvBits tuple, for every frame. There are only 256 possible first bytes, so look the fields up in a precomputed table. _serialize_frame() built the output for short frames with chained concatenation (header + mask + payload, then another bytearray copy). Append to the fresh header bytearray instead. Frames with extended payload lengths and payloads of other types keep the original expression, which is faster for large bytearrays on CPython. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.rst | 1 + src/wsproto/frame_protocol.py | 31 +++++++++----- tests/test_frame_protocol.py | 76 +++++++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 9 deletions(-) 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..4509636 100644 --- a/src/wsproto/frame_protocol.py +++ b/src/wsproto/frame_protocol.py @@ -224,6 +224,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 @@ -405,13 +417,7 @@ def parse_header(self) -> bool: 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: @@ -693,6 +699,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..b98a2e0 100644 --- a/tests/test_frame_protocol.py +++ b/tests/test_frame_protocol.py @@ -1208,3 +1208,79 @@ 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()) == [] From 448a0f2b980958e8a68144d72c423cecaa345f6f Mon Sep 17 00:00:00 2001 From: Dickson Date: Fri, 6 Nov 2026 02:58:49 -0500 Subject: [PATCH 2/3] Decode complete short data frames directly A complete, unfragmented text or binary frame with a short payload went through the staged parser: header parsing, a Frame, a second Frame from the message decoder, and for text a fresh incremental UTF-8 decoder, even though every byte is already buffered. When no extensions are in use and no frame or message is in progress, decode such a frame straight from the buffer and return the same Frame. Everything else (partial input, fragmented messages, control frames, extended lengths, reserved bits, wrong masking, invalid UTF-8) is left to the staged parser before any input is consumed, so errors and close codes are unchanged. Co-Authored-By: Claude Opus 5.5 --- src/wsproto/frame_protocol.py | 75 +++++++++++++++++++++++++++++++++++ tests/test_frame_protocol.py | 39 ++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/src/wsproto/frame_protocol.py b/src/wsproto/frame_protocol.py index 4509636..75a841a 100644 --- a/src/wsproto/frame_protocol.py +++ b/src/wsproto/frame_protocol.py @@ -526,6 +526,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)) @@ -569,8 +620,32 @@ 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: + 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: diff --git a/tests/test_frame_protocol.py b/tests/test_frame_protocol.py index b98a2e0..acaddfe 100644 --- a/tests/test_frame_protocol.py +++ b/tests/test_frame_protocol.py @@ -1284,3 +1284,42 @@ def test_bad_complete_text_preserves_parser_failure(client: bool) -> None: 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) == [] From 4875a851877864a66723cacd077e5d63779b7560 Mon Sep 17 00:00:00 2001 From: Dickson Date: Fri, 6 Nov 2026 02:58:49 -0500 Subject: [PATCH 3/3] Parse short frame headers without temporary slices Control frames and continuation frames still go through parse_header(), which consumes the header in several small slices. For a complete header with a 7-bit length and no extensions in use, decode it in place using a 16-entry opcode table. Partial headers, extended lengths, reserved bits, invalid opcodes and fragmented control frames take the original path before anything is consumed. received_frames() also skips the staged parser when the buffer is empty. It drops its reference to the previous frame before yielding, as the staged loop does, so a delivered payload is not kept alive while waiting for more data. Co-Authored-By: Claude Opus 5.5 --- src/wsproto/frame_protocol.py | 51 +++++++++++++++++++++++ tests/test_frame_protocol.py | 76 +++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) diff --git a/src/wsproto/frame_protocol.py b/src/wsproto/frame_protocol.py index 75a841a..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 @@ -412,6 +416,40 @@ 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() @@ -629,6 +667,19 @@ def _parse_more_gen(self) -> Generator[Frame | None, None, None]: ) 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 diff --git a/tests/test_frame_protocol.py b/tests/test_frame_protocol.py index acaddfe..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 @@ -1323,3 +1325,77 @@ def frame_inbound_complete(self, proto, fin): 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