Initial commit

2025-08-14 00:25:46 +02:00 · 2024-12-23 19:27:44 +06:00
commit e81df4c87e
4952 changed files with 1705479 additions and 0 deletions
--- a/venv/Lib/site-packages/aiohttp/_websocket/writer.py
+++ b/venv/Lib/site-packages/aiohttp/_websocket/writer.py
@@ -0,0 +1,177 @@
+"""WebSocket protocol versions 13 and 8."""
+
+import asyncio
+import random
+import zlib
+from functools import partial
+from typing import Any, Final, Optional, Union
+
+from ..base_protocol import BaseProtocol
+from ..client_exceptions import ClientConnectionResetError
+from ..compression_utils import ZLibCompressor
+from .helpers import (
+    MASK_LEN,
+    MSG_SIZE,
+    PACK_CLOSE_CODE,
+    PACK_LEN1,
+    PACK_LEN2,
+    PACK_LEN3,
+    PACK_RANDBITS,
+    websocket_mask,
+)
+from .models import WS_DEFLATE_TRAILING, WSMsgType
+
+DEFAULT_LIMIT: Final[int] = 2**16
+
+# For websockets, keeping latency low is extremely important as implementations
+# generally expect to be able to send and receive messages quickly.  We use a
+# larger chunk size than the default to reduce the number of executor calls
+# since the executor is a significant source of latency and overhead when
+# the chunks are small. A size of 5KiB was chosen because it is also the
+# same value python-zlib-ng choose to use as the threshold to release the GIL.
+
+WEBSOCKET_MAX_SYNC_CHUNK_SIZE = 5 * 1024
+
+
+class WebSocketWriter:
+    """WebSocket writer.
+
+    The writer is responsible for sending messages to the client. It is
+    created by the protocol when a connection is established. The writer
+    should avoid implementing any application logic and should only be
+    concerned with the low-level details of the WebSocket protocol.
+    """
+
+    def __init__(
+        self,
+        protocol: BaseProtocol,
+        transport: asyncio.Transport,
+        *,
+        use_mask: bool = False,
+        limit: int = DEFAULT_LIMIT,
+        random: random.Random = random.Random(),
+        compress: int = 0,
+        notakeover: bool = False,
+    ) -> None:
+        """Initialize a WebSocket writer."""
+        self.protocol = protocol
+        self.transport = transport
+        self.use_mask = use_mask
+        self.get_random_bits = partial(random.getrandbits, 32)
+        self.compress = compress
+        self.notakeover = notakeover
+        self._closing = False
+        self._limit = limit
+        self._output_size = 0
+        self._compressobj: Any = None  # actually compressobj
+
+    async def send_frame(
+        self, message: bytes, opcode: int, compress: Optional[int] = None
+    ) -> None:
+        """Send a frame over the websocket with message as its payload."""
+        if self._closing and not (opcode & WSMsgType.CLOSE):
+            raise ClientConnectionResetError("Cannot write to closing transport")
+
+        # RSV are the reserved bits in the frame header. They are used to
+        # indicate that the frame is using an extension.
+        # https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
+        rsv = 0
+        # Only compress larger packets (disabled)
+        # Does small packet needs to be compressed?
+        # if self.compress and opcode < 8 and len(message) > 124:
+        if (compress or self.compress) and opcode < 8:
+            # RSV1 (rsv = 0x40) is set for compressed frames
+            # https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.3.1
+            rsv = 0x40
+
+            if compress:
+                # Do not set self._compress if compressing is for this frame
+                compressobj = self._make_compress_obj(compress)
+            else:  # self.compress
+                if not self._compressobj:
+                    self._compressobj = self._make_compress_obj(self.compress)
+                compressobj = self._compressobj
+
+            message = (
+                await compressobj.compress(message)
+                + compressobj.flush(
+                    zlib.Z_FULL_FLUSH if self.notakeover else zlib.Z_SYNC_FLUSH
+                )
+            ).removesuffix(WS_DEFLATE_TRAILING)
+            # Its critical that we do not return control to the event
+            # loop until we have finished sending all the compressed
+            # data. Otherwise we could end up mixing compressed frames
+            # if there are multiple coroutines compressing data.
+
+        msg_length = len(message)
+
+        use_mask = self.use_mask
+        mask_bit = 0x80 if use_mask else 0
+
+        # Depending on the message length, the header is assembled differently.
+        # The first byte is reserved for the opcode and the RSV bits.
+        first_byte = 0x80 | rsv | opcode
+        if msg_length < 126:
+            header = PACK_LEN1(first_byte, msg_length | mask_bit)
+            header_len = 2
+        elif msg_length < 65536:
+            header = PACK_LEN2(first_byte, 126 | mask_bit, msg_length)
+            header_len = 4
+        else:
+            header = PACK_LEN3(first_byte, 127 | mask_bit, msg_length)
+            header_len = 10
+
+        if self.transport.is_closing():
+            raise ClientConnectionResetError("Cannot write to closing transport")
+
+        # https://datatracker.ietf.org/doc/html/rfc6455#section-5.3
+        # If we are using a mask, we need to generate it randomly
+        # and apply it to the message before sending it. A mask is
+        # a 32-bit value that is applied to the message using a
+        # bitwise XOR operation. It is used to prevent certain types
+        # of attacks on the websocket protocol. The mask is only used
+        # when aiohttp is acting as a client. Servers do not use a mask.
+        if use_mask:
+            mask = PACK_RANDBITS(self.get_random_bits())
+            message = bytearray(message)
+            websocket_mask(mask, message)
+            self.transport.write(header + mask + message)
+            self._output_size += MASK_LEN
+        elif msg_length > MSG_SIZE:
+            self.transport.write(header)
+            self.transport.write(message)
+        else:
+            self.transport.write(header + message)
+
+        self._output_size += header_len + msg_length
+
+        # It is safe to return control to the event loop when using compression
+        # after this point as we have already sent or buffered all the data.
+
+        # Once we have written output_size up to the limit, we call the
+        # drain helper which waits for the transport to be ready to accept
+        # more data. This is a flow control mechanism to prevent the buffer
+        # from growing too large. The drain helper will return right away
+        # if the writer is not paused.
+        if self._output_size > self._limit:
+            self._output_size = 0
+            if self.protocol._paused:
+                await self.protocol._drain_helper()
+
+    def _make_compress_obj(self, compress: int) -> ZLibCompressor:
+        return ZLibCompressor(
+            level=zlib.Z_BEST_SPEED,
+            wbits=-compress,
+            max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
+        )
+
+    async def close(self, code: int = 1000, message: Union[bytes, str] = b"") -> None:
+        """Close the websocket, sending the specified code and message."""
+        if isinstance(message, str):
+            message = message.encode("utf-8")
+        try:
+            await self.send_frame(
+                PACK_CLOSE_CODE(code) + message, opcode=WSMsgType.CLOSE
+            )
+        finally:
+            self._closing = True