"""Reader for WebSocket protocol versions 13 and 8.""" import asyncio import builtins import sys import weakref from collections import deque from typing import Final from ..base_protocol import BaseProtocol from ..compression_utils import TooManyMembersError, ZLibDecompressor from ..helpers import _EXC_SENTINEL, set_exception from ..log import ws_logger from ..streams import EofStream from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask from .models import ( WS_DEFLATE_TRAILING, WebSocketError, WSCloseCode, WSMessage, WSMessageTextBytes, WSMsgType, ) ALLOWED_CLOSE_CODES: set[int] = {int(i) for i in WSCloseCode} # States for the reader, used to parse the WebSocket frame # integer values are used so they can be cythonized READ_HEADER = 1 READ_PAYLOAD_LENGTH = 2 READ_PAYLOAD_MASK = 3 READ_PAYLOAD = 4 # Largest declared payload length the reader can represent: the compiled # reader stores it in a Py_ssize_t, which holds 2**31-1 on the 32-bit builds # (the win32 and armv7l wheels) and 2**63-1 everywhere else. # TODO: Remove when we drop 32 bit support (and from reader_c.pxd). MAX_PAYLOAD_LEN = sys.maxsize WS_MSG_TYPE_BINARY = WSMsgType.BINARY WS_MSG_TYPE_TEXT = WSMsgType.TEXT # WSMsgType values unpacked so they can by cythonized to ints OP_CODE_NOT_SET = -1 OP_CODE_CONTINUATION = WSMsgType.CONTINUATION.value OP_CODE_TEXT = WSMsgType.TEXT.value OP_CODE_BINARY = WSMsgType.BINARY.value OP_CODE_CLOSE = WSMsgType.CLOSE.value OP_CODE_PING = WSMsgType.PING.value OP_CODE_PONG = WSMsgType.PONG.value EMPTY_FRAME_ERROR = (True, b"") EMPTY_FRAME = (False, b"") COMPRESSED_NOT_SET = -1 COMPRESSED_FALSE = 0 COMPRESSED_TRUE = 1 TUPLE_NEW = tuple.__new__ # Overhead added to each message to ensure that tiny messages can't use # unreasonable amounts of memory. MSG_SIZE_OVERHEAD: Final[int] = 128 STALLED_READER_COLLECTED: Final[str] = ( "WebSocketReader was garbage collected while stalled; " "callers of set_parser() must hold a strong reference" ) cython_int = int # Typed to int in Python, but cython with use a signed int in the pxd class WebSocketDataQueue: """WebSocketDataQueue resumes and pauses an underlying stream. It is a destination for WebSocket data. """ def __init__( self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop ) -> None: self._size = 0 self._protocol = protocol self._limit = limit * 2 self._loop = loop self._eof = False self._waiter: asyncio.Future[None] | None = None self._exception: BaseException | None = None self._buffer: deque[tuple[WSMessage | WSMessageTextBytes, int]] = deque() self._get_buffer = self._buffer.popleft self._put_buffer = self._buffer.append self._stalled_reader: "weakref.ref[WebSocketReader] | None" = None def is_eof(self) -> bool: return self._eof def exception(self) -> BaseException | None: return self._exception def set_exception( self, exc: BaseException, exc_cause: builtins.BaseException = _EXC_SENTINEL, ) -> None: self._eof = True self._exception = exc if (waiter := self._waiter) is not None: self._waiter = None set_exception(waiter, exc, exc_cause) def _release_waiter(self) -> None: if (waiter := self._waiter) is None: return self._waiter = None if not waiter.done(): waiter.set_result(None) def feed_eof(self) -> None: self._eof = True self._release_waiter() self._exception = None # Break cyclic references def feed_data( self, data: "WSMessage | WSMessageTextBytes", size: "cython_int" ) -> None: self._size += size + MSG_SIZE_OVERHEAD self._put_buffer((data, size)) self._release_waiter() if self._size > self._limit and not self._protocol._reading_paused: self._protocol.pause_reading() async def read(self) -> WSMessage | WSMessageTextBytes: if not self._buffer and not self._eof: assert not self._waiter self._waiter = self._loop.create_future() try: await self._waiter except (asyncio.CancelledError, asyncio.TimeoutError): self._waiter = None raise return self._read_from_buffer() def _read_from_buffer(self) -> WSMessage | WSMessageTextBytes: if self._buffer: data, size = self._get_buffer() self._size -= size + MSG_SIZE_OVERHEAD if self._stalled_reader is not None and self._size <= self._limit // 2: # Resume parsing once the queue drains to the low-water mark. # Each resume re-slices the parser's unparsed tail, so waiting # for headroom makes a drain cost one copy per batch of # messages instead of one per message. if (reader := self._stalled_reader()) is not None: reader.feed_data(b"") else: # The stash died with the reader. Deliver what was already # queued, then surface the contract violation on the next # read instead of hanging. Log as well, since a caller that # stops reading early never sees the deferred exception. A # real failure that was already recorded stays the reported # cause. self._stalled_reader = None ws_logger.warning(STALLED_READER_COLLECTED) if self._exception is None: self.set_exception(RuntimeError(STALLED_READER_COLLECTED)) # Resuming the transport while a stash remains would admit # another socket read into the tail for every couple of messages # drained, moving the memory bound from the queue into the tail. if ( self._stalled_reader is None and self._size < self._limit and self._protocol._reading_paused ): self._protocol.resume_reading() return data if self._exception is not None: raise self._exception raise EofStream class WebSocketReader: def __init__( self, queue: WebSocketDataQueue, max_msg_size: int, compress: bool, decode_text: bool, ) -> None: self.queue = queue self._max_msg_size = max_msg_size self._decode_text = decode_text # Parked on the queue while parsing is stalled; created once so # stalling does not allocate. self._weak_self = weakref.ref(self) self._exc: Exception | None = None self._partial = bytearray() self._state = READ_HEADER self._opcode: int = OP_CODE_NOT_SET self._frame_fin = False self._frame_opcode: int = OP_CODE_NOT_SET # Reads of an in-flight frame, joined once when it completes. self._payload_fragments: list[bytes] = [] # Fold reads into _payload_buffer past this count to bound the object # count (bytes are bounded by max_msg_size). self._max_fragments = max(1024, max_msg_size // 256) if max_msg_size else 0 self._payload_buffer = bytearray() self._frame_payload_len = 0 self._tail: bytes = b"" self._has_mask = False self._frame_mask: bytes | None = None self._payload_bytes_to_read = 0 self._payload_len_flag = 0 self._compressed: int = COMPRESSED_NOT_SET self._decompressobj: ZLibDecompressor | None = None self._compress = compress def feed_eof(self) -> None: self.queue.feed_eof() # data can be bytearray on Windows because proactor event loop uses bytearray # and asyncio types this to Union[bytes, bytearray, memoryview] so we need # coerce data to bytes if it is not def feed_data(self, data: bytes | bytearray | memoryview) -> tuple[bool, bytes]: if type(data) is not bytes: data = bytes(data) if self._exc is not None: return True, data try: self._feed_data(data) except Exception as exc: self._exc = exc set_exception(self.queue, exc) return EMPTY_FRAME_ERROR return EMPTY_FRAME def _handle_frame( self, fin: bool, opcode: int | cython_int, # Union intended: Cython pxd uses C int payload: bytes | bytearray, compressed: int | cython_int, # Union intended: Cython pxd uses C int ) -> None: msg: WSMessage if opcode in {OP_CODE_TEXT, OP_CODE_BINARY, OP_CODE_CONTINUATION}: # Validate continuation frames before processing if opcode == OP_CODE_CONTINUATION and self._opcode == OP_CODE_NOT_SET: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Continuation frame for non started message", ) # load text/binary if not fin: # got partial frame payload if opcode != OP_CODE_CONTINUATION: # https://datatracker.ietf.org/doc/html/rfc6455#section-5.4 if self._opcode != OP_CODE_NOT_SET: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "The opcode in non-fin frame is expected " f"to be zero, got {opcode!r}", ) self._opcode = opcode self._partial += payload return has_partial = bool(self._partial) if opcode == OP_CODE_CONTINUATION: opcode = self._opcode self._opcode = OP_CODE_NOT_SET # previous frame was non finished # we should get continuation opcode elif has_partial: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "The opcode in non-fin frame is expected " f"to be zero, got {opcode!r}", ) assembled_payload: bytes | bytearray if has_partial: assembled_payload = self._partial + payload self._partial.clear() else: assembled_payload = payload # Decompress process must to be done after all packets # received. if compressed: if not self._decompressobj: self._decompressobj = ZLibDecompressor(suppress_deflate_header=True) # XXX: It's possible that the zlib backend (isal is known to # do this, maybe others too?) will return max_length bytes, # but internally buffer more data such that the payload is # >max_length, so we return one extra byte and if we're able # to do that, then the message is too big. try: payload_merged = self._decompressobj.decompress_sync( assembled_payload + WS_DEFLATE_TRAILING, ( self._max_msg_size + 1 if self._max_msg_size else self._max_msg_size ), ) except TooManyMembersError as exc: raise WebSocketError( WSCloseCode.MESSAGE_TOO_BIG, "Compressed message has too many deflate members", ) from exc if self._max_msg_size and len(payload_merged) > self._max_msg_size: raise WebSocketError( WSCloseCode.MESSAGE_TOO_BIG, f"Decompressed message exceeds size limit {self._max_msg_size}", ) elif type(assembled_payload) is bytes: payload_merged = assembled_payload else: payload_merged = bytes(assembled_payload) if opcode == OP_CODE_TEXT: if self._decode_text: try: text = payload_merged.decode("utf-8") except UnicodeDecodeError as exc: raise WebSocketError( WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message" ) from exc # XXX: The Text and Binary messages here can be a performance # bottleneck, so we use tuple.__new__ to improve performance. # This is not type safe, but many tests should fail in # test_client_ws_functional.py if this is wrong. self.queue.feed_data( TUPLE_NEW(WSMessage, (WS_MSG_TYPE_TEXT, text, "")), len(payload_merged), ) else: # Return raw bytes for TEXT messages when decode_text=False self.queue.feed_data( TUPLE_NEW( WSMessageTextBytes, (WS_MSG_TYPE_TEXT, payload_merged, "") ), len(payload_merged), ) else: self.queue.feed_data( TUPLE_NEW(WSMessage, (WS_MSG_TYPE_BINARY, payload_merged, "")), len(payload_merged), ) elif opcode == OP_CODE_CLOSE: if len(payload) >= 2: close_code = UNPACK_CLOSE_CODE(payload[:2])[0] if close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, f"Invalid close code: {close_code}", ) try: close_message = payload[2:].decode("utf-8") except UnicodeDecodeError as exc: raise WebSocketError( WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message" ) from exc msg = TUPLE_NEW(WSMessage, (WSMsgType.CLOSE, close_code, close_message)) elif payload: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, f"Invalid close frame: {fin} {opcode} {payload!r}", ) else: msg = TUPLE_NEW(WSMessage, (WSMsgType.CLOSE, 0, "")) self.queue.feed_data(msg, 0) elif opcode == OP_CODE_PING: msg = TUPLE_NEW(WSMessage, (WSMsgType.PING, payload, "")) self.queue.feed_data(msg, len(payload)) elif opcode == OP_CODE_PONG: msg = TUPLE_NEW(WSMessage, (WSMsgType.PONG, payload, "")) self.queue.feed_data(msg, len(payload)) else: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}" ) def _feed_data(self, data: bytes) -> None: """Return the next frame from the socket.""" self.queue._stalled_reader = None if self._tail: data, self._tail = self._tail + data, b"" start_pos = 0 data_len = len(data) data_cstr = data while True: if start_pos < data_len and self.queue._size > self.queue._limit: # Over the high-water mark with unparsed bytes left: stash the # remainder and stall. Gating on unparsed bytes keeps a read # that ended on a frame boundary from arming an empty stall, # which would hold the transport paused with nothing to drain. self.queue._stalled_reader = self._weak_self break # read header if self._state == READ_HEADER: if data_len - start_pos < 2: break first_byte = data_cstr[start_pos] second_byte = data_cstr[start_pos + 1] start_pos += 2 fin = (first_byte >> 7) & 1 rsv1 = (first_byte >> 6) & 1 rsv2 = (first_byte >> 5) & 1 rsv3 = (first_byte >> 4) & 1 opcode = first_byte & 0xF # frame-fin = %x0 ; more frames of this message follow # / %x1 ; final frame of this message # frame-rsv1 = %x0 ; # 1 bit, MUST be 0 unless negotiated otherwise # frame-rsv2 = %x0 ; # 1 bit, MUST be 0 unless negotiated otherwise # frame-rsv3 = %x0 ; # 1 bit, MUST be 0 unless negotiated otherwise # # Remove rsv1 from this test for deflate development if rsv2 or rsv3 or (rsv1 and not self._compress): raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Received frame with non-zero reserved bits", ) if opcode not in { OP_CODE_CONTINUATION, OP_CODE_TEXT, OP_CODE_BINARY, OP_CODE_CLOSE, OP_CODE_PING, OP_CODE_PONG, }: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}", ) if opcode > 0x7 and fin == 0: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Received fragmented control frame", ) has_mask = (second_byte >> 7) & 1 length = second_byte & 0x7F # Control frames MUST have a payload # length of 125 bytes or less if opcode > 0x7 and length > 125: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Control frame payload cannot be larger than 125 bytes", ) # Control frames (opcode > 0x7) may be interleaved between the # fragments of a data message and never carry the per-message # compressed bit, so they must not touch the compression state. # https://datatracker.ietf.org/doc/html/rfc6455#section-5.4 # https://datatracker.ietf.org/doc/html/rfc7692#section-6.1 if opcode > 0x7: if rsv1: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Received frame with non-zero reserved bits", ) else: # Set compress status if last package is FIN # OR set compress status if this is first fragment # Raise error if not first fragment with rsv1 = 0x1 if self._frame_fin or self._compressed == COMPRESSED_NOT_SET: self._compressed = COMPRESSED_TRUE if rsv1 else COMPRESSED_FALSE elif rsv1: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Received frame with non-zero reserved bits", ) self._frame_fin = bool(fin) self._frame_opcode = opcode self._has_mask = bool(has_mask) self._payload_len_flag = length self._state = READ_PAYLOAD_LENGTH # read payload length if self._state == READ_PAYLOAD_LENGTH: len_flag = self._payload_len_flag if len_flag == 126: if data_len - start_pos < 2: break first_byte = data_cstr[start_pos] second_byte = data_cstr[start_pos + 1] start_pos += 2 self._payload_bytes_to_read = first_byte << 8 | second_byte elif len_flag > 126: if data_len - start_pos < 8: break # The declared length is an unsigned 64-bit integer that # does not necessarily fit _payload_bytes_to_read. frame_len = UNPACK_LEN3(data, start_pos)[0] if frame_len > MAX_PAYLOAD_LEN: raise WebSocketError( WSCloseCode.MESSAGE_TOO_BIG, f"Message size {int(frame_len) + len(self._partial)} " f"exceeds limit {self._max_msg_size or MAX_PAYLOAD_LEN}", ) self._payload_bytes_to_read = frame_len start_pos += 8 else: self._payload_bytes_to_read = len_flag # Reject oversized data frames before buffering any payload # bytes. Control frames are capped at 125 bytes (checked in # READ_HEADER) so only text/binary/continuation need this. if self._max_msg_size and self._frame_opcode in { OP_CODE_TEXT, OP_CODE_BINARY, OP_CODE_CONTINUATION, }: # partial_len declared in reader_c.pxd to keep it in C. partial_len = len(self._partial) # payload_bytes_to_read is a signed Py_ssize_t C value, # use subtraction here to avoid an integer overflow. if self._payload_bytes_to_read >= self._max_msg_size - partial_len: raise WebSocketError( WSCloseCode.MESSAGE_TOO_BIG, f"Message size {int(self._payload_bytes_to_read) + partial_len} " f"exceeds limit {self._max_msg_size}", ) self._state = READ_PAYLOAD_MASK if self._has_mask else READ_PAYLOAD # read payload mask if self._state == READ_PAYLOAD_MASK: if data_len - start_pos < 4: break self._frame_mask = data_cstr[start_pos : start_pos + 4] start_pos += 4 self._state = READ_PAYLOAD if self._state == READ_PAYLOAD: chunk_len = data_len - start_pos if self._payload_bytes_to_read >= chunk_len: f_end_pos = data_len self._payload_bytes_to_read -= chunk_len else: f_end_pos = start_pos + self._payload_bytes_to_read self._payload_bytes_to_read = 0 had_fragments = self._frame_payload_len self._frame_payload_len += f_end_pos - start_pos f_start_pos = start_pos start_pos = f_end_pos if self._payload_bytes_to_read != 0: if f_start_pos < f_end_pos: # skip a header-only read self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos]) if ( self._max_fragments and len(self._payload_fragments) > self._max_fragments ): # Fold to bound the object count. Not a pause: nothing # resumes reading until the frame is queued. self._payload_buffer += b"".join(self._payload_fragments) self._payload_fragments.clear() break payload: bytes | bytearray if had_fragments: self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos]) if self._payload_buffer: # folded prefix self._payload_buffer += b"".join(self._payload_fragments) if self._has_mask: assert self._frame_mask is not None websocket_mask(self._frame_mask, self._payload_buffer) payload = self._payload_buffer self._payload_buffer = bytearray() # detach; payload aliases it elif self._has_mask: assert self._frame_mask is not None payload_bytearray = bytearray(b"".join(self._payload_fragments)) websocket_mask(self._frame_mask, payload_bytearray) payload = payload_bytearray else: payload = b"".join(self._payload_fragments) self._payload_fragments.clear() elif self._has_mask: assert self._frame_mask is not None payload_bytearray = data_cstr[f_start_pos:f_end_pos] # type: ignore[assignment] if type(payload_bytearray) is not bytearray: # Cython will do the conversion for us # but we need to do it for Python and we # will always get here in Python payload_bytearray = bytearray(payload_bytearray) websocket_mask(self._frame_mask, payload_bytearray) payload = payload_bytearray else: payload = data_cstr[f_start_pos:f_end_pos] self._handle_frame( self._frame_fin, self._frame_opcode, payload, self._compressed ) self._frame_payload_len = 0 self._state = READ_HEADER # XXX: Cython needs slices to be bounded, so we can't omit the slice end here. self._tail = data_cstr[start_pos:data_len] if start_pos < data_len else b""