Repository: Neoteroi/BlackSheep
Pull request #446 (2023-12-14 21:44:50)

Fix WS exception handling
@RobertoPrevato After some considerations, I decided to try this approach. Please, let me know, what you think about it. Is it good enough to you?

1. No matching WebSocket path? Close the connection, the ASGI server will respond with 403.
2. An exception raised during handshake? Close the connection, the ASGI server will respond with 403.
3. An exception raised after handshake? Close the connection with code 1011 (which seems to be equivalent to Internal Server Error, see https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1). We return the error reason too, but we take care to truncate it to prevent `control frame too long` errors.

We had a discussion about responding with a more precise status code if the connection was closed before accept, but it seems to be impossible in the ASGI framework. We are not in control of the initial HTTP request, so to speak. We cannot make the protocol server reject the connection with a status code different to 403, and if we try to pass an HTTP status code as a WebSocket close code, it just won't work. Clients and servers like Uvicorn/Daphne will only see this as an invalid WebSocket code.

RobertoPrevato: Hi @Klavionik
Thank You for helping with this! I appreciate very much your contributions. I am too tired now to review, but I can do it tomorrow. :)
RobertoPrevato: @Klavionik I reviewed and tested your changes. I think what you did is good, and I thank you for investigating and sharing how various ASGI servers handle this case. From my point of view the ASGI specification here doesn't make any sense.
Essentially if a WebSocket handler needs to return a Bad Request response, or an Unauthorized response with error details _before the handshake_ even completes (like Access Token expired, input parameter is wrong, etc) - it is not possible to provide a user friendly error message to the client. Always a dull 403 Forbidden error without details.

At the very least, I added a logging.exception line to your branch, to get some information about the Exception in the application logs:

![image](https://github.com/Neoteroi/BlackSheep/assets/2576032/ebd7db06-f600-4fd5-902e-524a0a76f660)


Patch:
diff --git a/blacksheep/server/application.py b/blacksheep/server/application.py
index de2e945..54cd4b6 100644
--- a/blacksheep/server/application.py
+++ b/blacksheep/server/application.py
@@ -32,7 +32,6 @@ from blacksheep.baseapp import BaseApplication, handle_not_found
 from blacksheep.common import extend
 from blacksheep.common.files.asyncfs import FilesHandler
 from blacksheep.contents import ASGIContent
-from blacksheep.exceptions import HTTPException
 from blacksheep.messages import Request, Response
 from blacksheep.middlewares import get_middlewares_chain
 from blacksheep.scribe import send_asgi_response
@@ -65,7 +64,7 @@ from blacksheep.server.routing import (
 )
 from blacksheep.server.routing import router as default_router
 from blacksheep.server.routing import validate_router
-from blacksheep.server.websocket import WebSocket
+from blacksheep.server.websocket import WebSocket, format_reason
 from blacksheep.sessions import SessionMiddleware, SessionSerializer
 from blacksheep.settings.di import di_settings
 from blacksheep.utils import ensure_bytes, join_fragments
@@ -739,31 +738,23 @@ class Application(BaseApplication):
             RouteMethod.GET_WS, scope["path"]
         )
 
-        if route:
-            ws.route_values = route.values
-            try:
-                return await route.handler(ws)
-            except UnauthorizedError as unauthorized_error:
-                # If the WebSocket connection was not accepted yet, we close the
-                # connection with an HTTP Status Code, otherwise we close the connection
-                # with a WebSocket status code
-                if ws.accepted:
-                    # Use a WebSocket error code, not an HTTP error code
-                    await ws.close(1005, "Unauthorized")
-                else:
-                    # Still in handshake phase, we close with an HTTP Status Code
-                    # https://asgi.readthedocs.io/en/latest/specs/www.html#close-send-event
-                    await ws.close(403, str(unauthorized_error))
-            except HTTPException as http_exception:
-                # Same like above
-                if ws.accepted:
-                    # Use a WebSocket error code, not an HTTP error code
-                    await ws.close(1005, str(http_exception))
-                else:
-                    # Still in handshake phase, we close with an HTTP Status Code
-                    # https://asgi.readthedocs.io/en/latest/specs/www.html#close-send-event
-                    await ws.close(http_exception.status, str(http_exception))
-        await ws.close()
+        if route is None:
+            return await ws.close()
+
+        ws.route_values = route.values
+
+        try:
+            return await route.handler(ws)
+        except Exception as exc:
+            logging.exception("Exception while handling WebSocket")
+            # If WebSocket connection accepted, close
+            # the connection using WebSocket Internal error code.
+            if ws.accepted:
+                return await ws.close(1011, reason=format_reason(str(exc)))
+
+            # Otherwise, just close the connection, the ASGI server
+            # will anyway respond 403 to the client.
+            return await ws.close()
 
     async def _handle_http(self, scope, receive, send):
         assert scope["type"] == "http"
diff --git a/blacksheep/server/websocket.py b/blacksheep/server/websocket.py
index 7fa12a0..d0d3deb 100644
--- a/blacksheep/server/websocket.py
+++ b/blacksheep/server/websocket.py
@@ -6,6 +6,8 @@ from blacksheep.messages import Request
 from blacksheep.server.asgi import get_full_path
 from blacksheep.settings.json import json_settings
 
+MAX_REASON_SIZE = 123
+
 
 class WebSocketState(Enum):
     CONNECTING = 0
@@ -180,3 +182,18 @@ class WebSocket(Request):
 
     async def close(self, code: int = 1000, reason: Optional[str] = None) -> None:
         await self._send({"type": "websocket.close", "code": code, "reason": reason})
+
+
+def format_reason(reason: str) -> str:
+    """
+    Ensures that the close reason is no longer than the max reason size.
+    (https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#reason)
+    """
+    reason_bytes = reason.encode()
+
+    if len(reason_bytes) <= MAX_REASON_SIZE:
+        return reason
+
+    ellipsis_ = b"..."
+    truncated_reason = reason_bytes[: MAX_REASON_SIZE - len(ellipsis_)] + ellipsis_
+    return truncated_reason.decode()


Test patch:
diff --git a/itests/app_2.py b/itests/app_2.py
index 91f7e5d..723cbb8 100644
--- a/itests/app_2.py
+++ b/itests/app_2.py
@@ -17,7 +17,6 @@ from openapidocs.v3 import Schema
 from pydantic import BaseModel
 
 from blacksheep import Response, TextContent, WebSocket
-from blacksheep.exceptions import BadRequest
 from blacksheep.server import Application
 from blacksheep.server.authentication import AuthenticationHandler
 from blacksheep.server.authorization import Policy, Requirement, auth
@@ -150,9 +149,15 @@ async def echo_text_admin_users(websocket: WebSocket):
         await websocket.send_text(msg)
 
 
-@app_2.router.ws("/websocket-echo-text-http-exp")
+@app_2.router.ws("/websocket-error-before-accept")
 async def echo_text_http_exp(websocket: WebSocket):
-    raise BadRequest("Example")
+    raise RuntimeError("Error before accept")
+
+
+@app_2.router.ws("/websocket-server-error")
+async def websocket_server_error(websocket: WebSocket):
+    await websocket.accept()
+    raise RuntimeError("Server error")
 
 
 @auth("authenticated")
diff --git a/itests/test_server.py b/itests/test_server.py
index 905d94a..c037556 100644
--- a/itests/test_server.py
+++ b/itests/test_server.py
@@ -8,7 +8,7 @@ from uuid import uuid4
 import pytest
 import websockets
 import yaml
-from websockets.exceptions import InvalidStatusCode
+from websockets.exceptions import ConnectionClosedError, InvalidStatusCode
 
 from .client_fixtures import get_static_path
 from .server_fixtures import *  # NoQA
@@ -793,7 +793,7 @@ async def test_websocket(server_host, server_port_4, route, data):
     "route",
     [
         "websocket-echo-text-auth",
-        "websocket-echo-text-http-exp",
+        "websocket-error-before-accept",
     ],
 )
 async def test_websocket_auth(server_host, server_port_2, route):
@@ -805,3 +805,16 @@ async def test_websocket_auth(server_host, server_port_2, route):
 
     assert error.value.status_code == 403
     assert "server rejected" in str(error.value)
+
+
+@pytest.mark.asyncio
+async def test_websocket_server_error(server_host, server_port_2):
+    uri = f"ws://{server_host}:{server_port_2}/websocket-server-error"
+
+    with pytest.raises(ConnectionClosedError) as error:
+        async with websockets.connect(uri) as ws:
+            async for _message in ws:
+                pass
+
+    assert error.value.code == 1011
+    assert error.value.reason == "Server error"
diff --git a/tests/test_websocket.py b/tests/test_websocket.py
index 4b95a38..05a3343 100644
--- a/tests/test_websocket.py
+++ b/tests/test_websocket.py
@@ -7,6 +7,7 @@ from blacksheep.server.websocket import (
     WebSocket,
     WebSocketDisconnectError,
     WebSocketState,
+    format_reason,
 )
 from blacksheep.testing.messages import MockReceive, MockSend
 from tests.utils.application import FakeApplication
@@ -442,3 +443,21 @@ async def test_application_websocket_binding_by_type_annotation():
         mock_receive,
         mock_send,
     )
+
+
+LONG_REASON = "WRY" * 41
+QIN = "秦"  # Qyn dynasty in Chinese, 3 bytes.
+TOO_LONG_REASON = QIN * 42
+TOO_LONG_REASON_TRUNC = TOO_LONG_REASON[:40] + "..."
+
+
+@pytest.mark.parametrize(
+    "inp,out",
+    [
+        ("Short reason", "Short reason"),
+        (LONG_REASON, LONG_REASON),
+        (TOO_LONG_REASON, TOO_LONG_REASON_TRUNC),
+    ],
+)
+def test_format_reason(inp, out):
+    assert format_reason(inp) == out