def test_upgrade_request() -> None:
server = WSConnection(SERVER)
server.initiate_upgrade_connection(
[
(b"Host", b"localhost"),
(b"Connection", b"Keep-Alive, Upgrade"),
(b"Upgrade", b"WebSocket"),
(b"Sec-WebSocket-Version", b"13"),
(b"Sec-WebSocket-Key", generate_nonce()),
(b"X-Foo", b"bar"),
],
"/",
)
event = next(server.events())
event = cast(Request, event)
assert event.extensions == []
assert event.host == "localhost"
assert event.subprotocols == []
assert event.target == "/"
headers = normed_header_dict(event.extra_headers)
assert b"host" not in headers
assert b"sec-websocket-extensions" not in headers
assert b"sec-websocket-protocol" not in headers
assert headers[b"connection"] == b"Keep-Alive, Upgrade"
assert headers[b"sec-websocket-version"] == b"13"
assert headers[b"upgrade"] == b"WebSocket"
assert headers[b"x-foo"] == b"bar"
async def init_for_server( # type: ignore[misc]
cls: Type["Transport"], stream: Stream, upgrade_request: Optional[H11Request] = None
) -> "Transport":
ws = WSConnection(ConnectionType.SERVER)
if upgrade_request:
ws.initiate_upgrade_connection(
headers=upgrade_request.headers, path=upgrade_request.target
)
transport = cls(stream, ws)
# Wait for client to init WebSocket handshake
event: Union[str, Event] = "Websocket handshake timeout"
with trio.move_on_after(WEBSOCKET_HANDSHAKE_TIMEOUT):
event = await transport._next_ws_event()
if isinstance(event, Request):
transport.logger.debug("Accepting WebSocket upgrade")
await transport._net_send(AcceptConnection())
return transport
transport.logger.warning("Unexpected event during WebSocket handshake", ws_event=event)
raise TransportError(f"Unexpected event during WebSocket handshake: {event}")
class WebsocketServer(HTTPServer, WebsocketMixin):
def __init__(
self,
app: Type[ASGIFramework],
loop: asyncio.AbstractEventLoop,
config: Config,
transport: asyncio.BaseTransport,
*,
upgrade_request: Optional[h11.Request] = None,
) -> None:
super().__init__(loop, config, transport, "wsproto")
self.stop_keep_alive_timeout()
self.app = app
self.connection = WSConnection(ConnectionType.SERVER)
self.app_queue: asyncio.Queue = asyncio.Queue()
self.response: Optional[dict] = None
self.scope: Optional[dict] = None
self.state = ASGIWebsocketState.HANDSHAKE
self.task: Optional[asyncio.Future] = None
self.buffer = WebsocketBuffer(self.config.websocket_max_message_size)
if upgrade_request is not None:
self.connection.initiate_upgrade_connection(
upgrade_request.headers, upgrade_request.target)
self.handle_events()
def connection_lost(self, error: Optional[Exception]) -> None:
if error is not None:
self.app_queue.put_nowait({"type": "websocket.disconnect"})
def eof_received(self) -> bool:
self.data_received(None)
return True
def data_received(self, data: Optional[bytes]) -> None:
self.connection.receive_data(data)
self.handle_events()
def handle_events(self) -> None:
for event in self.connection.events():
if isinstance(event, Request):
self.task = self.loop.create_task(self.handle_websocket(event))
self.task.add_done_callback(self.maybe_close)
elif isinstance(event, Message):
try:
self.buffer.extend(event)
except FrameTooLarge:
self.write(
self.connection.send(
CloseConnection(code=CloseReason.MESSAGE_TOO_BIG)))
self.app_queue.put_nowait({"type": "websocket.disconnect"})
self.close()
break
if event.message_finished:
self.app_queue.put_nowait(self.buffer.to_message())
self.buffer.clear()
elif isinstance(event, Ping):
self.write(self.connection.send(event.response()))
elif isinstance(event, CloseConnection):
if self.connection.state == ConnectionState.REMOTE_CLOSING:
self.write(self.connection.send(event.response()))
self.app_queue.put_nowait({"type": "websocket.disconnect"})
self.close()
break
def maybe_close(self, future: asyncio.Future) -> None:
# Close the connection iff a HTTP response was sent
if self.state == ASGIWebsocketState.HTTPCLOSED:
self.close()
async def asend(self, event: Event) -> None:
await self.drain()
self.write(self.connection.send(event))
async def asgi_put(self, message: dict) -> None:
await self.app_queue.put(message)
async def asgi_receive(self) -> dict:
"""Called by the ASGI instance to receive a message."""
return await self.app_queue.get()
@property
def scheme(self) -> str:
return "wss" if self.ssl_info is not None else "ws"
class WebsocketServer(HTTPServer, WebsocketMixin):
def __init__(
self,
app: ASGIFramework,
config: Config,
stream: trio.abc.Stream,
*,
upgrade_request: Optional[h11.Request] = None,
) -> None:
super().__init__(stream, "wsproto")
self.app = app
self.config = config
self.connection = WSConnection(ConnectionType.SERVER)
self.response: Optional[dict] = None
self.scope: Optional[dict] = None
self.send_lock = trio.Lock()
self.state = ASGIWebsocketState.HANDSHAKE
self.buffer = WebsocketBuffer(self.config.websocket_max_message_size)
self.app_send_channel, self.app_receive_channel = trio.open_memory_channel(
10)
if upgrade_request is not None:
self.connection.initiate_upgrade_connection(
upgrade_request.headers, upgrade_request.target)
async def handle_connection(self) -> None:
try:
request = await self.read_request()
async with trio.open_nursery() as nursery:
nursery.start_soon(self.read_messages)
await self.handle_websocket(request)
if self.state == ASGIWebsocketState.HTTPCLOSED:
raise MustCloseError()
except (trio.BrokenResourceError, trio.ClosedResourceError):
await self.asgi_put({"type": "websocket.disconnect"})
except MustCloseError:
pass
finally:
await self.aclose()
async def read_request(self) -> Request:
for event in self.connection.events():
if isinstance(event, Request):
return event
async def read_messages(self) -> None:
while True:
data = await self.stream.receive_some(MAX_RECV)
if data == b"":
data = None # wsproto expects None rather than b"" for EOF
self.connection.receive_data(data)
for event in self.connection.events():
if isinstance(event, Message):
try:
self.buffer.extend(event)
except FrameTooLarge:
await self.asend(
CloseConnection(code=CloseReason.MESSAGE_TOO_BIG))
await self.asgi_put({"type": "websocket.disconnect"})
raise MustCloseError()
if event.message_finished:
await self.asgi_put(self.buffer.to_message())
self.buffer.clear()
elif isinstance(event, Ping):
await self.asend(event.response())
elif isinstance(event, CloseConnection):
if self.connection.state == ConnectionState.REMOTE_CLOSING:
await self.asend(event.response())
await self.asgi_put({"type": "websocket.disconnect"})
raise MustCloseError()
async def asend(self, event: Event) -> None:
async with self.send_lock:
await self.stream.send_all(self.connection.send(event))
async def asgi_put(self, message: dict) -> None:
await self.app_send_channel.send(message)
async def asgi_receive(self) -> dict:
return await self.app_receive_channel.receive()
@property
def scheme(self) -> str:
return "wss" if self._is_ssl else "ws"
