Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/http_websocket.py: 27%
377 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-27 06:09 +0000
1"""WebSocket protocol versions 13 and 8."""
3import asyncio
4import functools
5import json
6import random
7import re
8import sys
9import zlib
10from enum import IntEnum
11from struct import Struct
12from typing import (
13 Any,
14 Callable,
15 Final,
16 List,
17 NamedTuple,
18 Optional,
19 Pattern,
20 Set,
21 Tuple,
22 Union,
23 cast,
24)
26from .base_protocol import BaseProtocol
27from .compression_utils import ZLibCompressor, ZLibDecompressor
28from .helpers import NO_EXTENSIONS
29from .streams import DataQueue
31__all__ = (
32 "WS_CLOSED_MESSAGE",
33 "WS_CLOSING_MESSAGE",
34 "WS_KEY",
35 "WebSocketReader",
36 "WebSocketWriter",
37 "WSMessage",
38 "WebSocketError",
39 "WSMsgType",
40 "WSCloseCode",
41)
44class WSCloseCode(IntEnum):
45 OK = 1000
46 GOING_AWAY = 1001
47 PROTOCOL_ERROR = 1002
48 UNSUPPORTED_DATA = 1003
49 ABNORMAL_CLOSURE = 1006
50 INVALID_TEXT = 1007
51 POLICY_VIOLATION = 1008
52 MESSAGE_TOO_BIG = 1009
53 MANDATORY_EXTENSION = 1010
54 INTERNAL_ERROR = 1011
55 SERVICE_RESTART = 1012
56 TRY_AGAIN_LATER = 1013
57 BAD_GATEWAY = 1014
60ALLOWED_CLOSE_CODES: Final[Set[int]] = {int(i) for i in WSCloseCode}
63class WSMsgType(IntEnum):
64 # websocket spec types
65 CONTINUATION = 0x0
66 TEXT = 0x1
67 BINARY = 0x2
68 PING = 0x9
69 PONG = 0xA
70 CLOSE = 0x8
72 # aiohttp specific types
73 CLOSING = 0x100
74 CLOSED = 0x101
75 ERROR = 0x102
78WS_KEY: Final[bytes] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
81UNPACK_LEN2 = Struct("!H").unpack_from
82UNPACK_LEN3 = Struct("!Q").unpack_from
83UNPACK_CLOSE_CODE = Struct("!H").unpack
84PACK_LEN1 = Struct("!BB").pack
85PACK_LEN2 = Struct("!BBH").pack
86PACK_LEN3 = Struct("!BBQ").pack
87PACK_CLOSE_CODE = Struct("!H").pack
88MSG_SIZE: Final[int] = 2**14
89DEFAULT_LIMIT: Final[int] = 2**16
92class WSMessage(NamedTuple):
93 type: WSMsgType
94 # To type correctly, this would need some kind of tagged union for each type.
95 data: Any
96 extra: Optional[str]
98 def json(self, *, loads: Callable[[Any], Any] = json.loads) -> Any:
99 """Return parsed JSON data.
101 .. versionadded:: 0.22
102 """
103 return loads(self.data)
106WS_CLOSED_MESSAGE = WSMessage(WSMsgType.CLOSED, None, None)
107WS_CLOSING_MESSAGE = WSMessage(WSMsgType.CLOSING, None, None)
110class WebSocketError(Exception):
111 """WebSocket protocol parser error."""
113 def __init__(self, code: int, message: str) -> None:
114 self.code = code
115 super().__init__(code, message)
117 def __str__(self) -> str:
118 return cast(str, self.args[1])
121class WSHandshakeError(Exception):
122 """WebSocket protocol handshake error."""
125native_byteorder: Final[str] = sys.byteorder
128# Used by _websocket_mask_python
129@functools.lru_cache()
130def _xor_table() -> List[bytes]:
131 return [bytes(a ^ b for a in range(256)) for b in range(256)]
134def _websocket_mask_python(mask: bytes, data: bytearray) -> None:
135 """Websocket masking function.
137 `mask` is a `bytes` object of length 4; `data` is a `bytearray`
138 object of any length. The contents of `data` are masked with `mask`,
139 as specified in section 5.3 of RFC 6455.
141 Note that this function mutates the `data` argument.
143 This pure-python implementation may be replaced by an optimized
144 version when available.
146 """
147 assert isinstance(data, bytearray), data
148 assert len(mask) == 4, mask
150 if data:
151 _XOR_TABLE = _xor_table()
152 a, b, c, d = (_XOR_TABLE[n] for n in mask)
153 data[::4] = data[::4].translate(a)
154 data[1::4] = data[1::4].translate(b)
155 data[2::4] = data[2::4].translate(c)
156 data[3::4] = data[3::4].translate(d)
159if NO_EXTENSIONS: # pragma: no cover
160 _websocket_mask = _websocket_mask_python
161else:
162 try:
163 from ._websocket import _websocket_mask_cython # type: ignore[import]
165 _websocket_mask = _websocket_mask_cython
166 except ImportError: # pragma: no cover
167 _websocket_mask = _websocket_mask_python
169_WS_DEFLATE_TRAILING: Final[bytes] = bytes([0x00, 0x00, 0xFF, 0xFF])
172_WS_EXT_RE: Final[Pattern[str]] = re.compile(
173 r"^(?:;\s*(?:"
174 r"(server_no_context_takeover)|"
175 r"(client_no_context_takeover)|"
176 r"(server_max_window_bits(?:=(\d+))?)|"
177 r"(client_max_window_bits(?:=(\d+))?)))*$"
178)
180_WS_EXT_RE_SPLIT: Final[Pattern[str]] = re.compile(r"permessage-deflate([^,]+)?")
183def ws_ext_parse(extstr: Optional[str], isserver: bool = False) -> Tuple[int, bool]:
184 if not extstr:
185 return 0, False
187 compress = 0
188 notakeover = False
189 for ext in _WS_EXT_RE_SPLIT.finditer(extstr):
190 defext = ext.group(1)
191 # Return compress = 15 when get `permessage-deflate`
192 if not defext:
193 compress = 15
194 break
195 match = _WS_EXT_RE.match(defext)
196 if match:
197 compress = 15
198 if isserver:
199 # Server never fail to detect compress handshake.
200 # Server does not need to send max wbit to client
201 if match.group(4):
202 compress = int(match.group(4))
203 # Group3 must match if group4 matches
204 # Compress wbit 8 does not support in zlib
205 # If compress level not support,
206 # CONTINUE to next extension
207 if compress > 15 or compress < 9:
208 compress = 0
209 continue
210 if match.group(1):
211 notakeover = True
212 # Ignore regex group 5 & 6 for client_max_window_bits
213 break
214 else:
215 if match.group(6):
216 compress = int(match.group(6))
217 # Group5 must match if group6 matches
218 # Compress wbit 8 does not support in zlib
219 # If compress level not support,
220 # FAIL the parse progress
221 if compress > 15 or compress < 9:
222 raise WSHandshakeError("Invalid window size")
223 if match.group(2):
224 notakeover = True
225 # Ignore regex group 5 & 6 for client_max_window_bits
226 break
227 # Return Fail if client side and not match
228 elif not isserver:
229 raise WSHandshakeError("Extension for deflate not supported" + ext.group(1))
231 return compress, notakeover
234def ws_ext_gen(
235 compress: int = 15, isserver: bool = False, server_notakeover: bool = False
236) -> str:
237 # client_notakeover=False not used for server
238 # compress wbit 8 does not support in zlib
239 if compress < 9 or compress > 15:
240 raise ValueError(
241 "Compress wbits must between 9 and 15, " "zlib does not support wbits=8"
242 )
243 enabledext = ["permessage-deflate"]
244 if not isserver:
245 enabledext.append("client_max_window_bits")
247 if compress < 15:
248 enabledext.append("server_max_window_bits=" + str(compress))
249 if server_notakeover:
250 enabledext.append("server_no_context_takeover")
251 # if client_notakeover:
252 # enabledext.append('client_no_context_takeover')
253 return "; ".join(enabledext)
256class WSParserState(IntEnum):
257 READ_HEADER = 1
258 READ_PAYLOAD_LENGTH = 2
259 READ_PAYLOAD_MASK = 3
260 READ_PAYLOAD = 4
263class WebSocketReader:
264 def __init__(
265 self, queue: DataQueue[WSMessage], max_msg_size: int, compress: bool = True
266 ) -> None:
267 self.queue = queue
268 self._max_msg_size = max_msg_size
270 self._exc: Optional[BaseException] = None
271 self._partial = bytearray()
272 self._state = WSParserState.READ_HEADER
274 self._opcode: Optional[int] = None
275 self._frame_fin = False
276 self._frame_opcode: Optional[int] = None
277 self._frame_payload = bytearray()
279 self._tail = b""
280 self._has_mask = False
281 self._frame_mask: Optional[bytes] = None
282 self._payload_length = 0
283 self._payload_length_flag = 0
284 self._compressed: Optional[bool] = None
285 self._decompressobj: Optional[ZLibDecompressor] = None
286 self._compress = compress
288 def feed_eof(self) -> None:
289 self.queue.feed_eof()
291 def feed_data(self, data: bytes) -> Tuple[bool, bytes]:
292 if self._exc:
293 return True, data
295 try:
296 return self._feed_data(data)
297 except Exception as exc:
298 self._exc = exc
299 self.queue.set_exception(exc)
300 return True, b""
302 def _feed_data(self, data: bytes) -> Tuple[bool, bytes]:
303 for fin, opcode, payload, compressed in self.parse_frame(data):
304 if compressed and not self._decompressobj:
305 self._decompressobj = ZLibDecompressor(suppress_deflate_header=True)
306 if opcode == WSMsgType.CLOSE:
307 if len(payload) >= 2:
308 close_code = UNPACK_CLOSE_CODE(payload[:2])[0]
309 if close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES:
310 raise WebSocketError(
311 WSCloseCode.PROTOCOL_ERROR,
312 f"Invalid close code: {close_code}",
313 )
314 try:
315 close_message = payload[2:].decode("utf-8")
316 except UnicodeDecodeError as exc:
317 raise WebSocketError(
318 WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
319 ) from exc
320 msg = WSMessage(WSMsgType.CLOSE, close_code, close_message)
321 elif payload:
322 raise WebSocketError(
323 WSCloseCode.PROTOCOL_ERROR,
324 f"Invalid close frame: {fin} {opcode} {payload!r}",
325 )
326 else:
327 msg = WSMessage(WSMsgType.CLOSE, 0, "")
329 self.queue.feed_data(msg, 0)
331 elif opcode == WSMsgType.PING:
332 self.queue.feed_data(
333 WSMessage(WSMsgType.PING, payload, ""), len(payload)
334 )
336 elif opcode == WSMsgType.PONG:
337 self.queue.feed_data(
338 WSMessage(WSMsgType.PONG, payload, ""), len(payload)
339 )
341 elif (
342 opcode not in (WSMsgType.TEXT, WSMsgType.BINARY)
343 and self._opcode is None
344 ):
345 raise WebSocketError(
346 WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}"
347 )
348 else:
349 # load text/binary
350 if not fin:
351 # got partial frame payload
352 if opcode != WSMsgType.CONTINUATION:
353 self._opcode = opcode
354 self._partial.extend(payload)
355 if self._max_msg_size and len(self._partial) >= self._max_msg_size:
356 raise WebSocketError(
357 WSCloseCode.MESSAGE_TOO_BIG,
358 "Message size {} exceeds limit {}".format(
359 len(self._partial), self._max_msg_size
360 ),
361 )
362 else:
363 # previous frame was non finished
364 # we should get continuation opcode
365 if self._partial:
366 if opcode != WSMsgType.CONTINUATION:
367 raise WebSocketError(
368 WSCloseCode.PROTOCOL_ERROR,
369 "The opcode in non-fin frame is expected "
370 "to be zero, got {!r}".format(opcode),
371 )
373 if opcode == WSMsgType.CONTINUATION:
374 assert self._opcode is not None
375 opcode = self._opcode
376 self._opcode = None
378 self._partial.extend(payload)
379 if self._max_msg_size and len(self._partial) >= self._max_msg_size:
380 raise WebSocketError(
381 WSCloseCode.MESSAGE_TOO_BIG,
382 "Message size {} exceeds limit {}".format(
383 len(self._partial), self._max_msg_size
384 ),
385 )
387 # Decompress process must to be done after all packets
388 # received.
389 if compressed:
390 assert self._decompressobj is not None
391 self._partial.extend(_WS_DEFLATE_TRAILING)
392 payload_merged = self._decompressobj.decompress_sync(
393 self._partial, self._max_msg_size
394 )
395 if self._decompressobj.unconsumed_tail:
396 left = len(self._decompressobj.unconsumed_tail)
397 raise WebSocketError(
398 WSCloseCode.MESSAGE_TOO_BIG,
399 "Decompressed message size {} exceeds limit {}".format(
400 self._max_msg_size + left, self._max_msg_size
401 ),
402 )
403 else:
404 payload_merged = bytes(self._partial)
406 self._partial.clear()
408 if opcode == WSMsgType.TEXT:
409 try:
410 text = payload_merged.decode("utf-8")
411 self.queue.feed_data(
412 WSMessage(WSMsgType.TEXT, text, ""), len(text)
413 )
414 except UnicodeDecodeError as exc:
415 raise WebSocketError(
416 WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
417 ) from exc
418 else:
419 self.queue.feed_data(
420 WSMessage(WSMsgType.BINARY, payload_merged, ""),
421 len(payload_merged),
422 )
424 return False, b""
426 def parse_frame(
427 self, buf: bytes
428 ) -> List[Tuple[bool, Optional[int], bytearray, Optional[bool]]]:
429 """Return the next frame from the socket."""
430 frames = []
431 if self._tail:
432 buf, self._tail = self._tail + buf, b""
434 start_pos = 0
435 buf_length = len(buf)
437 while True:
438 # read header
439 if self._state == WSParserState.READ_HEADER:
440 if buf_length - start_pos >= 2:
441 data = buf[start_pos : start_pos + 2]
442 start_pos += 2
443 first_byte, second_byte = data
445 fin = (first_byte >> 7) & 1
446 rsv1 = (first_byte >> 6) & 1
447 rsv2 = (first_byte >> 5) & 1
448 rsv3 = (first_byte >> 4) & 1
449 opcode = first_byte & 0xF
451 # frame-fin = %x0 ; more frames of this message follow
452 # / %x1 ; final frame of this message
453 # frame-rsv1 = %x0 ;
454 # 1 bit, MUST be 0 unless negotiated otherwise
455 # frame-rsv2 = %x0 ;
456 # 1 bit, MUST be 0 unless negotiated otherwise
457 # frame-rsv3 = %x0 ;
458 # 1 bit, MUST be 0 unless negotiated otherwise
459 #
460 # Remove rsv1 from this test for deflate development
461 if rsv2 or rsv3 or (rsv1 and not self._compress):
462 raise WebSocketError(
463 WSCloseCode.PROTOCOL_ERROR,
464 "Received frame with non-zero reserved bits",
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 # Set compress status if last package is FIN
485 # OR set compress status if this is first fragment
486 # Raise error if not first fragment with rsv1 = 0x1
487 if self._frame_fin or self._compressed is None:
488 self._compressed = True if rsv1 else False
489 elif rsv1:
490 raise WebSocketError(
491 WSCloseCode.PROTOCOL_ERROR,
492 "Received frame with non-zero reserved bits",
493 )
495 self._frame_fin = bool(fin)
496 self._frame_opcode = opcode
497 self._has_mask = bool(has_mask)
498 self._payload_length_flag = length
499 self._state = WSParserState.READ_PAYLOAD_LENGTH
500 else:
501 break
503 # read payload length
504 if self._state == WSParserState.READ_PAYLOAD_LENGTH:
505 length = self._payload_length_flag
506 if length == 126:
507 if buf_length - start_pos >= 2:
508 data = buf[start_pos : start_pos + 2]
509 start_pos += 2
510 length = UNPACK_LEN2(data)[0]
511 self._payload_length = length
512 self._state = (
513 WSParserState.READ_PAYLOAD_MASK
514 if self._has_mask
515 else WSParserState.READ_PAYLOAD
516 )
517 else:
518 break
519 elif length > 126:
520 if buf_length - start_pos >= 8:
521 data = buf[start_pos : start_pos + 8]
522 start_pos += 8
523 length = UNPACK_LEN3(data)[0]
524 self._payload_length = length
525 self._state = (
526 WSParserState.READ_PAYLOAD_MASK
527 if self._has_mask
528 else WSParserState.READ_PAYLOAD
529 )
530 else:
531 break
532 else:
533 self._payload_length = length
534 self._state = (
535 WSParserState.READ_PAYLOAD_MASK
536 if self._has_mask
537 else WSParserState.READ_PAYLOAD
538 )
540 # read payload mask
541 if self._state == WSParserState.READ_PAYLOAD_MASK:
542 if buf_length - start_pos >= 4:
543 self._frame_mask = buf[start_pos : start_pos + 4]
544 start_pos += 4
545 self._state = WSParserState.READ_PAYLOAD
546 else:
547 break
549 if self._state == WSParserState.READ_PAYLOAD:
550 length = self._payload_length
551 payload = self._frame_payload
553 chunk_len = buf_length - start_pos
554 if length >= chunk_len:
555 self._payload_length = length - chunk_len
556 payload.extend(buf[start_pos:])
557 start_pos = buf_length
558 else:
559 self._payload_length = 0
560 payload.extend(buf[start_pos : start_pos + length])
561 start_pos = start_pos + length
563 if self._payload_length == 0:
564 if self._has_mask:
565 assert self._frame_mask is not None
566 _websocket_mask(self._frame_mask, payload)
568 frames.append(
569 (self._frame_fin, self._frame_opcode, payload, self._compressed)
570 )
572 self._frame_payload = bytearray()
573 self._state = WSParserState.READ_HEADER
574 else:
575 break
577 self._tail = buf[start_pos:]
579 return frames
582class WebSocketWriter:
583 def __init__(
584 self,
585 protocol: BaseProtocol,
586 transport: asyncio.Transport,
587 *,
588 use_mask: bool = False,
589 limit: int = DEFAULT_LIMIT,
590 random: Any = random.Random(),
591 compress: int = 0,
592 notakeover: bool = False,
593 ) -> None:
594 self.protocol = protocol
595 self.transport = transport
596 self.use_mask = use_mask
597 self.randrange = random.randrange
598 self.compress = compress
599 self.notakeover = notakeover
600 self._closing = False
601 self._limit = limit
602 self._output_size = 0
603 self._compressobj: Any = None # actually compressobj
605 async def _send_frame(
606 self, message: bytes, opcode: int, compress: Optional[int] = None
607 ) -> None:
608 """Send a frame over the websocket with message as its payload."""
609 if self._closing and not (opcode & WSMsgType.CLOSE):
610 raise ConnectionResetError("Cannot write to closing transport")
612 rsv = 0
614 # Only compress larger packets (disabled)
615 # Does small packet needs to be compressed?
616 # if self.compress and opcode < 8 and len(message) > 124:
617 if (compress or self.compress) and opcode < 8:
618 if compress:
619 # Do not set self._compress if compressing is for this frame
620 compressobj = ZLibCompressor(level=zlib.Z_BEST_SPEED, wbits=-compress)
621 else: # self.compress
622 if not self._compressobj:
623 self._compressobj = ZLibCompressor(
624 level=zlib.Z_BEST_SPEED, wbits=-self.compress
625 )
626 compressobj = self._compressobj
628 message = await compressobj.compress(message)
629 message += compressobj.flush(
630 zlib.Z_FULL_FLUSH if self.notakeover else zlib.Z_SYNC_FLUSH
631 )
632 if message.endswith(_WS_DEFLATE_TRAILING):
633 message = message[:-4]
634 rsv = rsv | 0x40
636 msg_length = len(message)
638 use_mask = self.use_mask
639 if use_mask:
640 mask_bit = 0x80
641 else:
642 mask_bit = 0
644 if msg_length < 126:
645 header = PACK_LEN1(0x80 | rsv | opcode, msg_length | mask_bit)
646 elif msg_length < (1 << 16):
647 header = PACK_LEN2(0x80 | rsv | opcode, 126 | mask_bit, msg_length)
648 else:
649 header = PACK_LEN3(0x80 | rsv | opcode, 127 | mask_bit, msg_length)
650 if use_mask:
651 mask = self.randrange(0, 0xFFFFFFFF)
652 mask = mask.to_bytes(4, "big")
653 message = bytearray(message)
654 _websocket_mask(mask, message)
655 self._write(header + mask + message)
656 self._output_size += len(header) + len(mask) + len(message)
657 else:
658 if len(message) > MSG_SIZE:
659 self._write(header)
660 self._write(message)
661 else:
662 self._write(header + message)
664 self._output_size += len(header) + len(message)
666 if self._output_size > self._limit:
667 self._output_size = 0
668 await self.protocol._drain_helper()
670 def _write(self, data: bytes) -> None:
671 if self.transport.is_closing():
672 raise ConnectionResetError("Cannot write to closing transport")
673 self.transport.write(data)
675 async def pong(self, message: Union[bytes, str] = b"") -> None:
676 """Send pong message."""
677 if isinstance(message, str):
678 message = message.encode("utf-8")
679 await self._send_frame(message, WSMsgType.PONG)
681 async def ping(self, message: Union[bytes, str] = b"") -> None:
682 """Send ping message."""
683 if isinstance(message, str):
684 message = message.encode("utf-8")
685 await self._send_frame(message, WSMsgType.PING)
687 async def send(
688 self,
689 message: Union[str, bytes],
690 binary: bool = False,
691 compress: Optional[int] = None,
692 ) -> None:
693 """Send a frame over the websocket with message as its payload."""
694 if isinstance(message, str):
695 message = message.encode("utf-8")
696 if binary:
697 await self._send_frame(message, WSMsgType.BINARY, compress)
698 else:
699 await self._send_frame(message, WSMsgType.TEXT, compress)
701 async def close(self, code: int = 1000, message: Union[bytes, str] = b"") -> None:
702 """Close the websocket, sending the specified code and message."""
703 if isinstance(message, str):
704 message = message.encode("utf-8")
705 try:
706 await self._send_frame(
707 PACK_CLOSE_CODE(code) + message, opcode=WSMsgType.CLOSE
708 )
709 finally:
710 self._closing = True