fix: handle WebSocket connections in audit log exception handlers

Author: sricharan-varanasi

src/apps/audit/fields.py

@@ -1,26 +1,27 @@
 from asgi_correlation_id.context import correlation_id
 from ddtrace.trace import tracer
-from fastapi import Request
 from fastapi.routing import APIRoute
 from starlette.exceptions import HTTPException as StarletteHTTPException
+from starlette.requests import HTTPConnection, Request
 
 from apps.shared.exception import BaseError
 
 from .enums import EventOutcome
 
 
-def http_audit_fields(request: Request, error: BaseError | StarletteHTTPException | None = None) -> dict:
-    """Audit fields derived from HTTP request/error."""
-    route = request.scope.get("route")
+def http_audit_fields(conn: HTTPConnection, error: BaseError | StarletteHTTPException | None = None) -> dict:
+    """Audit fields derived from an HTTP or WebSocket connection and error."""
+    route = conn.scope.get("route")
     span = tracer.current_span()
     fields = {
-        "client_ip": request.client and request.client.host,
+        "client_ip": conn.client and conn.client.host,
         "http_request_id": correlation_id.get(),
-        "http_request_method": request.method,
+        # WebSocket connections have no HTTP method.
+        "http_request_method": conn.method if isinstance(conn, Request) else None,
         "http_response_status_code": isinstance(route, APIRoute) and route.status_code or 200,
-        "url_path": request.url.path,
-        "url_query": request.url.query or None,
-        "user_agent": request.headers.get("user-agent"),
+        "url_path": conn.url.path,
+        "url_query": conn.url.query or None,
+        "user_agent": conn.headers.get("user-agent"),
         "trace_id": span and str(span.trace_id),
     }
     if error is not None:

src/apps/authentication/tests/test_auth.py

@@ -7,6 +7,8 @@
 import pytest
 from pytest_mock import MockerFixture
 from sqlalchemy.ext.asyncio import AsyncSession
+from starlette.types import Message
+from starlette.websockets import WebSocket
 
 from apps.audit import EventAction, EventOutcome
 from apps.authentication.domain.login import UserLoginRequest
@@ -25,6 +27,7 @@
 from apps.users.cruds.user import UsersCRUD
 from apps.users.domain import User, UserCreate, UserCreateRequest
 from config import settings
+from infrastructure.http.exceptions import session_token_invalid_error_handler
 
 TEST_PASSWORD = "Test12345!"
 
@@ -113,6 +116,50 @@ async def test_user_not_found(self, client: TestClient, mocker: MockerFixture):
         assert resp.status_code == http.HTTPStatus.UNAUTHORIZED
         assert resp.json()["result"][0]["message"] == SessionTokenInvalidError.message
 
+    async def test_ws_session_invalid_audit_event(self, mocker: MockerFixture):
+        """A `SessionTokenInvalidError` raised from a WebSocket connection (e.g. `/ws/alerts`) must
+
+        log the `user:session:invalid` audit event without crashing. `http_audit_fields` used to
+        read the HTTP-only `request.method`, which a `WebSocket` lacks (regression, M2-10698).
+        """
+        audit_log = mocker.patch("infrastructure.http.exceptions.log")
+        user_id = uuid.uuid4()
+
+        async def receive() -> Message:
+            return {"type": "websocket.connect"}
+
+        async def send(message: Message) -> None:
+            return None
+
+        websocket = WebSocket(
+            {
+                "type": "websocket",
+                "scheme": "ws",
+                "server": ("test.com", 80),
+                "path": "/ws/alerts",
+                "query_string": b"",
+                "root_path": "",
+                "headers": [(b"host", b"test.com"), (b"user-agent", b"pytest-ws-client")],
+                "client": ("10.1.2.3", 54321),
+            },
+            receive=receive,
+            send=send,
+        )
+        error = SessionTokenInvalidError(user_id=user_id)
+
+        resp = await session_token_invalid_error_handler(websocket, error)
+
+        audit_log.assert_awaited_once()
+        event = audit_log.call_args[0][0]
+        assert event.event_action == EventAction.USER_SESSION_INVALID
+        assert event.event_outcome == EventOutcome.FAILURE
+        assert event.user_id == user_id
+        assert event.http_request_method is None  # WebSocket has no HTTP method
+        assert event.url_path == "/ws/alerts"
+        assert event.client_ip == "10.1.2.3"
+        assert event.user_agent == "pytest-ws-client"
+        assert resp.status_code == http.HTTPStatus.UNAUTHORIZED
+
     async def test_delete_access_token(self, client: TestClient, user: User, mocker: MockerFixture):
         audit_log = mocker.patch("apps.authentication.api.auth.log")
         client.login(user)

src/infrastructure/http/exceptions.py

@@ -5,7 +5,7 @@
 from fastapi.exceptions import RequestValidationError
 from starlette import status
 from starlette.exceptions import HTTPException as StarletteHTTPException
-from starlette.requests import Request
+from starlette.requests import HTTPConnection, Request
 from starlette.responses import JSONResponse, Response
 
 from apps.audit import AuditEvent, EventAction, http_audit_fields, log
@@ -21,7 +21,7 @@ def _set_trace_exception(exc: Exception) -> None:
         span.set_exc_info(type(exc), exc, exc.__traceback__)
 
 
-def custom_base_errors_handler(_: Request, error: BaseError) -> JSONResponse:
+def custom_base_errors_handler(_: HTTPConnection, error: BaseError) -> JSONResponse:
     """This function is called if the BaseError was raised."""
 
     logger.error(error.error, exc_info=error)
@@ -53,8 +53,12 @@ def custom_base_errors_handler(_: Request, error: BaseError) -> JSONResponse:
     )
 
 
-async def session_token_invalid_error_handler(request: Request, error: SessionTokenInvalidError) -> JSONResponse:
-    """user:session:invalid audit event on 401 from invalid session token in `Authorization` header."""
+async def session_token_invalid_error_handler(request: HTTPConnection, error: SessionTokenInvalidError) -> JSONResponse:
+    """user:session:invalid audit event on 401 from invalid session token.
+
+    ``request`` may be a ``WebSocket`` (e.g. the `/ws/alerts` handshake) as well as an HTTP
+    ``Request``, so only fields common to both connection types may be used here.
+    """
     await log(
         AuditEvent(
             event_action=EventAction.USER_SESSION_INVALID,