From d1a9230ab598d4a0807d54882b596d2da8e1634f Mon Sep 17 00:00:00 2001
From: Will Dillon <william@housedillon.com>
Date: Sun, 5 Jul 2026 02:45:27 +0000
Subject: [PATCH] feat: allow API key authentication in WebSocket handler

The PacketWebSocket.opened() method only accepts JWT tokens for
authentication. API keys (used via X-API-Key header for REST endpoints)
are silently rejected, forcing WS clients to obtain a JWT through a
separate login flow.

Add fallback to token_manager.verify_token() when JWT verification
fails or when an X-API-Key header is present. This lets API token
holders connect directly to the WebSocket without needing a JWT.

Auth priority:
1. JWT in ?token= query parameter (existing behavior)
2. API key in X-API-Key header
3. API key in ?token= query parameter (falls through from JWT reject)
---
 .../data_acquisition/websocket_handler.py     | 58 +++++++++++--------
 1 file changed, 35 insertions(+), 23 deletions(-)

diff --git a/repeater/data_acquisition/websocket_handler.py b/repeater/data_acquisition/websocket_handler.py
index 19ba2c8..17aeb1f 100644
--- a/repeater/data_acquisition/websocket_handler.py
+++ b/repeater/data_acquisition/websocket_handler.py
@@ -30,10 +30,9 @@ _websocket_plugin = None
 class PacketWebSocket(WebSocket):
     def opened(self):
         """Called when a WebSocket connection is established"""
-        # Authenticate using JWT provided as query parameter (token=)
         jwt_handler = cherrypy.config.get("jwt_handler")
+        token_manager = cherrypy.config.get("token_manager")
 
-        # Get query string from environ
         qs = ""
         if hasattr(self, "environ"):
             qs = self.environ.get("QUERY_STRING", "")
@@ -42,38 +41,51 @@ class PacketWebSocket(WebSocket):
         token = params.get("token", [None])[0]
         client_id = params.get("client_id", [None])[0]
 
+        api_key = self.environ.get("HTTP_X_API_KEY", "") if hasattr(self, "environ") else ""
+
         if not jwt_handler:
             logger.warning("WebSocket connection rejected: no JWT handler configured")
             self.close(code=1011, reason="server configuration error")
             return
 
-        if not token:
+        if not token and not api_key:
             logger.warning("WebSocket connection rejected: missing token")
             self.close(code=1008, reason="unauthorized")
             return
 
-        try:
-            payload = jwt_handler.verify_jwt(token)
-            if not payload:
-                logger.warning("WebSocket connection rejected: invalid token")
-                self.close(code=1008, reason="unauthorized")
-                return
-        except Exception as e:
-            logger.warning(f"WebSocket auth error: {e}")
-            self.close(code=1008, reason="unauthorized")
-            return
+        if token:
+            try:
+                payload = jwt_handler.verify_jwt(token)
+                if payload:
+                    if client_id and payload.get("client_id") and payload.get("client_id") != client_id:
+                        logger.warning("WebSocket connection rejected: client_id mismatch")
+                        self.close(code=1008, reason="unauthorized")
+                        return
+                    self.user = payload.get("sub")
+                    _connected_clients.add(self)
+                    logger.info(
+                        f"WebSocket connected ({self.user or 'unknown user'}). Total clients: {len(_connected_clients)}"
+                    )
+                    return
+            except Exception as e:
+                logger.warning(f"WebSocket JWT auth error: {e}")
 
-        if client_id and payload.get("client_id") and payload.get("client_id") != client_id:
-            logger.warning("WebSocket connection rejected: client_id mismatch")
-            self.close(code=1008, reason="unauthorized")
-            return
+        api_token = api_key or token
+        if api_token and token_manager:
+            try:
+                token_info = token_manager.verify_token(api_token)
+                if token_info:
+                    self.user = f"api_token:{token_info.get('name', 'unknown')}"
+                    _connected_clients.add(self)
+                    logger.info(
+                        f"WebSocket connected (API token: {token_info.get('name', 'unknown')}). Total clients: {len(_connected_clients)}"
+                    )
+                    return
+            except Exception as e:
+                logger.warning(f"WebSocket API key auth error: {e}")
 
-        # Auth success - store user and add to connected clients
-        self.user = payload.get("sub")  # type: ignore[attr-defined]
-        _connected_clients.add(self)
-        logger.info(
-            f"WebSocket connected ({self.user or 'unknown user'}). Total clients: {len(_connected_clients)}"
-        )
+        logger.warning("WebSocket connection rejected: no valid authentication")
+        self.close(code=1008, reason="unauthorized")
 
     def closed(self, code, reason=None):
         """Called when a WebSocket connection is closed"""