Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/_websocket/reader_py.py: 16%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Reader for WebSocket protocol versions 13 and 8."""
3import asyncio
4import builtins
5import sys
6import weakref
7from collections import deque
9from ..base_protocol import BaseProtocol
10from ..compression_utils import TooManyMembersError, ZLibDecompressor
11from ..helpers import _EXC_SENTINEL, set_exception
12from ..log import ws_logger
13from ..streams import EofStream
14from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask
15from .models import (
16 WS_DEFLATE_TRAILING,
17 WebSocketError,
18 WSCloseCode,
19 WSMessage,
20 WSMessageBinary,
21 WSMessageClose,
22 WSMessagePing,
23 WSMessagePong,
24 WSMessageText,
25 WSMessageTextBytes,
26 WSMsgType,
27)
29# ABNORMAL_CLOSURE is used internally, should never be accepted from a client.
30# https://datatracker.ietf.org/doc/html/rfc6455#section-7.4.1
31ALLOWED_CLOSE_CODES = {int(i) for i in WSCloseCode} - {
32 int(WSCloseCode.ABNORMAL_CLOSURE)
33}
35# States for the reader, used to parse the WebSocket frame
36# integer values are used so they can be cythonized
37READ_HEADER = 1
38READ_PAYLOAD_LENGTH = 2
39READ_PAYLOAD_MASK = 3
40READ_PAYLOAD = 4
42# Largest declared payload length the reader can represent: the compiled
43# reader stores it in a Py_ssize_t, which holds 2**31-1 on the 32-bit builds
44# (the win32 and armv7l wheels) and 2**63-1 everywhere else.
45# TODO: Remove when we drop 32 bit support (and from reader_c.pxd).
46MAX_PAYLOAD_LEN = sys.maxsize
48WS_MSG_TYPE_BINARY = WSMsgType.BINARY
49WS_MSG_TYPE_TEXT = WSMsgType.TEXT
51# WSMsgType values unpacked so they can by cythonized to ints
52OP_CODE_NOT_SET = -1
53OP_CODE_CONTINUATION = WSMsgType.CONTINUATION.value
54OP_CODE_TEXT = WSMsgType.TEXT.value
55OP_CODE_BINARY = WSMsgType.BINARY.value
56OP_CODE_CLOSE = WSMsgType.CLOSE.value
57OP_CODE_PING = WSMsgType.PING.value
58OP_CODE_PONG = WSMsgType.PONG.value
60EMPTY_FRAME_ERROR = (True, b"")
61EMPTY_FRAME = (False, b"")
63COMPRESSED_NOT_SET = -1
64COMPRESSED_FALSE = 0
65COMPRESSED_TRUE = 1
67TUPLE_NEW = tuple.__new__
69# Overhead added to each message to ensure that tiny messages can't use
70# unreasonable amounts of memory.
71MSG_SIZE_OVERHEAD = 128
73STALLED_READER_COLLECTED = (
74 "WebSocketReader was garbage collected while stalled; "
75 "callers of set_parser() must hold a strong reference"
76)
78cython_int = int # Typed to int in Python, but cython with use a signed int in the pxd
81class WebSocketDataQueue:
82 """WebSocketDataQueue resumes and pauses an underlying stream.
84 It is a destination for WebSocket data.
85 """
87 def __init__(
88 self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop
89 ) -> None:
90 self._size = 0
91 self._protocol = protocol
92 self._limit = limit * 2
93 self._loop = loop
94 self._eof = False
95 self._waiter: asyncio.Future[None] | None = None
96 self._exception: type[BaseException] | BaseException | None = None
97 self._buffer: deque[WSMessage] = deque()
98 self._get_buffer = self._buffer.popleft
99 self._put_buffer = self._buffer.append
100 self._stalled_reader: "weakref.ref[WebSocketReader] | None" = None
102 def is_eof(self) -> bool:
103 return self._eof
105 def exception(self) -> type[BaseException] | BaseException | None:
106 return self._exception
108 def set_exception(
109 self,
110 exc: type[BaseException] | BaseException,
111 exc_cause: builtins.BaseException = _EXC_SENTINEL,
112 ) -> None:
113 self._eof = True
114 self._exception = exc
115 if (waiter := self._waiter) is not None:
116 self._waiter = None
117 set_exception(waiter, exc, exc_cause)
119 def _release_waiter(self) -> None:
120 if (waiter := self._waiter) is None:
121 return
122 self._waiter = None
123 if not waiter.done():
124 waiter.set_result(None)
126 def feed_eof(self) -> None:
127 self._eof = True
128 self._release_waiter()
129 self._exception = None # Break cyclic references
131 def feed_data(self, data: "WSMessage") -> None:
132 # Unbox into the typed local before adding, so Cython keeps the sum in
133 # C instead of boxing MSG_SIZE_OVERHEAD for a Python-level add.
134 size = data.size
135 self._size += size + MSG_SIZE_OVERHEAD
136 self._put_buffer(data)
137 self._release_waiter()
138 if self._size > self._limit and not self._protocol._reading_paused:
139 self._protocol.pause_reading()
141 async def read(self) -> WSMessage:
142 if not self._buffer and not self._eof:
143 assert not self._waiter
144 self._waiter = self._loop.create_future()
145 try:
146 await self._waiter
147 except (asyncio.CancelledError, asyncio.TimeoutError):
148 self._waiter = None
149 raise
150 return self._read_from_buffer()
152 def _read_from_buffer(self) -> WSMessage:
153 if self._buffer:
154 data = self._get_buffer()
155 size = data.size
156 self._size -= size + MSG_SIZE_OVERHEAD
157 if self._stalled_reader is not None and self._size <= self._limit // 2:
158 # Resume parsing once the queue drains to the low-water mark.
159 # Each resume re-slices the parser's unparsed tail, so waiting
160 # for headroom makes a drain cost one copy per batch of
161 # messages instead of one per message.
162 if (reader := self._stalled_reader()) is not None:
163 reader.feed_data(b"")
164 else:
165 # The stash died with the reader. Deliver what was already
166 # queued, then surface the contract violation on the next
167 # read instead of hanging. Log as well, since a caller that
168 # stops reading early never sees the deferred exception. A
169 # real failure that was already recorded stays the reported
170 # cause.
171 self._stalled_reader = None
172 ws_logger.warning(STALLED_READER_COLLECTED)
173 if self._exception is None:
174 self.set_exception(RuntimeError(STALLED_READER_COLLECTED))
175 # Resuming the transport while a stash remains would admit
176 # another socket read into the tail for every couple of messages
177 # drained, moving the memory bound from the queue into the tail.
178 if (
179 self._stalled_reader is None
180 and self._size < self._limit
181 and self._protocol._reading_paused
182 ):
183 self._protocol.resume_reading()
184 return data
185 if self._exception is not None:
186 raise self._exception
187 raise EofStream
190class WebSocketReader:
191 def __init__(
192 self,
193 queue: WebSocketDataQueue,
194 max_msg_size: int,
195 compress: bool,
196 decode_text: bool,
197 ) -> None:
198 self.queue = queue
199 self._max_msg_size = max_msg_size
200 self._decode_text = decode_text
201 # Parked on the queue while parsing is stalled; created once so
202 # stalling does not allocate.
203 self._weak_self = weakref.ref(self)
205 self._exc: Exception | None = None
206 self._partial = bytearray()
207 self._state = READ_HEADER
209 self._opcode: int = OP_CODE_NOT_SET
210 self._frame_fin = False
211 self._frame_opcode: int = OP_CODE_NOT_SET
212 # Reads of an in-flight frame, joined once when it completes.
213 self._payload_fragments: list[bytes] = []
214 # Fold reads into _payload_buffer past this count to bound the object
215 # count (bytes are bounded by max_msg_size).
216 self._max_fragments = max(1024, max_msg_size // 256) if max_msg_size else 0
217 self._payload_buffer = bytearray()
218 self._frame_payload_len = 0
220 self._tail: bytes = b""
221 self._has_mask = False
222 self._frame_mask: bytes | None = None
223 self._payload_bytes_to_read = 0
224 self._payload_len_flag = 0
225 self._compressed: int = COMPRESSED_NOT_SET
226 self._decompressobj: ZLibDecompressor | None = None
227 self._compress = compress
229 def feed_eof(self) -> None:
230 self.queue.feed_eof()
232 # data can be bytearray on Windows because proactor event loop uses bytearray
233 # and asyncio types this to Union[bytes, bytearray, memoryview] so we need
234 # coerce data to bytes if it is not
235 def feed_data(self, data: bytes | bytearray | memoryview) -> tuple[bool, bytes]:
236 if type(data) is not bytes:
237 data = bytes(data)
239 if self._exc is not None:
240 return True, data
242 try:
243 self._feed_data(data)
244 except Exception as exc:
245 self._exc = exc
246 set_exception(self.queue, exc)
247 return EMPTY_FRAME_ERROR
249 return EMPTY_FRAME
251 def _handle_frame(
252 self,
253 fin: bool,
254 opcode: int | cython_int, # Union intended: Cython pxd uses C int
255 payload: bytes | bytearray,
256 compressed: int | cython_int, # Union intended: Cython pxd uses C int
257 ) -> None:
258 msg: WSMessage
259 if opcode in {OP_CODE_TEXT, OP_CODE_BINARY, OP_CODE_CONTINUATION}:
260 # Validate continuation frames before processing
261 if opcode == OP_CODE_CONTINUATION and self._opcode == OP_CODE_NOT_SET:
262 raise WebSocketError(
263 WSCloseCode.PROTOCOL_ERROR,
264 "Continuation frame for non started message",
265 )
267 # load text/binary
268 if not fin:
269 # got partial frame payload
270 if opcode != OP_CODE_CONTINUATION:
271 # https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
272 if self._opcode != OP_CODE_NOT_SET:
273 raise WebSocketError(
274 WSCloseCode.PROTOCOL_ERROR,
275 "The opcode in non-fin frame is expected "
276 f"to be zero, got {opcode!r}",
277 )
278 self._opcode = opcode
279 self._partial += payload
280 return
282 has_partial = bool(self._partial)
283 if opcode == OP_CODE_CONTINUATION:
284 opcode = self._opcode
285 self._opcode = OP_CODE_NOT_SET
286 # previous frame was non finished
287 # we should get continuation opcode
288 elif has_partial:
289 raise WebSocketError(
290 WSCloseCode.PROTOCOL_ERROR,
291 "The opcode in non-fin frame is expected "
292 f"to be zero, got {opcode!r}",
293 )
295 assembled_payload: bytes | bytearray
296 if has_partial:
297 assembled_payload = self._partial + payload
298 self._partial.clear()
299 else:
300 assembled_payload = payload
302 # Decompress process must to be done after all packets
303 # received.
304 if compressed:
305 if not self._decompressobj:
306 self._decompressobj = ZLibDecompressor(suppress_deflate_header=True)
307 # XXX: It's possible that the zlib backend (isal is known to
308 # do this, maybe others too?) will return max_length bytes,
309 # but internally buffer more data such that the payload is
310 # >max_length, so we return one extra byte and if we're able
311 # to do that, then the message is too big.
312 try:
313 payload_merged = self._decompressobj.decompress_sync(
314 assembled_payload + WS_DEFLATE_TRAILING,
315 (
316 self._max_msg_size + 1
317 if self._max_msg_size
318 else self._max_msg_size
319 ),
320 )
321 except TooManyMembersError as exc:
322 raise WebSocketError(
323 WSCloseCode.MESSAGE_TOO_BIG,
324 "Compressed message has too many deflate members",
325 ) from exc
326 if self._max_msg_size and len(payload_merged) > self._max_msg_size:
327 raise WebSocketError(
328 WSCloseCode.MESSAGE_TOO_BIG,
329 f"Decompressed message exceeds size limit {self._max_msg_size}",
330 )
331 elif type(assembled_payload) is bytes:
332 payload_merged = assembled_payload
333 else:
334 payload_merged = bytes(assembled_payload)
336 size = len(payload_merged)
337 if opcode == OP_CODE_TEXT:
338 if self._decode_text:
339 try:
340 text = payload_merged.decode("utf-8")
341 except UnicodeDecodeError as exc:
342 raise WebSocketError(
343 WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
344 ) from exc
346 # XXX: The Text and Binary messages here can be a performance
347 # bottleneck, so we use tuple.__new__ to improve performance.
348 # This is not type safe, but many tests should fail in
349 # test_client_ws_functional.py if this is wrong.
350 msg = TUPLE_NEW(WSMessageText, (text, size, "", WS_MSG_TYPE_TEXT))
351 else:
352 # Return raw bytes for TEXT messages when decode_text=False
353 msg = TUPLE_NEW(
354 WSMessageTextBytes, (payload_merged, size, "", WS_MSG_TYPE_TEXT)
355 )
356 else:
357 msg = TUPLE_NEW(
358 WSMessageBinary, (payload_merged, size, "", WS_MSG_TYPE_BINARY)
359 )
361 self.queue.feed_data(msg)
362 elif opcode == OP_CODE_CLOSE:
363 payload_len = len(payload)
364 if payload_len >= 2:
365 close_code = UNPACK_CLOSE_CODE(payload[:2])[0]
366 # https://datatracker.ietf.org/doc/html/rfc6455#section-7.4.2
367 if close_code > 4999 or (
368 close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES
369 ):
370 raise WebSocketError(
371 WSCloseCode.PROTOCOL_ERROR,
372 f"Invalid close code: {close_code}",
373 )
374 try:
375 close_message = payload[2:].decode("utf-8")
376 except UnicodeDecodeError as exc:
377 raise WebSocketError(
378 WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
379 ) from exc
380 msg = WSMessageClose(
381 data=close_code, size=payload_len, extra=close_message
382 )
383 elif payload:
384 raise WebSocketError(
385 WSCloseCode.PROTOCOL_ERROR,
386 f"Invalid close frame: {fin} {opcode} {payload!r}",
387 )
388 else:
389 msg = WSMessageClose(data=0, size=payload_len, extra="")
391 self.queue.feed_data(msg)
392 elif opcode == OP_CODE_PING:
393 self.queue.feed_data(
394 WSMessagePing(data=bytes(payload), size=len(payload), extra="")
395 )
396 elif opcode == OP_CODE_PONG:
397 self.queue.feed_data(
398 WSMessagePong(data=bytes(payload), size=len(payload), extra="")
399 )
400 else:
401 raise WebSocketError(
402 WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}"
403 )
405 def _feed_data(self, data: bytes) -> None:
406 """Return the next frame from the socket."""
407 self.queue._stalled_reader = None
408 if self._tail:
409 data, self._tail = self._tail + data, b""
411 start_pos = 0
412 data_len = len(data)
413 data_cstr = data
415 while True:
416 if start_pos < data_len and self.queue._size > self.queue._limit:
417 # Over the high-water mark with unparsed bytes left: stash the
418 # remainder and stall. Gating on unparsed bytes keeps a read
419 # that ended on a frame boundary from arming an empty stall,
420 # which would hold the transport paused with nothing to drain.
421 self.queue._stalled_reader = self._weak_self
422 break
424 # read header
425 if self._state == READ_HEADER:
426 if data_len - start_pos < 2:
427 break
428 first_byte = data_cstr[start_pos]
429 second_byte = data_cstr[start_pos + 1]
430 start_pos += 2
432 fin = (first_byte >> 7) & 1
433 rsv1 = (first_byte >> 6) & 1
434 rsv2 = (first_byte >> 5) & 1
435 rsv3 = (first_byte >> 4) & 1
436 opcode = first_byte & 0xF
438 # frame-fin = %x0 ; more frames of this message follow
439 # / %x1 ; final frame of this message
440 # frame-rsv1 = %x0 ;
441 # 1 bit, MUST be 0 unless negotiated otherwise
442 # frame-rsv2 = %x0 ;
443 # 1 bit, MUST be 0 unless negotiated otherwise
444 # frame-rsv3 = %x0 ;
445 # 1 bit, MUST be 0 unless negotiated otherwise
446 #
447 # Remove rsv1 from this test for deflate development
448 if rsv2 or rsv3 or (rsv1 and not self._compress):
449 raise WebSocketError(
450 WSCloseCode.PROTOCOL_ERROR,
451 "Received frame with non-zero reserved bits",
452 )
454 if opcode not in {
455 OP_CODE_CONTINUATION,
456 OP_CODE_TEXT,
457 OP_CODE_BINARY,
458 OP_CODE_CLOSE,
459 OP_CODE_PING,
460 OP_CODE_PONG,
461 }:
462 raise WebSocketError(
463 WSCloseCode.PROTOCOL_ERROR,
464 f"Unexpected opcode={opcode!r}",
465 )
467 if opcode > 0x7 and fin == 0:
468 raise WebSocketError(
469 WSCloseCode.PROTOCOL_ERROR,
470 "Received fragmented control frame",
471 )
473 has_mask = (second_byte >> 7) & 1
474 length = second_byte & 0x7F
476 # Control frames MUST have a payload
477 # length of 125 bytes or less
478 if opcode > 0x7 and length > 125:
479 raise WebSocketError(
480 WSCloseCode.PROTOCOL_ERROR,
481 "Control frame payload cannot be larger than 125 bytes",
482 )
484 # Control frames (opcode > 0x7) may be interleaved between the
485 # fragments of a data message and never carry the per-message
486 # compressed bit, so they must not touch the compression state.
487 # https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
488 # https://datatracker.ietf.org/doc/html/rfc7692#section-6.1
489 if opcode > 0x7:
490 if rsv1:
491 raise WebSocketError(
492 WSCloseCode.PROTOCOL_ERROR,
493 "Received frame with non-zero reserved bits",
494 )
495 else:
496 # Set compress status if last package is FIN
497 # OR set compress status if this is first fragment
498 # Raise error if not first fragment with rsv1 = 0x1
499 if self._frame_fin or self._compressed == COMPRESSED_NOT_SET:
500 self._compressed = COMPRESSED_TRUE if rsv1 else COMPRESSED_FALSE
501 elif rsv1:
502 raise WebSocketError(
503 WSCloseCode.PROTOCOL_ERROR,
504 "Received frame with non-zero reserved bits",
505 )
506 self._frame_fin = bool(fin)
508 self._frame_opcode = opcode
509 self._has_mask = bool(has_mask)
510 self._payload_len_flag = length
511 self._state = READ_PAYLOAD_LENGTH
513 # read payload length
514 if self._state == READ_PAYLOAD_LENGTH:
515 len_flag = self._payload_len_flag
516 if len_flag == 126:
517 if data_len - start_pos < 2:
518 break
519 first_byte = data_cstr[start_pos]
520 second_byte = data_cstr[start_pos + 1]
521 start_pos += 2
522 self._payload_bytes_to_read = first_byte << 8 | second_byte
523 elif len_flag > 126:
524 if data_len - start_pos < 8:
525 break
526 # The declared length is an unsigned 64-bit integer that
527 # does not necessarily fit _payload_bytes_to_read.
528 frame_len = UNPACK_LEN3(data, start_pos)[0]
529 if frame_len > MAX_PAYLOAD_LEN:
530 raise WebSocketError(
531 WSCloseCode.MESSAGE_TOO_BIG,
532 f"Message size {int(frame_len) + len(self._partial)} "
533 f"exceeds limit {self._max_msg_size or MAX_PAYLOAD_LEN}",
534 )
535 self._payload_bytes_to_read = frame_len
536 start_pos += 8
537 else:
538 self._payload_bytes_to_read = len_flag
540 # Reject oversized data frames before buffering any payload
541 # bytes. Control frames are capped at 125 bytes (checked in
542 # READ_HEADER) so only text/binary/continuation need this.
543 if self._max_msg_size and self._frame_opcode in {
544 OP_CODE_TEXT,
545 OP_CODE_BINARY,
546 OP_CODE_CONTINUATION,
547 }:
548 # partial_len declared in reader_c.pxd to keep it in C.
549 partial_len = len(self._partial)
550 # payload_bytes_to_read is a signed Py_ssize_t C value,
551 # use subtraction here to avoid an integer overflow.
552 if self._payload_bytes_to_read >= self._max_msg_size - partial_len:
553 raise WebSocketError(
554 WSCloseCode.MESSAGE_TOO_BIG,
555 f"Message size {int(self._payload_bytes_to_read) + partial_len} "
556 f"exceeds limit {self._max_msg_size}",
557 )
559 self._state = READ_PAYLOAD_MASK if self._has_mask else READ_PAYLOAD
561 # read payload mask
562 if self._state == READ_PAYLOAD_MASK:
563 if data_len - start_pos < 4:
564 break
565 self._frame_mask = data_cstr[start_pos : start_pos + 4]
566 start_pos += 4
567 self._state = READ_PAYLOAD
569 if self._state == READ_PAYLOAD:
570 chunk_len = data_len - start_pos
571 if self._payload_bytes_to_read >= chunk_len:
572 f_end_pos = data_len
573 self._payload_bytes_to_read -= chunk_len
574 else:
575 f_end_pos = start_pos + self._payload_bytes_to_read
576 self._payload_bytes_to_read = 0
578 had_fragments = self._frame_payload_len
579 self._frame_payload_len += f_end_pos - start_pos
580 f_start_pos = start_pos
581 start_pos = f_end_pos
583 if self._payload_bytes_to_read != 0:
584 if f_start_pos < f_end_pos: # skip a header-only read
585 self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
586 if (
587 self._max_fragments
588 and len(self._payload_fragments) > self._max_fragments
589 ):
590 # Fold to bound the object count. Not a pause: nothing
591 # resumes reading until the frame is queued.
592 self._payload_buffer += b"".join(self._payload_fragments)
593 self._payload_fragments.clear()
594 break
596 payload: bytes | bytearray
597 if had_fragments:
598 self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
599 if self._payload_buffer: # folded prefix
600 self._payload_buffer += b"".join(self._payload_fragments)
601 if self._has_mask:
602 assert self._frame_mask is not None
603 websocket_mask(self._frame_mask, self._payload_buffer)
604 payload = self._payload_buffer
605 self._payload_buffer = bytearray() # detach; payload aliases it
606 elif self._has_mask:
607 assert self._frame_mask is not None
608 payload_bytearray = bytearray(b"".join(self._payload_fragments))
609 websocket_mask(self._frame_mask, payload_bytearray)
610 payload = payload_bytearray
611 else:
612 payload = b"".join(self._payload_fragments)
613 self._payload_fragments.clear()
614 elif self._has_mask:
615 assert self._frame_mask is not None
616 payload_bytearray = data_cstr[f_start_pos:f_end_pos] # type: ignore[assignment]
617 if type(payload_bytearray) is not bytearray:
618 # Cython will do the conversion for us
619 # but we need to do it for Python and we
620 # will always get here in Python
621 payload_bytearray = bytearray(payload_bytearray)
622 websocket_mask(self._frame_mask, payload_bytearray)
623 payload = payload_bytearray
624 else:
625 payload = data_cstr[f_start_pos:f_end_pos]
627 self._handle_frame(
628 self._frame_fin, self._frame_opcode, payload, self._compressed
629 )
630 self._frame_payload_len = 0
631 self._state = READ_HEADER
633 # XXX: Cython needs slices to be bounded, so we can't omit the slice end here.
634 self._tail = data_cstr[start_pos:data_len] if start_pos < data_len else b""