diff --git a/packages/kaya-core/src/kaya/core/__init__.py b/packages/kaya-core/src/kaya/core/__init__.py index 76b56e6..ba76117 100644 --- a/packages/kaya-core/src/kaya/core/__init__.py +++ b/packages/kaya-core/src/kaya/core/__init__.py @@ -4,7 +4,7 @@ from ._http_context import HttpContext from ._mixin import KayaMixin from ._tree import Tree, PathIterator from ._path_handler import PathHandler, Matches -from ._websocket import WebSocket, WebSocketMessage +from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage __all__ = [ @@ -18,5 +18,6 @@ __all__ = [ 'Matches', 'PathIterator', 'WebSocket', + 'WebSocketDisconnect', 'WebSocketMessage' ] diff --git a/packages/kaya-core/src/kaya/core/_app.py b/packages/kaya-core/src/kaya/core/_app.py index baf55b6..2d76b19 100644 --- a/packages/kaya-core/src/kaya/core/_app.py +++ b/packages/kaya-core/src/kaya/core/_app.py @@ -9,7 +9,7 @@ from ._http_method import HttpMethod from ._mixin import KayaMixin from ._path_handler import Context from ._types import StrOrStrings -from ._websocket import WebSocket, WebSocketMessage +from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage from ._asgi import AsgiContext, AsgiWebSocket from ._tree import Tree from ._types.asgi import LifespanScope, HTTPScope as ASGIHTTPScope, WebSocketScope as ASGIWebSocketScope @@ -164,6 +164,12 @@ class KayaApp(AbstractKayaApp): await handler.handle_request(ws, captured) else: await ws.close(1000) + except WebSocketDisconnect: + # The peer went away while the handler was using the websocket + # (e.g. sending on a vanished transport). This is normal + # termination, not an application error, so it is not recorded + # on the socket nor re-raised. + pass except Exception as exc: ws.exception = exc raise diff --git a/packages/kaya-core/src/kaya/core/_asgi.py b/packages/kaya-core/src/kaya/core/_asgi.py index 0612612..156e858 100644 --- a/packages/kaya-core/src/kaya/core/_asgi.py +++ b/packages/kaya-core/src/kaya/core/_asgi.py @@ -17,7 +17,7 @@ from pwo import Maybe from pathlib import Path from ._http_method import HttpMethod from ._http_context import HttpContext -from ._websocket import WebSocket, WebSocketMessage +from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage from ._types import StrOrStrings from ._types.asgi import HTTPScope, WebSocketScope @@ -179,7 +179,16 @@ class AsgiWebSocket(WebSocket): # Emit a list rather than a tuple: wsproto's AcceptConnection (used # by httpx_ws and others) requires list concatenation. message['headers'] = list(encode_headers(headers)) - await self._send(message) + await self._send_message(message) + + async def _send_message(self, message: Mapping[str, Any]) -> None: + try: + await self._send(message) + except OSError as exc: + # ASGI servers (e.g. uvicorn) report a vanished peer with an + # OSError subclass; surface it as a kaya-level disconnect so + # handlers do not need to know about server internals. + raise WebSocketDisconnect(code=1006) from exc async def receive(self) -> WebSocketMessage: message = await self._receive() @@ -200,13 +209,13 @@ class AsgiWebSocket(WebSocket): return WebSocketMessage(kind='close') async def send_text(self, data: str) -> None: - await self._send({'type': 'websocket.send', 'text': data}) + await self._send_message({'type': 'websocket.send', 'text': data}) async def send_bytes(self, data: bytes) -> None: - await self._send({'type': 'websocket.send', 'bytes': data}) + await self._send_message({'type': 'websocket.send', 'bytes': data}) async def close(self, code: int = 1000) -> None: - await self._send({'type': 'websocket.close', 'code': code}) + await self._send_message({'type': 'websocket.close', 'code': code}) async def __anext__(self) -> WebSocketMessage: message = await self.receive() diff --git a/packages/kaya-core/src/kaya/core/_websocket.py b/packages/kaya-core/src/kaya/core/_websocket.py index 2b223a1..9ec1147 100644 --- a/packages/kaya-core/src/kaya/core/_websocket.py +++ b/packages/kaya-core/src/kaya/core/_websocket.py @@ -13,6 +13,41 @@ class WebSocketMessage: data: WebSocketData = None +class WebSocketDisconnect(Exception): + """Raised when using a websocket whose peer has disconnected. + + Concrete websocket implementations raise this from ``send_text()`` and + ``send_bytes()`` when the transport reports that the peer is gone (for + example when a browser tab is closed without a close handshake), so + handlers can catch it and stop producing output: + + .. code-block:: python + + try: + await ws.send_text('hello') + except WebSocketDisconnect: + pass + + Reading from a disconnected websocket is not an error: ``receive()`` + reports a ``close`` message and iterating with ``async for`` simply ends, + which is what most handlers should rely on instead. + + ``code`` is a websocket close code; ``1006`` means the connection was + closed abnormally, without a close handshake. + """ + + code: int + reason: str + + def __init__(self, code: int = 1006, reason: str = '') -> None: + self.code = code + self.reason = reason + message = f'WebSocket disconnected (code={code})' + if reason: + message = f'{message}: {reason}' + super().__init__(message) + + class WebSocket(ABC): path: str query_string: str diff --git a/packages/kaya-core/tests/test_websocket.py b/packages/kaya-core/tests/test_websocket.py index 14120f3..36b090e 100644 --- a/packages/kaya-core/tests/test_websocket.py +++ b/packages/kaya-core/tests/test_websocket.py @@ -1,9 +1,9 @@ import unittest import httpx from pwo import async_test -from httpx_ws import aconnect_ws, WebSocketDisconnect +from httpx_ws import aconnect_ws, WebSocketDisconnect as HttpxWebSocketDisconnect from httpx_ws.transport import ASGIWebSocketTransport -from kaya.core import KayaApp, WebSocket +from kaya.core import KayaApp, WebSocket, WebSocketDisconnect, WebSocketMessage from kaya.core._asgi import AsgiWebSocket @@ -66,7 +66,7 @@ class WebSocketTest(unittest.TestCase): transport = ASGIWebSocketTransport(app=self.app) async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: - with self.assertRaises(WebSocketDisconnect) as cm: + with self.assertRaises(HttpxWebSocketDisconnect) as cm: async with aconnect_ws("/unknown", client) as ws: pass @@ -151,3 +151,113 @@ class WebSocketTest(unittest.TestCase): } ws = AsgiWebSocket(scope, receive, send) self.assertEqual('ws', ws.scheme) + + +class FakeWebSocket(WebSocket): + """Minimal in-memory websocket for exercising ``KayaApp`` directly.""" + + def __init__(self, path: str) -> None: + self.path = path + self.query_string = '' + self.scheme = 'ws' + self.headers: dict[str, list[str]] = {} + self.client = ('127.0.0.1', 12345) + self.server = ('127.0.0.1', 80) + self.closed: list[int] = [] + + async def accept(self, headers=None) -> None: + pass + + async def receive(self) -> WebSocketMessage: + return WebSocketMessage(kind='close') + + async def send_text(self, data: str) -> None: + pass + + async def send_bytes(self, data: bytes) -> None: + pass + + async def close(self, code: int = 1000) -> None: + self.closed.append(code) + + async def __anext__(self) -> WebSocketMessage: + raise StopAsyncIteration + + +class WebSocketDisconnectTest(unittest.TestCase): + + def setUp(self): + self.app = KayaApp() + + @self.app.websocket('/disconnect') + async def disconnect(ws: WebSocket) -> None: + await ws.accept() + await ws.send_text('before') + raise WebSocketDisconnect() + + @self.app.websocket('/failure') + async def failure(ws: WebSocket) -> None: + await ws.accept() + raise RuntimeError('boom') + + @async_test + async def test_disconnect_terminates_quietly(self): + after_hook_calls = [] + + async def after(ws: WebSocket) -> None: + after_hook_calls.append(ws) + + self.app.add_after_websocket_hook(after) + ws = FakeWebSocket('/disconnect') + await self.app.handle_websocket(ws) + self.assertEqual([ws], after_hook_calls) + self.assertIsNone(ws.exception) + + @async_test + async def test_other_errors_are_recorded_and_reraised(self): + ws = FakeWebSocket('/failure') + with self.assertRaises(RuntimeError): + await self.app.handle_websocket(ws) + self.assertIsInstance(ws.exception, RuntimeError) + + @async_test + async def test_asgi_send_after_disconnect_raises_websocket_disconnect(self): + async def send(message): + raise OSError('client gone') + + async def receive(): + return {'type': 'websocket.disconnect', 'code': 1006} + + ws = AsgiWebSocket(self._scope(), receive, send) + with self.assertRaises(WebSocketDisconnect) as ctx: + await ws.send_text('hello') + self.assertEqual(1006, ctx.exception.code) + with self.assertRaises(WebSocketDisconnect): + await ws.send_bytes(b'hello') + with self.assertRaises(WebSocketDisconnect): + await ws.close() + + @async_test + async def test_asgi_accept_after_disconnect_raises_websocket_disconnect(self): + async def send(message): + raise OSError('client gone') + + async def receive(): + return {'type': 'websocket.connect'} + + ws = AsgiWebSocket(self._scope(), receive, send) + with self.assertRaises(WebSocketDisconnect) as ctx: + await ws.accept() + self.assertEqual(1006, ctx.exception.code) + + @staticmethod + def _scope() -> dict: + return { + 'type': 'websocket', + 'path': '/echo', + 'query_string': b'', + 'scheme': 'ws', + 'client': ('127.0.0.1', 12345), + 'server': ('127.0.0.1', 80), + 'headers': [], + } diff --git a/packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py b/packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py index 774cbe6..4806388 100644 --- a/packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py +++ b/packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py @@ -17,13 +17,21 @@ from typing import ( from granian._granian import ( # type: ignore[attr-defined] RSGIHTTPProtocol, RSGIHTTPScope, + RSGIProtocolClosed, RSGIWebsocketProtocol, RSGIWebsocketScope, RSGIWebsocketTransport, ) from pwo import Maybe -from kaya.core import AbstractKayaApp, HttpContext, HttpMethod, WebSocket, WebSocketMessage +from kaya.core import ( + AbstractKayaApp, + HttpContext, + HttpMethod, + WebSocket, + WebSocketDisconnect, + WebSocketMessage, +) from kaya.core._types import StrOrStrings @@ -109,6 +117,7 @@ class RsgiContext(HttpContext): class RsgiWebSocket(WebSocket): _protocol: RSGIWebsocketProtocol _transport: Optional[RSGIWebsocketTransport] + _closed: bool scheme: str path: str query_string: str @@ -124,6 +133,7 @@ class RsgiWebSocket(WebSocket): ) self._protocol = protocol self._transport = None + self._closed = False self.scheme = scope.scheme self.path = scope.path self.query_string = scope.query_string @@ -144,13 +154,28 @@ class RsgiWebSocket(WebSocket): async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: # RSGI's websocket accept() takes no arguments (https://github.com/emmett-framework/granian/blob/master/docs/spec/RSGI.md), # so handshake headers (e.g. Set-Cookie) cannot be sent on RSGI; they are silently ignored here. - self._transport = await self._protocol.accept() + try: + self._transport = await self._protocol.accept() + except RSGIProtocolClosed as exc: + self._closed = True + raise WebSocketDisconnect(code=1006) from exc async def receive(self) -> WebSocketMessage: if self._transport is None: raise RuntimeError('WebSocket connection has not been accepted yet') - message = await self._transport.receive() + if self._closed: + return WebSocketMessage(kind='close') + try: + message = await self._transport.receive() + except RSGIProtocolClosed: + # Granian raises this when the peer is gone without a close + # handshake (browser tab closed, network dropped, ...). Surface it + # as a close message so receive loops terminate normally instead + # of seeing an exception. + self._closed = True + return WebSocketMessage(kind='close') if message.kind == 0: + self._closed = True return WebSocketMessage(kind='close') elif message.kind == 1: return WebSocketMessage(kind='binary', data=message.data) @@ -162,15 +187,32 @@ class RsgiWebSocket(WebSocket): async def send_text(self, data: str) -> None: if self._transport is None: raise RuntimeError('WebSocket connection has not been accepted yet') - await self._transport.send_str(data) + if self._closed: + raise WebSocketDisconnect(code=1006) + try: + await self._transport.send_str(data) + except RSGIProtocolClosed as exc: + self._closed = True + raise WebSocketDisconnect(code=1006) from exc async def send_bytes(self, data: bytes) -> None: if self._transport is None: raise RuntimeError('WebSocket connection has not been accepted yet') - await self._transport.send_bytes(data) + if self._closed: + raise WebSocketDisconnect(code=1006) + try: + await self._transport.send_bytes(data) + except RSGIProtocolClosed as exc: + self._closed = True + raise WebSocketDisconnect(code=1006) from exc async def close(self, code: int = 1000) -> None: - self._protocol.close(code) + if self._closed: + return + try: + self._protocol.close(code) + except RSGIProtocolClosed: + self._closed = True async def __anext__(self) -> WebSocketMessage: message = await self.receive() diff --git a/packages/kaya-rsgi/tests/test_rsgi.py b/packages/kaya-rsgi/tests/test_rsgi.py index 096850b..26b2a36 100644 --- a/packages/kaya-rsgi/tests/test_rsgi.py +++ b/packages/kaya-rsgi/tests/test_rsgi.py @@ -1,22 +1,150 @@ import unittest +from types import SimpleNamespace +from typing import Any + +from granian.rsgi import ProtocolClosed +from pwo import async_test + +from kaya.core import WebSocketDisconnect from kaya.rsgi import RsgiWebSocket +class FakeScope: + scheme = 'ws' + path = '/ws' + query_string = '' + headers: dict[str, list[str]] = {} + client = '127.0.0.1:12345' + server = '127.0.0.1:80' + + +class FakeTransport: + """A transport whose peer vanished without a close handshake.""" + + def __init__(self) -> None: + self.receive_calls = 0 + self.sent: list[Any] = [] + + async def receive(self) -> Any: + self.receive_calls += 1 + raise ProtocolClosed() + + async def send_str(self, data: str) -> None: + self.sent.append(data) + raise ProtocolClosed() + + async def send_bytes(self, data: bytes) -> None: + self.sent.append(data) + raise ProtocolClosed() + + +class CleanCloseTransport(FakeTransport): + """A transport that reports a close message (kind 0).""" + + async def receive(self) -> Any: + self.receive_calls += 1 + return SimpleNamespace(kind=0, data=None) + + +class FakeProtocol: + def __init__(self, transport: FakeTransport) -> None: + self._transport = transport + self.closed: list[int] = [] + + async def accept(self) -> FakeTransport: + return self._transport + + def close(self, code: int) -> None: + self.closed.append(code) + + +class ClosedOnAcceptProtocol(FakeProtocol): + """A protocol whose peer vanished during the handshake.""" + + async def accept(self) -> FakeTransport: + raise ProtocolClosed() + + class RsgiWebSocketTest(unittest.TestCase): def test_misconfigured_granian(self): - class FakeScope: - scheme = 'ws' - path = '/ws' - query_string = '' - headers = {} - client = '127.0.0.1:12345' - server = '127.0.0.1:80' - - class FakeProtocol: + class FakeProtocolWithoutWebsockets: pass with self.assertRaises(RuntimeError) as ctx: - RsgiWebSocket(FakeScope(), FakeProtocol()) # type: ignore[arg-type] + RsgiWebSocket(FakeScope(), FakeProtocolWithoutWebsockets()) # type: ignore[arg-type] self.assertIn('Granian was not configured for websockets', str(ctx.exception)) + + async def _socket(self, transport: FakeTransport) -> tuple[RsgiWebSocket, FakeProtocol]: + protocol = FakeProtocol(transport) + ws = RsgiWebSocket(FakeScope(), protocol) # type: ignore[arg-type] + await ws.accept() + return ws, protocol + + @async_test + async def test_receive_returns_close_when_transport_is_gone(self): + ws, _ = await self._socket(FakeTransport()) + message = await ws.receive() + self.assertEqual('close', message.kind) + + @async_test + async def test_iteration_ends_when_transport_is_gone(self): + ws, _ = await self._socket(FakeTransport()) + messages = [message async for message in ws] + self.assertEqual([], messages) + + @async_test + async def test_receive_remembers_the_close(self): + transport = FakeTransport() + ws, _ = await self._socket(transport) + self.assertEqual('close', (await ws.receive()).kind) + self.assertEqual('close', (await ws.receive()).kind) + self.assertEqual(1, transport.receive_calls) + + @async_test + async def test_send_text_raises_disconnect_when_transport_is_gone(self): + ws, _ = await self._socket(FakeTransport()) + with self.assertRaises(WebSocketDisconnect) as ctx: + await ws.send_text('hello') + self.assertEqual(1006, ctx.exception.code) + + @async_test + async def test_send_bytes_raises_disconnect_when_transport_is_gone(self): + ws, _ = await self._socket(FakeTransport()) + with self.assertRaises(WebSocketDisconnect) as ctx: + await ws.send_bytes(b'hello') + self.assertEqual(1006, ctx.exception.code) + + @async_test + async def test_sends_after_disconnect_do_not_touch_the_transport(self): + transport = FakeTransport() + ws, _ = await self._socket(transport) + with self.assertRaises(WebSocketDisconnect): + await ws.send_text('first') + with self.assertRaises(WebSocketDisconnect): + await ws.send_text('second') + self.assertEqual(['first'], transport.sent) + + @async_test + async def test_send_after_a_clean_close_raises_disconnect(self): + transport = CleanCloseTransport() + ws, _ = await self._socket(transport) + self.assertEqual('close', (await ws.receive()).kind) + with self.assertRaises(WebSocketDisconnect): + await ws.send_text('hello') + self.assertEqual([], transport.sent) + + @async_test + async def test_accept_raises_disconnect_when_transport_is_gone(self): + ws = RsgiWebSocket(FakeScope(), ClosedOnAcceptProtocol(FakeTransport())) # type: ignore[arg-type] + with self.assertRaises(WebSocketDisconnect) as ctx: + await ws.accept() + self.assertEqual(1006, ctx.exception.code) + + @async_test + async def test_close_is_noop_once_the_transport_is_gone(self): + ws, protocol = await self._socket(FakeTransport()) + await ws.receive() + await ws.close(1000) + self.assertEqual([], protocol.closed)