Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/_websocket/reader_py.py: 17%
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
5from collections import deque
6from typing import Final
8from ..base_protocol import BaseProtocol
9from ..compression_utils import ZLibDecompressor
10from ..helpers import _EXC_SENTINEL, set_exception
11from ..streams import EofStream
12from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask
13from .models import (
14 WS_DEFLATE_TRAILING,
15 WebSocketError,
16 WSCloseCode,
17 WSMessage,
18 WSMessageBinary,
19 WSMessageClose,
20 WSMessagePing,
21 WSMessagePong,
22 WSMessageText,
23 WSMessageTextBytes,
24 WSMsgType,
25)
27ALLOWED_CLOSE_CODES: Final[set[int]] = {int(i) for i in WSCloseCode}
29# States for the reader, used to parse the WebSocket frame
30# integer values are used so they can be cythonized
31READ_HEADER = 1
32READ_PAYLOAD_LENGTH = 2
33READ_PAYLOAD_MASK = 3
34READ_PAYLOAD = 4
36WS_MSG_TYPE_BINARY = WSMsgType.BINARY
37WS_MSG_TYPE_TEXT = WSMsgType.TEXT
39# WSMsgType values unpacked so they can by cythonized to ints
40OP_CODE_NOT_SET = -1
41OP_CODE_CONTINUATION = WSMsgType.CONTINUATION.value
42OP_CODE_TEXT = WSMsgType.TEXT.value
43OP_CODE_BINARY = WSMsgType.BINARY.value
44OP_CODE_CLOSE = WSMsgType.CLOSE.value
45OP_CODE_PING = WSMsgType.PING.value
46OP_CODE_PONG = WSMsgType.PONG.value
48EMPTY_FRAME_ERROR = (True, b"")
49EMPTY_FRAME = (False, b"")
51COMPRESSED_NOT_SET = -1
52COMPRESSED_FALSE = 0
53COMPRESSED_TRUE = 1
55TUPLE_NEW = tuple.__new__
57cython_int = int # Typed to int in Python, but cython with use a signed int in the pxd
60class WebSocketDataQueue:
61 """WebSocketDataQueue resumes and pauses an underlying stream.
63 It is a destination for WebSocket data.
64 """
66 def __init__(
67 self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop
68 ) -> None:
69 self._size = 0
70 self._protocol = protocol
71 self._limit = limit * 2
72 self._loop = loop
73 self._eof = False
74 self._waiter: asyncio.Future[None] | None = None
75 self._exception: type[BaseException] | BaseException | None = None
76 self._buffer: deque[WSMessage] = deque()
77 self._get_buffer = self._buffer.popleft
78 self._put_buffer = self._buffer.append
80 def is_eof(self) -> bool:
81 return self._eof
83 def exception(self) -> type[BaseException] | BaseException | None:
84 return self._exception
86 def set_exception(
87 self,
88 exc: type[BaseException] | BaseException,
89 exc_cause: builtins.BaseException = _EXC_SENTINEL,
90 ) -> None:
91 self._eof = True
92 self._exception = exc
93 if (waiter := self._waiter) is not None:
94 self._waiter = None
95 set_exception(waiter, exc, exc_cause)
97 def _release_waiter(self) -> None:
98 if (waiter := self._waiter) is None:
99 return
100 self._waiter = None
101 if not waiter.done():
102 waiter.set_result(None)
104 def feed_eof(self) -> None:
105 self._eof = True
106 self._release_waiter()
107 self._exception = None # Break cyclic references
109 def feed_data(self, data: "WSMessage") -> None:
110 size = data.size
111 self._size += size
112 self._put_buffer(data)
113 self._release_waiter()
114 if self._size > self._limit and not self._protocol._reading_paused:
115 self._protocol.pause_reading()
117 async def read(self) -> WSMessage:
118 if not self._buffer and not self._eof:
119 assert not self._waiter
120 self._waiter = self._loop.create_future()
121 try:
122 await self._waiter
123 except (asyncio.CancelledError, asyncio.TimeoutError):
124 self._waiter = None
125 raise
126 return self._read_from_buffer()
128 def _read_from_buffer(self) -> WSMessage:
129 if self._buffer:
130 data = self._get_buffer()
131 size = data.size
132 self._size -= size
133 if self._size < self._limit and self._protocol._reading_paused:
134 self._protocol.resume_reading()
135 return data
136 if self._exception is not None:
137 raise self._exception
138 raise EofStream
141class WebSocketReader:
142 def __init__(
143 self,
144 queue: WebSocketDataQueue,
145 max_msg_size: int,
146 compress: bool,
147 decode_text: bool,
148 ) -> None:
149 self.queue = queue
150 self._max_msg_size = max_msg_size
151 self._decode_text = decode_text
153 self._exc: Exception | None = None
154 self._partial = bytearray()
155 self._state = READ_HEADER
157 self._opcode: int = OP_CODE_NOT_SET
158 self._frame_fin = False
159 self._frame_opcode: int = OP_CODE_NOT_SET
160 self._payload_fragments: list[bytes] = []
161 self._frame_payload_len = 0
163 self._tail: bytes = b""
164 self._has_mask = False
165 self._frame_mask: bytes | None = None
166 self._payload_bytes_to_read = 0
167 self._payload_len_flag = 0
168 self._compressed: int = COMPRESSED_NOT_SET
169 self._decompressobj: ZLibDecompressor | None = None
170 self._compress = compress
172 def feed_eof(self) -> None:
173 self.queue.feed_eof()
175 # data can be bytearray on Windows because proactor event loop uses bytearray
176 # and asyncio types this to Union[bytes, bytearray, memoryview] so we need
177 # coerce data to bytes if it is not
178 def feed_data(self, data: bytes | bytearray | memoryview) -> tuple[bool, bytes]:
179 if type(data) is not bytes:
180 data = bytes(data)
182 if self._exc is not None:
183 return True, data
185 try:
186 self._feed_data(data)
187 except Exception as exc:
188 self._exc = exc
189 set_exception(self.queue, exc)
190 return EMPTY_FRAME_ERROR
192 return EMPTY_FRAME
194 def _handle_frame(
195 self,
196 fin: bool,
197 opcode: int | cython_int, # Union intended: Cython pxd uses C int
198 payload: bytes | bytearray,
199 compressed: int | cython_int, # Union intended: Cython pxd uses C int
200 ) -> None:
201 msg: WSMessage
202 if opcode in {OP_CODE_TEXT, OP_CODE_BINARY, OP_CODE_CONTINUATION}:
203 # Validate continuation frames before processing
204 if opcode == OP_CODE_CONTINUATION and self._opcode == OP_CODE_NOT_SET:
205 raise WebSocketError(
206 WSCloseCode.PROTOCOL_ERROR,
207 "Continuation frame for non started message",
208 )
210 # load text/binary
211 if not fin:
212 # got partial frame payload
213 if opcode != OP_CODE_CONTINUATION:
214 self._opcode = opcode
215 self._partial += payload
216 return
218 has_partial = bool(self._partial)
219 if opcode == OP_CODE_CONTINUATION:
220 opcode = self._opcode
221 self._opcode = OP_CODE_NOT_SET
222 # previous frame was non finished
223 # we should get continuation opcode
224 elif has_partial:
225 raise WebSocketError(
226 WSCloseCode.PROTOCOL_ERROR,
227 "The opcode in non-fin frame is expected "
228 f"to be zero, got {opcode!r}",
229 )
231 assembled_payload: bytes | bytearray
232 if has_partial:
233 assembled_payload = self._partial + payload
234 self._partial.clear()
235 else:
236 assembled_payload = payload
238 # Decompress process must to be done after all packets
239 # received.
240 if compressed:
241 if not self._decompressobj:
242 self._decompressobj = ZLibDecompressor(suppress_deflate_header=True)
243 # XXX: It's possible that the zlib backend (isal is known to
244 # do this, maybe others too?) will return max_length bytes,
245 # but internally buffer more data such that the payload is
246 # >max_length, so we return one extra byte and if we're able
247 # to do that, then the message is too big.
248 payload_merged = self._decompressobj.decompress_sync(
249 assembled_payload + WS_DEFLATE_TRAILING,
250 (
251 self._max_msg_size + 1
252 if self._max_msg_size
253 else self._max_msg_size
254 ),
255 )
256 if self._max_msg_size and len(payload_merged) > self._max_msg_size:
257 raise WebSocketError(
258 WSCloseCode.MESSAGE_TOO_BIG,
259 f"Decompressed message exceeds size limit {self._max_msg_size}",
260 )
261 elif type(assembled_payload) is bytes:
262 payload_merged = assembled_payload
263 else:
264 payload_merged = bytes(assembled_payload)
266 size = len(payload_merged)
267 if opcode == OP_CODE_TEXT:
268 if self._decode_text:
269 try:
270 text = payload_merged.decode("utf-8")
271 except UnicodeDecodeError as exc:
272 raise WebSocketError(
273 WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
274 ) from exc
276 # XXX: The Text and Binary messages here can be a performance
277 # bottleneck, so we use tuple.__new__ to improve performance.
278 # This is not type safe, but many tests should fail in
279 # test_client_ws_functional.py if this is wrong.
280 msg = TUPLE_NEW(WSMessageText, (text, size, "", WS_MSG_TYPE_TEXT))
281 else:
282 # Return raw bytes for TEXT messages when decode_text=False
283 msg = TUPLE_NEW(
284 WSMessageTextBytes, (payload_merged, size, "", WS_MSG_TYPE_TEXT)
285 )
286 else:
287 msg = TUPLE_NEW(
288 WSMessageBinary, (payload_merged, size, "", WS_MSG_TYPE_BINARY)
289 )
291 self.queue.feed_data(msg)
292 elif opcode == OP_CODE_CLOSE:
293 payload_len = len(payload)
294 if payload_len >= 2:
295 close_code = UNPACK_CLOSE_CODE(payload[:2])[0]
296 # https://datatracker.ietf.org/doc/html/rfc6455#section-7.4.2
297 if close_code > 4999 or (
298 close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES
299 ):
300 raise WebSocketError(
301 WSCloseCode.PROTOCOL_ERROR,
302 f"Invalid close code: {close_code}",
303 )
304 try:
305 close_message = payload[2:].decode("utf-8")
306 except UnicodeDecodeError as exc:
307 raise WebSocketError(
308 WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
309 ) from exc
310 msg = WSMessageClose(
311 data=close_code, size=payload_len, extra=close_message
312 )
313 elif payload:
314 raise WebSocketError(
315 WSCloseCode.PROTOCOL_ERROR,
316 f"Invalid close frame: {fin} {opcode} {payload!r}",
317 )
318 else:
319 msg = WSMessageClose(data=0, size=payload_len, extra="")
321 self.queue.feed_data(msg)
322 elif opcode == OP_CODE_PING:
323 self.queue.feed_data(
324 WSMessagePing(data=bytes(payload), size=len(payload), extra="")
325 )
326 elif opcode == OP_CODE_PONG:
327 self.queue.feed_data(
328 WSMessagePong(data=bytes(payload), size=len(payload), extra="")
329 )
330 else:
331 raise WebSocketError(
332 WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}"
333 )
335 def _feed_data(self, data: bytes) -> None:
336 """Return the next frame from the socket."""
337 if self._tail:
338 data, self._tail = self._tail + data, b""
340 start_pos: int = 0
341 data_len = len(data)
342 data_cstr = data
344 while True:
345 # read header
346 if self._state == READ_HEADER:
347 if data_len - start_pos < 2:
348 break
349 first_byte = data_cstr[start_pos]
350 second_byte = data_cstr[start_pos + 1]
351 start_pos += 2
353 fin = (first_byte >> 7) & 1
354 rsv1 = (first_byte >> 6) & 1
355 rsv2 = (first_byte >> 5) & 1
356 rsv3 = (first_byte >> 4) & 1
357 opcode = first_byte & 0xF
359 # frame-fin = %x0 ; more frames of this message follow
360 # / %x1 ; final frame of this message
361 # frame-rsv1 = %x0 ;
362 # 1 bit, MUST be 0 unless negotiated otherwise
363 # frame-rsv2 = %x0 ;
364 # 1 bit, MUST be 0 unless negotiated otherwise
365 # frame-rsv3 = %x0 ;
366 # 1 bit, MUST be 0 unless negotiated otherwise
367 #
368 # Remove rsv1 from this test for deflate development
369 if rsv2 or rsv3 or (rsv1 and not self._compress):
370 raise WebSocketError(
371 WSCloseCode.PROTOCOL_ERROR,
372 "Received frame with non-zero reserved bits",
373 )
375 if opcode not in {
376 OP_CODE_CONTINUATION,
377 OP_CODE_TEXT,
378 OP_CODE_BINARY,
379 OP_CODE_CLOSE,
380 OP_CODE_PING,
381 OP_CODE_PONG,
382 }:
383 raise WebSocketError(
384 WSCloseCode.PROTOCOL_ERROR,
385 f"Unexpected opcode={opcode!r}",
386 )
388 if opcode > 0x7 and fin == 0:
389 raise WebSocketError(
390 WSCloseCode.PROTOCOL_ERROR,
391 "Received fragmented control frame",
392 )
394 has_mask = (second_byte >> 7) & 1
395 length = second_byte & 0x7F
397 # Control frames MUST have a payload
398 # length of 125 bytes or less
399 if opcode > 0x7 and length > 125:
400 raise WebSocketError(
401 WSCloseCode.PROTOCOL_ERROR,
402 "Control frame payload cannot be larger than 125 bytes",
403 )
405 # Set compress status if last package is FIN
406 # OR set compress status if this is first fragment
407 # Raise error if not first fragment with rsv1 = 0x1
408 if self._frame_fin or self._compressed == COMPRESSED_NOT_SET:
409 self._compressed = COMPRESSED_TRUE if rsv1 else COMPRESSED_FALSE
410 elif rsv1:
411 raise WebSocketError(
412 WSCloseCode.PROTOCOL_ERROR,
413 "Received frame with non-zero reserved bits",
414 )
416 # Control frames (opcode > 0x7) may be interleaved between the
417 # fragments of a data message.
418 # https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
419 if opcode <= 0x7:
420 self._frame_fin = bool(fin)
421 self._frame_opcode = opcode
422 self._has_mask = bool(has_mask)
423 self._payload_len_flag = length
424 self._state = READ_PAYLOAD_LENGTH
426 # read payload length
427 if self._state == READ_PAYLOAD_LENGTH:
428 len_flag = self._payload_len_flag
429 if len_flag == 126:
430 if data_len - start_pos < 2:
431 break
432 first_byte = data_cstr[start_pos]
433 second_byte = data_cstr[start_pos + 1]
434 start_pos += 2
435 self._payload_bytes_to_read = first_byte << 8 | second_byte
436 elif len_flag > 126:
437 if data_len - start_pos < 8:
438 break
439 self._payload_bytes_to_read = UNPACK_LEN3(data, start_pos)[0]
440 start_pos += 8
441 else:
442 self._payload_bytes_to_read = len_flag
444 # Reject oversized data frames before buffering any payload
445 # bytes. Control frames are capped at 125 bytes (checked in
446 # READ_HEADER) so only text/binary/continuation need this.
447 if self._max_msg_size and self._frame_opcode in {
448 OP_CODE_TEXT,
449 OP_CODE_BINARY,
450 OP_CODE_CONTINUATION,
451 }:
452 projected_size = self._payload_bytes_to_read + len(self._partial)
453 if projected_size >= self._max_msg_size:
454 raise WebSocketError(
455 WSCloseCode.MESSAGE_TOO_BIG,
456 f"Message size {projected_size} "
457 f"exceeds limit {self._max_msg_size}",
458 )
460 self._state = READ_PAYLOAD_MASK if self._has_mask else READ_PAYLOAD
462 # read payload mask
463 if self._state == READ_PAYLOAD_MASK:
464 if data_len - start_pos < 4:
465 break
466 self._frame_mask = data_cstr[start_pos : start_pos + 4]
467 start_pos += 4
468 self._state = READ_PAYLOAD
470 if self._state == READ_PAYLOAD:
471 chunk_len = data_len - start_pos
472 if self._payload_bytes_to_read >= chunk_len:
473 f_end_pos = data_len
474 self._payload_bytes_to_read -= chunk_len
475 else:
476 f_end_pos = start_pos + self._payload_bytes_to_read
477 self._payload_bytes_to_read = 0
479 had_fragments = self._frame_payload_len
480 self._frame_payload_len += f_end_pos - start_pos
481 f_start_pos = start_pos
482 start_pos = f_end_pos
484 if self._payload_bytes_to_read != 0:
485 # If we don't have a complete frame, we need to save the
486 # data for the next call to feed_data.
487 self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
488 break
490 payload: bytes | bytearray
491 if had_fragments:
492 # We have to join the payload fragments get the payload
493 self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
494 if self._has_mask:
495 assert self._frame_mask is not None
496 payload_bytearray = bytearray(b"".join(self._payload_fragments))
497 websocket_mask(self._frame_mask, payload_bytearray)
498 payload = payload_bytearray
499 else:
500 payload = b"".join(self._payload_fragments)
501 self._payload_fragments.clear()
502 elif self._has_mask:
503 assert self._frame_mask is not None
504 payload_bytearray = data_cstr[f_start_pos:f_end_pos] # type: ignore[assignment]
505 if type(payload_bytearray) is not bytearray: # pragma: no branch
506 # Cython will do the conversion for us
507 # but we need to do it for Python and we
508 # will always get here in Python
509 payload_bytearray = bytearray(payload_bytearray)
510 websocket_mask(self._frame_mask, payload_bytearray)
511 payload = payload_bytearray
512 else:
513 payload = data_cstr[f_start_pos:f_end_pos]
515 self._handle_frame(
516 self._frame_fin, self._frame_opcode, payload, self._compressed
517 )
518 self._frame_payload_len = 0
519 self._state = READ_HEADER
521 # XXX: Cython needs slices to be bounded, so we can't omit the slice end here.
522 self._tail = data_cstr[start_pos:data_len] if start_pos < data_len else b""