api: Make authentication work with WebSockets, too

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>

diff --git a/src/api/auth.py b/src/api/auth.py
index 1170a5d64d87ca91648853ca5863fa7fecce0d0a..bb7c46517c27e57e156406dd34f67fe2a22c8fb6 100644 (file)
--- a/src/api/auth.py
+++ b/src/api/auth.py
@@ -25,6 +25,7 @@ import kerberos
 import os
 import pydantic
 import socket
+import starlette
 import typing
 
 from . import apiv1
@@ -51,8 +52,18 @@ router = fastapi.APIRouter(
        tags=["Authentication"],
 )
 
+class HTTPBearer(fastapi.security.HTTPBearer):
+       """
+               This is a custom implementation of HTTPBearer which also supports WebSockets
+       """
+       # This works because fastapi.Request and fastapi.WebSocket both inherit from
+       # starlette's HTTPConnection.
+       async def __call__(self, request: starlette.requests.HTTPConnection) \
+                       -> typing.Optional[fastapi.security.HTTPAuthorizationCredentials]:
+               return await super().__call__(request)
+
 # Class to extract the Bearer token from the request headers
-security = fastapi.security.HTTPBearer()
+security = HTTPBearer()
 
 async def get_current_principal(
        credentials: fastapi.security.HTTPAuthorizationCredentials = fastapi.Depends(security)