Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/_websocket/writer.py: 23%
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"""WebSocket protocol versions 13 and 8."""
3import asyncio
4import random
5import sys
6from asyncio.base_events import BaseEventLoop
7from functools import partial
8from typing import Final
10from ..base_protocol import BaseProtocol
11from ..client_exceptions import ClientConnectionResetError
12from ..compression_utils import ZLibBackend, ZLibCompressor
13from ..helpers import DEFAULT_CHUNK_SIZE
14from .helpers import (
15 MASK_LEN,
16 MSG_SIZE,
17 PACK_CLOSE_CODE,
18 PACK_LEN1,
19 PACK_LEN2,
20 PACK_LEN3,
21 PACK_RANDBITS,
22 websocket_mask,
23)
24from .models import WS_DEFLATE_TRAILING, WSMsgType
26# WebSocket opcode boundary: opcodes 0-7 are data frames, 8-15 are control frames
27# Control frames (ping, pong, close) are never compressed
28WS_CONTROL_FRAME_OPCODE: Final[int] = 8
30# For websockets, keeping latency low is extremely important as implementations
31# generally expect to be able to send and receive messages quickly. We use a
32# larger chunk size to reduce the number of executor calls and avoid task
33# creation overhead, since both are significant sources of latency when chunks
34# are small. A size of 16KiB was chosen as a balance between avoiding task
35# overhead and not blocking the event loop too long with synchronous compression.
37WEBSOCKET_MAX_SYNC_CHUNK_SIZE = 16 * 1024
40class WebSocketWriter:
41 """WebSocket writer.
43 The writer is responsible for sending messages to the client. It is
44 created by the protocol when a connection is established. The writer
45 should avoid implementing any application logic and should only be
46 concerned with the low-level details of the WebSocket protocol.
47 """
49 def __init__(
50 self,
51 protocol: BaseProtocol,
52 transport: asyncio.WriteTransport,
53 *,
54 use_mask: bool = False,
55 limit: int = DEFAULT_CHUNK_SIZE,
56 random: random.Random = random.Random(),
57 compress: int = 0,
58 notakeover: bool = False,
59 ) -> None:
60 """Initialize a WebSocket writer."""
61 self.protocol = protocol
62 self.transport = transport
63 self.use_mask = use_mask
64 self.get_random_bits = partial(random.getrandbits, 32)
65 self.compress = compress
66 self.notakeover = notakeover
67 self._closing = False
68 self._limit = limit
69 self._output_size = 0
70 self._compressobj: ZLibCompressor | None = None
71 self._send_lock = asyncio.Lock()
72 self._background_tasks: set[asyncio.Task[None]] = set()
74 async def send_frame(
75 self, message: bytes, opcode: int, compress: int | None = None
76 ) -> None:
77 """Send a frame over the websocket with message as its payload."""
78 if self._closing and not (opcode & WSMsgType.CLOSE):
79 raise ClientConnectionResetError("Cannot write to closing transport")
81 if not (compress or self.compress) or opcode >= WS_CONTROL_FRAME_OPCODE:
82 # Non-compressed frames don't need lock or shield
83 self._write_websocket_frame(message, opcode, 0)
84 elif len(message) <= WEBSOCKET_MAX_SYNC_CHUNK_SIZE:
85 # Small compressed payloads - compress synchronously in event loop
86 # We need the lock even though sync compression has no await points.
87 # This prevents small frames from interleaving with large frames that
88 # compress in the executor, avoiding compressor state corruption.
89 async with self._send_lock:
90 self._send_compressed_frame_sync(message, opcode, compress)
91 else:
92 # Large compressed frames need shield to prevent corruption
93 # For large compressed frames, the entire compress+send
94 # operation must be atomic. If cancelled after compression but
95 # before send, the compressor state would be advanced but data
96 # not sent, corrupting subsequent frames.
97 # Create a task to shield from cancellation
98 # The lock is acquired inside the shielded task so the entire
99 # operation (lock + compress + send) completes atomically.
100 # Use eager_start to avoid scheduling overhead
101 coro = self._send_compressed_frame_async_locked(message, opcode, compress)
102 if sys.version_info >= (3, 14):
103 loop = asyncio.get_running_loop()
104 if isinstance(loop, BaseEventLoop):
105 send_task = asyncio.create_task(coro, eager_start=True)
106 else:
107 send_task = asyncio.Task(coro, loop=loop, eager_start=True)
108 elif sys.version_info >= (3, 12):
109 send_task = asyncio.Task(
110 coro, loop=asyncio.get_running_loop(), eager_start=True
111 )
112 else:
113 send_task = asyncio.create_task(coro)
114 # Keep a strong reference to prevent garbage collection
115 self._background_tasks.add(send_task)
116 send_task.add_done_callback(self._background_tasks.discard)
117 await asyncio.shield(send_task)
119 # It is safe to return control to the event loop when using compression
120 # after this point as we have already sent or buffered all the data.
121 # Once we have written output_size up to the limit, we call the
122 # drain helper which waits for the transport to be ready to accept
123 # more data. This is a flow control mechanism to prevent the buffer
124 # from growing too large. The drain helper will return right away
125 # if the writer is not paused.
126 if self._output_size > self._limit:
127 self._output_size = 0
128 if self.protocol._paused:
129 await self.protocol._drain_helper()
131 def _write_websocket_frame(self, message: bytes, opcode: int, rsv: int) -> None:
132 """
133 Write a websocket frame to the transport.
135 This method handles frame header construction, masking, and writing to transport.
136 It does not handle compression or flow control - those are the responsibility
137 of the caller.
138 """
139 msg_length = len(message)
141 use_mask = self.use_mask
142 mask_bit = 0x80 if use_mask else 0
144 # Depending on the message length, the header is assembled differently.
145 # The first byte is reserved for the opcode and the RSV bits.
146 first_byte = 0x80 | rsv | opcode
147 if msg_length < 126:
148 header = PACK_LEN1(first_byte, msg_length | mask_bit)
149 header_len = 2
150 elif msg_length < 65536:
151 header = PACK_LEN2(first_byte, 126 | mask_bit, msg_length)
152 header_len = 4
153 else:
154 header = PACK_LEN3(first_byte, 127 | mask_bit, msg_length)
155 header_len = 10
157 if self.transport.is_closing():
158 raise ClientConnectionResetError("Cannot write to closing transport")
160 # https://datatracker.ietf.org/doc/html/rfc6455#section-5.3
161 # If we are using a mask, we need to generate it randomly
162 # and apply it to the message before sending it. A mask is
163 # a 32-bit value that is applied to the message using a
164 # bitwise XOR operation. It is used to prevent certain types
165 # of attacks on the websocket protocol. The mask is only used
166 # when aiohttp is acting as a client. Servers do not use a mask.
167 if use_mask:
168 mask = PACK_RANDBITS(self.get_random_bits())
169 message_arr = bytearray(message)
170 websocket_mask(mask, message_arr)
171 self.transport.write(header + mask + message_arr)
172 self._output_size += MASK_LEN
173 elif msg_length > MSG_SIZE:
174 self.transport.write(header)
175 self.transport.write(message)
176 else:
177 self.transport.write(header + message)
179 self._output_size += header_len + msg_length
181 def _get_compressor(self, compress: int | None) -> ZLibCompressor:
182 """Get or create a compressor object for the given compression level."""
183 if compress:
184 # Do not set self._compress if compressing is for this frame
185 return ZLibCompressor(
186 level=ZLibBackend.Z_BEST_SPEED,
187 wbits=-compress,
188 max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
189 )
190 if not self._compressobj:
191 self._compressobj = ZLibCompressor(
192 level=ZLibBackend.Z_BEST_SPEED,
193 wbits=-self.compress,
194 max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
195 )
196 return self._compressobj
198 def _send_compressed_frame_sync(
199 self, message: bytes, opcode: int, compress: int | None
200 ) -> None:
201 """
202 Synchronous send for small compressed frames.
204 This is used for small compressed payloads that compress synchronously in the event loop.
205 Since there are no await points, this is inherently cancellation-safe.
206 """
207 # RSV are the reserved bits in the frame header. They are used to
208 # indicate that the frame is using an extension.
209 # https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
210 compressobj = self._get_compressor(compress)
211 # (0x40) RSV1 is set for compressed frames
212 # https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.3.1
213 self._write_websocket_frame(
214 (
215 compressobj.compress_sync(message)
216 + compressobj.flush(
217 ZLibBackend.Z_FULL_FLUSH
218 if self.notakeover
219 else ZLibBackend.Z_SYNC_FLUSH
220 )
221 ).removesuffix(WS_DEFLATE_TRAILING),
222 opcode,
223 0x40,
224 )
226 async def _send_compressed_frame_async_locked(
227 self, message: bytes, opcode: int, compress: int | None
228 ) -> None:
229 """
230 Async send for large compressed frames with lock.
232 Acquires the lock and compresses large payloads asynchronously in
233 the executor. The lock is held for the entire operation to ensure
234 the compressor state is not corrupted by concurrent sends.
236 MUST be run shielded from cancellation. If cancelled after
237 compression but before sending, the compressor state would be
238 advanced but data not sent, corrupting subsequent frames.
239 """
240 async with self._send_lock:
241 # RSV are the reserved bits in the frame header. They are used to
242 # indicate that the frame is using an extension.
243 # https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
244 compressobj = self._get_compressor(compress)
245 # (0x40) RSV1 is set for compressed frames
246 # https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.3.1
247 self._write_websocket_frame(
248 (
249 await compressobj.compress(message)
250 + compressobj.flush(
251 ZLibBackend.Z_FULL_FLUSH
252 if self.notakeover
253 else ZLibBackend.Z_SYNC_FLUSH
254 )
255 ).removesuffix(WS_DEFLATE_TRAILING),
256 opcode,
257 0x40,
258 )
260 async def close(self, code: int = 1000, message: bytes | str = b"") -> None:
261 """Close the websocket, sending the specified code and message."""
262 if isinstance(message, str):
263 message = message.encode("utf-8")
264 try:
265 await self.send_frame(
266 PACK_CLOSE_CODE(code) + message, opcode=WSMsgType.CLOSE
267 )
268 finally:
269 self._closing = True