refactor: Rewrite server middleware with websocket support

The builtin fastapi support for middleware only supports http/rest
requests. To enable the same middleware for websockets, a new
implementation of the starlette middleware is required and it makes
sense to use the same implementation for both protocols.

For rest requests, there should not be any change in behaviour from this
change.

Author: tpoliaw

src/blueapi/service/main.py

@@ -25,18 +25,19 @@
     get_tracer,
     start_as_current_span,
 )
-from opentelemetry.context import attach
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
-from opentelemetry.propagate import get_global_textmap
 from opentelemetry.trace import get_tracer_provider
 from pydantic import ValidationError
 from pydantic.json_schema import SkipJsonSchema
 from starlette.responses import JSONResponse
 from super_state_machine.errors import TransitionError
 
-from blueapi import __version__
 from blueapi.config import ApplicationConfig, OIDCConfig, Tag
 from blueapi.service import interface
+from blueapi.service.middleware import (
+    ObservabilityContextPropagator,
+    VersionHeaders,
+)
 from blueapi.worker import TrackableTask, WorkerState
 from blueapi.worker.event import TaskStatusEnum
 
@@ -124,8 +125,9 @@ def get_app(config: ApplicationConfig):
     app.include_router(secure_router, dependencies=dependencies)
     app.add_exception_handler(KeyError, on_key_error_404)
     app.add_exception_handler(jwt.PyJWTError, on_token_error_401)
-    app.middleware("http")(add_version_headers)
-    app.middleware("http")(inject_propagated_observability_context)
+
+    app.add_middleware(ObservabilityContextPropagator)
+    app.add_middleware(VersionHeaders)
     app.middleware("http")(log_request_details)
     if config.api.cors:
         app.add_middleware(
@@ -569,15 +571,6 @@ def start(config: ApplicationConfig):
     )
 
 
-async def add_version_headers(
-    request: Request, call_next: Callable[[Request], Awaitable[Response]]
-):
-    response = await call_next(request)
-    response.headers["X-API-Version"] = ApplicationConfig.REST_API_VERSION
-    response.headers["X-BlueAPI-Version"] = __version__
-    return response
-
-
 async def log_request_details(
     request: Request, call_next: Callable[[Request], Awaitable[StreamingResponse]]
 ) -> Response:
@@ -599,25 +592,3 @@ async def log_request_details(
         LOGGER.info(log_message, extra=extra)
 
     return response
-
-
-async def inject_propagated_observability_context(
-    request: Request, call_next: Callable[[Request], Awaitable[Response]]
-) -> Response:
-    """Middleware to extract any propagated observability context from the
-    HTTP headers and attach it to the local one.
-    """
-    headers = request.headers
-    if ApplicationConfig.CONTEXT_HEADER in headers:
-        carrier = {
-            ApplicationConfig.CONTEXT_HEADER: headers[ApplicationConfig.CONTEXT_HEADER]
-        }
-        if ApplicationConfig.VENDOR_CONTEXT_HEADER in headers:
-            carrier[ApplicationConfig.VENDOR_CONTEXT_HEADER] = headers[
-                ApplicationConfig.VENDOR_CONTEXT_HEADER
-            ]
-        ctx = get_global_textmap().extract(carrier)
-
-        attach(ctx)
-    response = await call_next(request)
-    return response

src/blueapi/service/middleware.py

@@ -0,0 +1,58 @@
+import logging
+
+from opentelemetry.context import attach
+from opentelemetry.propagate import get_global_textmap
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from blueapi import __version__
+from blueapi.config import ApplicationConfig
+
+OBS_LOGGER = logging.getLogger("blueapi.service.middleware.observability")
+
+CONTEXT_HEADER = ApplicationConfig.CONTEXT_HEADER.encode()
+VENDOR_CONTEXT_HEADER = ApplicationConfig.VENDOR_CONTEXT_HEADER.encode()
+
+API_VERSION = (b"x-api-version", ApplicationConfig.REST_API_VERSION.encode("utf-8"))
+VERSION = (b"x-blueapi-version", __version__.encode("utf-8"))
+
+
+class VersionHeaders:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope.get("type") not in ("websocket", "http"):
+            return await self.app(scope, receive, send)
+
+        async def local_send(message: Message):
+            if message["type"] in ("websocket.accept", "http.response.start"):
+                message["headers"].append(VERSION)
+                message["headers"].append(API_VERSION)
+            await send(message)
+
+        await self.app(scope, receive, local_send)
+
+
+class ObservabilityContextPropagator:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope["type"] not in ("http", "websocket"):
+            return await self.app(scope, receive, send)
+
+        ctx = None
+        v_ctx = None
+        for key, val in scope.get("headers", ()):
+            if key == CONTEXT_HEADER:
+                ctx = val.decode()
+            elif key == VENDOR_CONTEXT_HEADER:
+                v_ctx = val.decode()
+        if ctx:
+            OBS_LOGGER.debug("Propagating observability context: %s, %s", ctx, v_ctx)
+            carrier = {ApplicationConfig.CONTEXT_HEADER: ctx}
+            if v_ctx:
+                carrier[ApplicationConfig.VENDOR_CONTEXT_HEADER] = v_ctx
+            attach(get_global_textmap().extract(carrier))
+
+        await self.app(scope, receive, send)

tests/unit_tests/service/test_main.py

@@ -8,15 +8,15 @@
 from blueapi import __version__
 from blueapi.config import ApplicationConfig
 from blueapi.service.main import (
-    add_version_headers,
     get_passthrough_headers,
     log_request_details,
 )
+from blueapi.service.middleware import VersionHeaders
 
 
 async def test_add_version_header():
     app = FastAPI()
-    app.middleware("http")(add_version_headers)
+    app.add_middleware(VersionHeaders)
 
     @app.get("/")
     async def root():