from __future__ import annotations import typing if typing.TYPE_CHECKING: from ...._async.response import AsyncHTTPResponse from wsproto import ConnectionType, WSConnection from wsproto.events import ( AcceptConnection, BytesMessage, CloseConnection, Ping, Pong, Request, TextMessage, ) from wsproto.extensions import PerMessageDeflate from wsproto.utilities import ProtocolError as WebSocketProtocolError from ....backend import HttpVersion from ....exceptions import ProtocolError from .protocol import AsyncExtensionFromHTTP class AsyncWebSocketExtensionFromHTTP(AsyncExtensionFromHTTP): def __init__(self) -> None: super().__init__() self._protocol = WSConnection(ConnectionType.CLIENT) self._request_headers: dict[str, str] | None = None self._remote_shutdown: bool = False @staticmethod def supported_svn() -> set[HttpVersion]: return {HttpVersion.h11} @staticmethod def implementation() -> str: return "wsproto" async def start(self, response: AsyncHTTPResponse) -> None: await super().start(response) fake_http_response = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n" fake_http_response += b"Sec-Websocket-Accept: " accept_token: str | None = response.headers.get("Sec-Websocket-Accept") if accept_token is None: raise ProtocolError( "The WebSocket HTTP extension requires 'Sec-Websocket-Accept' header in the server response but was not present." ) fake_http_response += accept_token.encode() + b"\r\n" if "sec-websocket-extensions" in response.headers: fake_http_response += ( b"Sec-Websocket-Extensions: " + response.headers.get("sec-websocket-extensions").encode() # type: ignore[union-attr] + b"\r\n" ) fake_http_response += b"\r\n" try: self._protocol.receive_data(fake_http_response) except WebSocketProtocolError as e: raise ProtocolError from e # Defensive: should never occur! event = next(self._protocol.events()) if not isinstance(event, AcceptConnection): raise RuntimeError( "The WebSocket state-machine did not pass the handshake phase when expected." ) def headers(self, http_version: HttpVersion) -> dict[str, str]: """Specific HTTP headers required (request) before the 101 status response.""" if self._request_headers is not None: return self._request_headers try: raw_data_to_socket = self._protocol.send( Request( host="example.com", target="/", extensions=(PerMessageDeflate(),) ) ) except WebSocketProtocolError as e: raise ProtocolError from e # Defensive: should never occur! raw_headers = raw_data_to_socket.split(b"\r\n")[2:-2] request_headers: dict[str, str] = {} for raw_header in raw_headers: k, v = raw_header.decode().split(": ") request_headers[k.lower()] = v if http_version != HttpVersion.h11: del request_headers["upgrade"] del request_headers["connection"] request_headers[":protocol"] = "websocket" request_headers[":method"] = "CONNECT" self._request_headers = request_headers return request_headers async def close(self) -> None: """End/Notify close for sub protocol.""" if self._dsa is not None: if self._police_officer is not None: async with self._police_officer.borrow(self._response): if self._remote_shutdown is False: try: data_to_send: bytes = self._protocol.send( CloseConnection(0) ) except WebSocketProtocolError: pass else: async with self._write_error_catcher(): await self._dsa.sendall(data_to_send) await self._dsa.close() self._dsa = None else: self._dsa = None if self._response is not None: if self._police_officer is not None: self._police_officer.forget(self._response) else: await self._response.close() self._response = None self._police_officer = None async def next_payload(self) -> str | bytes | None: """Unpack the next received message/payload from remote.""" if self._dsa is None or self._response is None or self._police_officer is None: raise OSError("The HTTP extension is closed or uninitialized") async with self._police_officer.borrow(self._response): for event in self._protocol.events(): if isinstance(event, TextMessage): return event.data elif isinstance(event, BytesMessage): return event.data elif isinstance(event, CloseConnection): self._remote_shutdown = True await self.close() return None elif isinstance(event, Ping): try: data_to_send: bytes = self._protocol.send(event.response()) except WebSocketProtocolError as e: await self.close() raise ProtocolError from e async with self._write_error_catcher(): await self._dsa.sendall(data_to_send) while True: async with self._read_error_catcher(): data, eot, _ = await self._dsa.recv_extended(None) try: self._protocol.receive_data(data) except WebSocketProtocolError as e: raise ProtocolError from e for event in self._protocol.events(): if isinstance(event, TextMessage): return event.data elif isinstance(event, BytesMessage): return event.data elif isinstance(event, CloseConnection): self._remote_shutdown = True await self.close() return None elif isinstance(event, Ping): data_to_send = self._protocol.send(event.response()) async with self._write_error_catcher(): await self._dsa.sendall(data_to_send) elif isinstance(event, Pong): continue async def send_payload(self, buf: str | bytes) -> None: """Dispatch a buffer to remote.""" if self._dsa is None or self._response is None or self._police_officer is None: raise OSError("The HTTP extension is closed or uninitialized") async with self._police_officer.borrow(self._response): try: if isinstance(buf, str): data_to_send: bytes = self._protocol.send(TextMessage(buf)) else: data_to_send = self._protocol.send(BytesMessage(buf)) except WebSocketProtocolError as e: raise ProtocolError from e async with self._write_error_catcher(): await self._dsa.sendall(data_to_send) async def ping(self) -> None: if self._dsa is None or self._response is None or self._police_officer is None: raise OSError("The HTTP extension is closed or uninitialized") async with self._police_officer.borrow(self._response): try: data_to_send: bytes = self._protocol.send(Ping()) except WebSocketProtocolError as e: raise ProtocolError from e async with self._write_error_catcher(): await self._dsa.sendall(data_to_send) @staticmethod def supported_schemes() -> set[str]: return {"ws", "wss"} @staticmethod def scheme_to_http_scheme(scheme: str) -> str: return {"ws": "http", "wss": "https"}[scheme] class AsyncWebSocketExtensionFromMultiplexedHTTP(AsyncWebSocketExtensionFromHTTP): """ Plugin that support doing WebSocket over HTTP 2 and 3. This implement RFC8441. Beware that this isn't actually supported by much server around internet. """ @staticmethod def implementation() -> str: return "rfc8441" @staticmethod def supported_svn() -> set[HttpVersion]: return {HttpVersion.h11, HttpVersion.h2, HttpVersion.h3}