src/genro_asgi/websocket.py¶
Source from this local checkout, regenerated when the reader rebuilds.
Line links use #L<number>; a GitHub line range opens its first line.
1 # Copyright 2025 Softwell S.r.l.2 #3 # Licensed under the Apache License, Version 2.0 (the "License");4 # you may not use this file except in compliance with the License.5 # You may obtain a copy of the License at6 #7 # https://www.apache.org/licenses/LICENSE-2.08 #9 # Unless required by applicable law or agreed to in writing, software10 # distributed under the License is distributed on an "AS IS" BASIS,11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12 # See the License for the specific language governing permissions and13 # limitations under the License.14 15 """The websocket facade: one ASGI socket as an object, and nothing above it.16 17 ``WebSocket`` wraps the three things an ASGI server hands a websocket18 application — the scope, ``receive`` and ``send`` — and gives them a shape a19 reader can follow: ``accept()`` consumes the connect and answers it,20 ``receive_text()`` and ``receive_bytes()`` read one message, ``send_text()``21 and ``send_bytes()`` write one, ``close()`` ends the connection once, and22 iterating the object yields the incoming texts until the client leaves.23 24 **It knows nothing of WSX.** The protocol lives in ``wsx.py`` and the motor in25 the server; this object is the transport, so the admitted raw seam — an26 application that wants the socket itself — is served by the same class27 (`internals/10_server/055_websocket/decisions.md`).28 29 **The state is one boolean.** ``connected`` is true between the accept and the30 end, and everything that depends on the state reads it: an accept happens once,31 a close writes once, and a read or a write with nothing accepted raises. There32 is no exported state type, because the three readers of the state are those33 three rules and each is a question with a yes or a no (owner, 2026-09-06; the34 precedent is ``WorkerConnector.connected``).35 36 **A disconnect is an exception, never a value.** Every read raises37 ``WebSocketDisconnect`` when the client is gone, so a read loop never has to38 check what it got back, and the iterator ends on it.39 40 **The handshake facts are read once, in the constructor.** The path, the41 headers, the cookies and the subprotocols the client offered come off the42 scope — headers lowercased and TYTX-hydrated, cookies split out of the43 ``Cookie`` header, exactly as ``Request`` does for HTTP.44 45 ``WebSocketRegistry`` is the server's picture of what is connected: the live46 sockets, and the ``page_id → socket`` association a page writes when it opens47 its channel, so the server can address one page later. It is NEUTRAL — it does48 not know the SPA and validates nothing: whether a page belongs to the49 connection asking for it is judged by the application that holds the pool.50 """51 52 from __future__ import annotations53 54 from http.cookies import SimpleCookie55 from typing import Any, AsyncIterator56 57 from genro_tytx import from_tytx58 59 from .exceptions import WebSocketDisconnect60 from .types import Receive, Scope, Send61 62 __all__ = ["WebSocket", "WebSocketRegistry"]63 64 65 class WebSocket:66 """One ASGI websocket connection, as an object."""67 68 def __init__(self, scope: Scope, receive: Receive, send: Send) -> None:69 """Args:70 scope: the ASGI websocket scope of the handshake.71 receive: the ASGI receive callable.72 send: the ASGI send callable.73 """74 self.scope = scope75 self.asgi_receive = receive76 self.asgi_send = send77 self.accepted_subprotocol: str | None = None78 self._connected = False79 self._closed = False80 self._headers: dict[str, Any] = {}81 self._cookies: dict[str, str] = {}82 self.read_handshake()83 84 def read_handshake(self) -> None:85 """Fill the header and cookie maps off the scope.86 87 Keys are lowercased and values TYTX-hydrated; ``cookie`` stays out of88 the map and becomes ``cookies``, the way ``Request`` reads an HTTP89 request. Acts on the instance; called by ``__init__``.90 """91 cookie_header = ""92 for name, value in self.scope.get("headers") or []:93 key = name.decode("latin-1").lower()94 text = value.decode("latin-1")95 if key == "cookie":96 cookie_header = text97 else:98 self._headers[key] = from_tytx(text)99 if cookie_header:100 morsels: SimpleCookie = SimpleCookie()101 morsels.load(cookie_header)102 self._cookies = {name: morsel.value for name, morsel in morsels.items()}103 104 @property105 def connected(self) -> bool:106 """Whether this socket is accepted and not yet closed."""107 return self._connected108 109 @property110 def path(self) -> str:111 """The path of the handshake — what names the home application."""112 return str(self.scope.get("path", "/"))113 114 @property115 def headers(self) -> dict[str, Any]:116 """The handshake headers, lowercase keys, values hydrated by TYTX."""117 return self._headers118 119 @property120 def cookies(self) -> dict[str, str]:121 """The cookies of the handshake, from its ``Cookie`` header."""122 return self._cookies123 124 @property125 def subprotocols(self) -> tuple[str, ...]:126 """The subprotocols the client offered, in the order it offered them."""127 return tuple(self.scope.get("subprotocols") or ())128 129 async def accept(130 self, subprotocol: str | None = None, headers: dict[str, str] | None = None131 ) -> None:132 """Consume the connect and accept the connection.133 134 Args:135 subprotocol: the one to negotiate, when the client offered any.136 headers: response headers of the handshake — the one place a137 websocket can carry a ``Set-Cookie``.138 139 Raises:140 RuntimeError: this socket was already accepted, or the first141 message on the wire was not ``websocket.connect``.142 143 Sets ``connected``.144 """145 if self._connected or self._closed:146 raise RuntimeError("this socket cannot accept: it was accepted already")147 message = await self.asgi_receive()148 if message["type"] != "websocket.connect":149 raise RuntimeError(f"expected websocket.connect, got {message['type']}")150 accept: dict[str, Any] = {"type": "websocket.accept"}151 if subprotocol is not None:152 accept["subprotocol"] = subprotocol153 self.accepted_subprotocol = subprotocol154 if headers is not None:155 accept["headers"] = [156 (name.encode("latin-1"), value.encode("latin-1"))157 for name, value in headers.items()158 ]159 await self.asgi_send(accept)160 self._connected = True161 162 async def refuse(self, code: int = 1008, reason: str = "") -> None:163 """Turn the handshake away without accepting it.164 165 Args:166 code: the close code the client sees.167 reason: the text that travels with it.168 169 Raises:170 RuntimeError: this socket was accepted already — turning away what171 is already in is a ``close``, not a refusal.172 173 The connect is consumed first: a close written before it is read leaves174 that message on the wire. Sets ``connected`` to false, so nothing can175 be written afterwards.176 """177 if self._connected or self._closed:178 raise RuntimeError("this socket cannot refuse: it was accepted already")179 await self.asgi_receive()180 self._closed = True181 await self.asgi_send({"type": "websocket.close", "code": code, "reason": reason})182 183 async def close(self, code: int = 1000, reason: str = "") -> None:184 """End the connection, once.185 186 Args:187 code: the websocket close code.188 reason: the text that travels with it.189 190 Raises:191 RuntimeError: nothing was accepted yet — before the accept a192 handshake is turned away with ``refuse``, which is what a193 hostile Origin gets. Everything judged AFTER the accept — the194 home application's cookie, an invalid credential, a server that195 is not running — is accepted first and closed here with its196 code, so the browser can read why.197 198 Sets ``connected`` to false. Calling it again writes nothing.199 """200 if self._closed:201 return202 if not self._connected:203 raise RuntimeError("this socket cannot close: it is not accepted")204 self._connected = False205 self._closed = True206 await self.asgi_send({"type": "websocket.close", "code": code, "reason": reason})207 208 async def receive_text(self) -> str:209 """The next message, as text.210 211 Returns:212 The message's text.213 214 Raises:215 RuntimeError: this socket is not connected.216 TypeError: the message carried bytes.217 WebSocketDisconnect: the client is gone.218 """219 message = await self.read_message()220 if message.get("text") is None:221 raise TypeError("this message is binary: read it with receive_bytes()")222 return str(message["text"])223 224 async def receive_bytes(self) -> bytes:225 """The next message, as bytes.226 227 Returns:228 The message's bytes.229 230 Raises:231 RuntimeError: this socket is not connected.232 TypeError: the message carried text.233 WebSocketDisconnect: the client is gone.234 """235 message = await self.read_message()236 if message.get("bytes") is None:237 raise TypeError("this message is text: read it with receive_text()")238 return bytes(message["bytes"])239 240 async def read_message(self) -> dict[str, Any]:241 """One raw ASGI message, with the disconnect turned into an exception.242 243 Returns:244 The ``websocket.receive`` message as it came.245 246 Raises:247 RuntimeError: this socket is not connected.248 WebSocketDisconnect: the client is gone; ``connected`` is false249 from here on.250 """251 if not self._connected:252 raise RuntimeError("this socket is not connected")253 message = await self.asgi_receive()254 if message["type"] == "websocket.disconnect":255 self._connected = False256 self._closed = True257 raise WebSocketDisconnect(message.get("code", 1000), message.get("reason", ""))258 return dict(message)259 260 async def send_text(self, text: str) -> None:261 """Write one text message.262 263 Raises:264 RuntimeError: this socket is not connected.265 """266 await self.write_message({"type": "websocket.send", "text": text})267 268 async def send_bytes(self, data: bytes) -> None:269 """Write one binary message.270 271 Raises:272 RuntimeError: this socket is not connected.273 """274 await self.write_message({"type": "websocket.send", "bytes": data})275 276 async def write_message(self, message: dict[str, Any]) -> None:277 """Write one raw ASGI message.278 279 Args:280 message: the ``websocket.send`` message to write.281 282 Raises:283 RuntimeError: this socket is not connected — nothing is written to284 a socket nobody accepted, and nothing after a close.285 """286 if not self._connected:287 raise RuntimeError("this socket is not connected")288 await self.asgi_send(message)289 290 async def __aiter__(self) -> AsyncIterator[str]:291 """The incoming texts, until the client leaves.292 293 The disconnect ends the loop instead of raising: a read loop's ordinary294 end is the client going away.295 """296 while True:297 try:298 yield await self.receive_text()299 except WebSocketDisconnect:300 return301 302 303 class WebSocketRegistry:304 """The live sockets of one server, and which one each page speaks on."""305 306 def __init__(self) -> None:307 self._sockets: list[WebSocket] = []308 self._page_sockets: dict[str, WebSocket] = {}309 310 def register(self, socket: WebSocket) -> None:311 """Take one accepted socket into the picture.312 313 Args:314 socket: the facade of a connection that was just accepted.315 """316 self._sockets.append(socket)317 318 def unregister(self, socket: WebSocket) -> None:319 """Take one socket out, and every page that still speaks on IT.320 321 Args:322 socket: the connection that ended.323 324 A page whose association has moved to another socket — a reconnection325 that happened before this one closed — is left alone: the comparison is326 on the socket itself, never on the page. A socket that was never327 registered is no error: the ``finally`` of a handshake that failed328 before the accept comes through here too.329 """330 if socket in self._sockets:331 self._sockets.remove(socket)332 for page_id in [page for page, bound in self._page_sockets.items() if bound is socket]:333 del self._page_sockets[page_id]334 335 def bind_page(self, page_id: str, socket: WebSocket) -> None:336 """Say that this page speaks on this socket.337 338 Args:339 page_id: the page opening its channel.340 socket: the connection its messages arrive on.341 342 A page already bound is REBOUND, with no error: a browser that lost its343 socket and opened a new one says so again, and the association follows344 it. One socket carries as many pages as the browser has under that345 connection.346 """347 self._page_sockets[page_id] = socket348 349 def get_page_socket(self, page_id: str) -> WebSocket | None:350 """The socket that page speaks on, or ``None`` when it speaks on none.351 352 Args:353 page_id: the page to address.354 355 Returns:356 Its socket, or ``None`` — the page never opened a channel, its357 socket is gone, or the page itself is.358 """359 return self._page_sockets.get(page_id)360 361 def snapshot(self) -> list[WebSocket]:362 """The live sockets, in the order they were accepted."""363 return list(self._sockets)