1"""WebSockets: the accept response, upgrade detection, and the `ws` a view holds.
2
3A websocket is an ordinary `View` with an `async def websocket(self, ws)`
4handler. The request runs through the normal pipeline (middleware,
5`before_request`, URL kwargs), so auth is inherited; the view then
6returns a `WebSocketResponse`, which the server frames as a bodiless
7101 before handing the raw connection to a `WebSocket` and running the
8view's coroutine for the life of the socket.
9
10`WebSocket` is the whole surface a view sees:
11
12 async for message in ws # str or bytes; ends quietly on any close
13 await ws.send(str | bytes) # raises WebSocketClosed once closed
14 await ws.close(code=1000, reason="")
15 ws.subprotocol # the negotiated one, or None
16
17Framing is in `websocket_frames`; this module is the message layer on
18top of it: fragment reassembly, control frames, the close handshake,
19the keepalive ping loop, and the size cap — plus the handshake pieces
20(upgrade detection, subprotocol selection, the accept response).
21"""
22
23import asyncio
24import base64
25from collections.abc import Callable, Coroutine
26from types import TracebackType
27from typing import TYPE_CHECKING, Any, Protocol, Self
28
29from .response import Response
30from .websocket_frames import (
31 CLOSE_ABNORMAL,
32 CLOSE_GOING_AWAY,
33 CLOSE_INVALID_PAYLOAD,
34 CLOSE_MESSAGE_TOO_BIG,
35 CLOSE_NO_STATUS,
36 CLOSE_NORMAL,
37 CLOSE_POLICY_VIOLATION,
38 CLOSE_PROTOCOL_ERROR,
39 OP_BINARY,
40 OP_CLOSE,
41 OP_CONTINUATION,
42 OP_PING,
43 OP_PONG,
44 OP_TEXT,
45 IncompleteFrame,
46 PayloadTooLarge,
47 ProtocolError,
48 compute_accept_key,
49 encode_close,
50 encode_frame,
51 parse_close_payload,
52 read_frame,
53)
54
55if TYPE_CHECKING:
56 from .request import Request
57
58# Keepalive: a PING every PING_INTERVAL seconds, and the socket is failed
59# with 1001 when no PONG arrives within PING_TIMEOUT of it. The same
60# defaults as the `websockets` library, and well inside the 55 seconds a
61# router like Heroku's allows a connection to sit idle.
62PING_INTERVAL = 20.0
63PING_TIMEOUT = 20.0
64
65# How long `close()` waits for the peer to answer our CLOSE frame before
66# closing the transport anyway. The server tightens this during drain.
67CLOSE_TIMEOUT = 5.0
68
69# Outbound frames go to the transport in chunks of this size, and each
70# chunk gets WRITE_TIMEOUT to reach the kernel's send buffer. A live peer
71# drains its receive window in milliseconds, however slow its link; one
72# that has stopped reading never does, and without a bound a `send` would
73# sit in `drain()` holding the write lock, keeping the keepalive from ever
74# closing the socket. Chunking makes the bound "no progress for this
75# long" rather than "the whole message within this long".
76WRITE_CHUNK = 64 * 1024
77WRITE_TIMEOUT = PING_TIMEOUT
78
79# Messages the reader may hold for a view that is not iterating yet.
80# Bounded so a push-only view (one that only ever sends) cannot make the
81# reader buffer without limit; PING/PONG keep flowing while there is room.
82RECEIVE_QUEUE_SIZE = 16
83
84
85class WebSocketTransport(Protocol):
86 """What `WebSocket` needs from the connection underneath it.
87
88 The server's `Connection` satisfies this as-is; the test client
89 supplies the same class over its end of a socketpair.
90 """
91
92 async def recv(self, n: int) -> bytes: ...
93
94 async def sendall(self, data: bytes) -> None: ...
95
96 def close(self) -> None: ...
97
98 def abort(self) -> None: ...
99
100
101class WebSocketClosed(Exception):
102 """The socket is closed: raised by `send()` once no more frames can go out.
103
104 A closed peer is a normal ending, not an error — a view that keeps
105 sending after the browser left sees this and stops.
106 """
107
108 def __init__(self, code: int = CLOSE_ABNORMAL, reason: str = "") -> None:
109 self.code = code
110 self.reason = reason
111 super().__init__(f"WebSocket closed ({code}) {reason}".rstrip())
112
113
114class WebSocket:
115 """One open websocket, from the 101 to the transport close.
116
117 The server enters it (`async with ws:` starts the reader and the
118 keepalive pings, leaving it stops them) and hands it to the view,
119 which iterates it, sends on it, and may close it.
120 """
121
122 def __init__(
123 self,
124 transport: WebSocketTransport,
125 *,
126 subprotocol: str | None = None,
127 max_message_size: int,
128 ping_interval: float = PING_INTERVAL,
129 ping_timeout: float = PING_TIMEOUT,
130 ) -> None:
131 self._transport = transport
132 self.subprotocol = subprotocol
133 self._max_message_size = max_message_size
134 self._ping_interval = ping_interval
135 self._ping_timeout = ping_timeout
136
137 # Complete messages for the view. `_activity` wakes an iterator
138 # blocked on an empty queue when a message lands or the socket ends.
139 self._messages: asyncio.Queue[str | bytes] = asyncio.Queue(
140 maxsize=RECEIVE_QUEUE_SIZE
141 )
142 self._activity = asyncio.Event()
143 self._write_lock = asyncio.Lock()
144
145 # Set once the socket is over, whichever side ended it and however:
146 # the peer's CLOSE, ours timing out, a failure, or shutdown.
147 self._closed = asyncio.Event()
148 # Our CLOSE frame has gone out (or failed to). Between that and
149 # `_closed` being set, `send` must already refuse and report the
150 # code we sent, which is why `_send_close` records it too.
151 self._close_sent = False
152 self._close_code = CLOSE_ABNORMAL
153 self._close_reason = ""
154 self._pong_received = asyncio.Event()
155 # True while the reader is parked on a full receive queue: the
156 # view has stopped iterating, so nothing more (PONGs included) is
157 # being read. A ping timeout in that state is the view's doing,
158 # not the peer's, and is reported as such.
159 self._receive_full = False
160
161 self._reader_task: asyncio.Task[None] | None = None
162 self._ping_task: asyncio.Task[None] | None = None
163
164 # ------------------------------------------------------------------
165 # The view-facing surface
166 # ------------------------------------------------------------------
167
168 @property
169 def closed(self) -> bool:
170 return self._closed.is_set()
171
172 @property
173 def close_code(self) -> int:
174 """The close code once `closed`; 1006 when the socket ended without a CLOSE frame."""
175 return self._close_code
176
177 @property
178 def close_reason(self) -> str:
179 return self._close_reason
180
181 async def wait_closed(self) -> None:
182 """Block until the socket is over, however it ends."""
183 await self._closed.wait()
184
185 def __aiter__(self) -> WebSocket:
186 return self
187
188 async def __anext__(self) -> str | bytes:
189 while True:
190 try:
191 return self._messages.get_nowait()
192 except asyncio.QueueEmpty:
193 pass
194 if self.closed:
195 raise StopAsyncIteration
196 self._activity.clear()
197 await self._activity.wait()
198
199 async def send(self, message: str | bytes) -> None:
200 if self.closed or self._close_sent:
201 raise WebSocketClosed(self._close_code, self._close_reason)
202 if isinstance(message, str):
203 opcode, payload = OP_TEXT, message.encode()
204 elif isinstance(message, bytes | bytearray | memoryview):
205 opcode, payload = OP_BINARY, bytes(message)
206 else:
207 raise TypeError(
208 f"WebSocket messages are str or bytes, not {type(message).__name__}"
209 )
210 try:
211 await self._write(encode_frame(opcode, payload))
212 except OSError:
213 self._end(CLOSE_ABNORMAL)
214 raise WebSocketClosed(self._close_code, self._close_reason) from None
215
216 async def close(
217 self,
218 code: int = CLOSE_NORMAL,
219 reason: str = "",
220 *,
221 timeout: float = CLOSE_TIMEOUT,
222 ) -> None:
223 """Start the closing handshake and finish it within `timeout`.
224
225 Sends our CLOSE once, waits for the peer's (the reader answers
226 it), then closes the transport. Safe to call more than once and
227 from a task other than the one iterating.
228 """
229 if self.closed:
230 return
231 try:
232 if not self._close_sent:
233 await asyncio.wait_for(self._send_close(code, reason), timeout)
234 await asyncio.wait_for(self._closed.wait(), timeout)
235 except TimeoutError:
236 pass
237 self._end(code, reason)
238
239 # ------------------------------------------------------------------
240 # Server-facing lifecycle
241 # ------------------------------------------------------------------
242
243 async def __aenter__(self) -> Self:
244 """Begin reading frames and pinging. Entered once, after the 101."""
245 self._reader_task = asyncio.create_task(self._read_loop())
246 self._ping_task = asyncio.create_task(self._ping_loop())
247 return self
248
249 async def __aexit__(
250 self,
251 exc_type: type[BaseException] | None,
252 exc: BaseException | None,
253 tb: TracebackType | None,
254 ) -> None:
255 """Stop the background tasks and make sure the transport is closed."""
256 self._end()
257 tasks = [t for t in (self._ping_task, self._reader_task) if t is not None]
258 for task in tasks:
259 task.cancel()
260 await asyncio.gather(*tasks, return_exceptions=True)
261
262 # ------------------------------------------------------------------
263 # Internals
264 # ------------------------------------------------------------------
265
266 async def _write(self, data: bytes) -> None:
267 async with self._write_lock:
268 complete = False
269 try:
270 for start in range(0, max(len(data), 1), WRITE_CHUNK):
271 await asyncio.wait_for(
272 self._transport.sendall(data[start : start + WRITE_CHUNK]),
273 WRITE_TIMEOUT,
274 )
275 complete = True
276 except TimeoutError:
277 # The peer stopped reading. Abort rather than close: a
278 # close would wait for the same buffer to flush.
279 self._transport.abort()
280 raise ConnectionResetError("Peer stopped reading") from None
281 except asyncio.CancelledError:
282 if self.closed:
283 # Teardown cancelled a writer the peer never drained;
284 # a close would hang on it.
285 self._transport.abort()
286 elif len(data) > WRITE_CHUNK and not complete:
287 # A frame that was queued in pieces is now missing its
288 # tail; anything sent after it would be read as payload.
289 # A single-chunk frame is queued atomically and is fine.
290 self._transport.abort()
291 self._end(CLOSE_ABNORMAL)
292 raise
293
294 async def _send_close(self, code: int, reason: str) -> None:
295 if self._close_sent:
296 return
297 self._close_sent = True
298 self._close_code = code
299 self._close_reason = reason
300 try:
301 await self._write(encode_close(code, reason))
302 except OSError:
303 self._end(CLOSE_ABNORMAL)
304
305 def _end(self, code: int | None = None, reason: str = "") -> None:
306 """The socket is over: record why (if told), wake everyone, drop the transport."""
307 if self.closed:
308 return
309 if code is not None:
310 self._close_code = code
311 self._close_reason = reason
312 self._closed.set()
313 self._activity.set()
314 try:
315 # A close flushes what is queued. The paths that know the peer
316 # has stopped draining (a write timeout, a CLOSE that could not
317 # go out, a teardown cancel of a writer still holding the lock)
318 # abort the transport themselves.
319 self._transport.close()
320 except OSError:
321 pass
322
323 async def _fail(self, code: int, reason: str) -> None:
324 """Fail the connection (RFC 6455 §7.1.7): a best-effort CLOSE, then drop it.
325
326 Never waits on a stuck writer: the CLOSE gets `CLOSE_TIMEOUT` to go
327 out behind whatever holds the write lock, then the socket ends
328 regardless.
329 """
330 try:
331 await asyncio.wait_for(self._send_close(code, reason), CLOSE_TIMEOUT)
332 except TimeoutError:
333 # A 125-byte frame that cannot go out in that long is stuck
334 # behind a writer the peer is not draining.
335 self._transport.abort()
336 self._end(code, reason)
337
338 async def _read_loop(self) -> None:
339 fragment_opcode: int | None = None
340 fragment = bytearray()
341 try:
342 while not self.closed:
343 frame = await read_frame(self._transport.recv, self._max_message_size)
344 opcode = frame.opcode
345
346 if opcode == OP_PING:
347 await self._write(encode_frame(OP_PONG, frame.payload))
348 continue
349 if opcode == OP_PONG:
350 self._pong_received.set()
351 continue
352 if opcode == OP_CLOSE:
353 peer = parse_close_payload(frame.payload)
354 echo = CLOSE_NORMAL if peer.code == CLOSE_NO_STATUS else peer.code
355 await self._send_close(echo, "")
356 self._end(peer.code, peer.reason)
357 return
358
359 if opcode == OP_CONTINUATION:
360 if fragment_opcode is None:
361 raise ProtocolError("Continuation frame without a message")
362 fragment += frame.payload
363 elif opcode in (OP_TEXT, OP_BINARY):
364 if fragment_opcode is not None:
365 raise ProtocolError("New message inside a fragmented one")
366 fragment_opcode = opcode
367 fragment = bytearray(frame.payload)
368 else:
369 raise ProtocolError(f"Unexpected opcode {opcode:#x}")
370
371 if len(fragment) > self._max_message_size:
372 raise PayloadTooLarge(len(fragment), self._max_message_size)
373
374 if frame.fin:
375 payload = bytes(fragment)
376 message: str | bytes
377 if fragment_opcode == OP_TEXT:
378 message = payload.decode()
379 else:
380 message = payload
381 fragment_opcode = None
382 fragment = bytearray()
383 if self._messages.full():
384 self._receive_full = True
385 await self._messages.put(message)
386 self._receive_full = False
387 self._activity.set()
388 except PayloadTooLarge:
389 await self._fail(CLOSE_MESSAGE_TOO_BIG, "Message too big")
390 except UnicodeDecodeError:
391 await self._fail(CLOSE_INVALID_PAYLOAD, "Invalid UTF-8")
392 except ProtocolError as exc:
393 await self._fail(CLOSE_PROTOCOL_ERROR, str(exc))
394 except IncompleteFrame, OSError:
395 # The peer went away without a CLOSE — an abnormal close
396 # (the default code) from our side, and a perfectly normal
397 # ending for the view. The finally records it.
398 pass
399 finally:
400 self._end()
401
402 async def _ping_loop(self) -> None:
403 while not self.closed:
404 await asyncio.sleep(self._ping_interval)
405 if self.closed:
406 return
407 self._pong_received.clear()
408 try:
409 # Bounded too: a `send` stuck on a dead peer holds the write
410 # lock, and the keepalive must not queue behind it.
411 await asyncio.wait_for(
412 self._write(encode_frame(OP_PING, b"")), self._ping_timeout
413 )
414 await asyncio.wait_for(self._pong_received.wait(), self._ping_timeout)
415 except TimeoutError:
416 # Before OSError: the builtin TimeoutError is an OSError.
417 if self._receive_full:
418 await self._fail(
419 CLOSE_POLICY_VIOLATION,
420 "Receive buffer full: the handler stopped reading",
421 )
422 else:
423 await self._fail(CLOSE_GOING_AWAY, "Ping timeout")
424 return
425 except OSError:
426 self._end(CLOSE_ABNORMAL)
427 return
428
429
430class WebSocketResponse(Response):
431 """The accepted upgrade: a bodiless 101 the framework constructs for a view.
432
433 It flows through `after_response` like any response (so middleware
434 cookies and headers land on the 101), and the server, on seeing it,
435 frames the header block, opens a `WebSocket` on the connection and
436 runs `handler(ws)` — the view's `websocket()` — until the socket
437 ends, when its resource closers run. It is not a streaming response:
438 no view code has run when middleware sees it, so anything a
439 middleware releases at the end of an ordinary request can be released
440 here too and re-acquired inside the socket. Application code never
441 builds one; `Response` refuses 1xx status codes from callers, and
442 this class is the framework's own 101.
443 """
444
445 _default_status_code = 101
446
447 def __init__(
448 self,
449 *,
450 request: Request,
451 handler: Callable[[WebSocket], Coroutine[Any, Any, None]],
452 subprotocol: str | None,
453 max_message_size: int,
454 ) -> None:
455 super().__init__()
456 self.request = request
457 self.handler = handler
458 self.subprotocol = subprotocol
459 self.max_message_size = max_message_size
460
461 self.headers["Sec-WebSocket-Accept"] = compute_accept_key(
462 request.headers["Sec-WebSocket-Key"].strip()
463 )
464 if subprotocol:
465 self.headers["Sec-WebSocket-Protocol"] = subprotocol
466
467
468def is_websocket_upgrade(request: Request) -> bool:
469 """True for a well-formed RFC 6455 opening handshake, else an ordinary GET.
470
471 Anything short of the full set — wrong method, missing `Upgrade`,
472 a `Connection` without the `upgrade` token, a version other than 13,
473 a key that is not 16 base64 bytes — is served as a normal GET; the
474 browser then fails the socket on its own. (The RFC's 426 for other
475 versions is deliberately not implemented.)
476 """
477 if request.method != "GET":
478 return False
479 if request.headers.get("Upgrade", "").strip().lower() != "websocket":
480 return False
481 connection_tokens = {
482 token.strip().lower()
483 for token in request.headers.get("Connection", "").split(",")
484 }
485 if "upgrade" not in connection_tokens:
486 return False
487 if request.headers.get("Sec-WebSocket-Version", "").strip() != "13":
488 return False
489 key = request.headers.get("Sec-WebSocket-Key", "").strip()
490 try:
491 return len(base64.b64decode(key, validate=True)) == 16
492 except ValueError:
493 return False
494
495
496def select_subprotocol(
497 offered_header: str | None, supported: tuple[str, ...]
498) -> str | None:
499 """Pick the subprotocol to echo back, or None to echo nothing.
500
501 `offered_header` is the request's `Sec-WebSocket-Protocol` value; repeated
502 header lines arrive comma-joined. The client's order is the client's
503 preference order (RFC 6455 Section 4.1), so the first offered name that we
504 support wins even if `supported` lists another one first. Subprotocol names
505 are tokens and compare case-sensitively.
506 """
507 if not offered_header:
508 return None
509 for offered in offered_header.split(","):
510 name = offered.strip()
511 if name in supported:
512 return name
513 return None