adam/pyMC_Repeater
1
0
Fork 0
mirror of https://github.com/pyMC-dev/pyMC_Repeater.git synced 2026-07-26 19:42:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: allow API key authentication in WebSocket handler

Browse Source 
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)
This commit is contained in:
Will Dillon
2026-07-05 02:45:27 +00:00
committed by Rightup
parent 667169c7ff
commit d1a9230ab5
1 changed files with 35 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
repeater/data_acquisition/websocket_handler.py
+35 -23
View File
@@ -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"""