Add WebSocket authentication via api_key

- When --api_key is set, WebSocket connections require auth too
- Supports Authorization: Bearer <key> header or ?token=<key> query param
- Unauthenticated connections receive HTTP 401 before upgrade
- Uses websockets process_request callback (no resource allocation before auth)
- Added 5 unit tests for WebSocket auth handler
This commit is contained in:
Aaron Boxer
2026-04-17 10:21:14 -04:00
committed by Aaron Boxer
parent b648bcb2a4
commit 5334ea0f7a
3 changed files with 63 additions and 3 deletions
+17 -1
View File
@@ -769,6 +769,21 @@ class TranscriptionServer:
logging.info(f"✅ OpenAI-Compatible API started on http://0.0.0.0:{rest_port}")
# Original WebSocket server (always supported)
extra_ws_kwargs = {}
if api_key:
def _ws_auth(path, request_headers):
auth = request_headers.get("Authorization", "")
token_param = None
# Check query string for token parameter
if "?" in path:
from urllib.parse import urlparse, parse_qs
parsed = urlparse(path)
token_param = parse_qs(parsed.query).get("token", [None])[0]
if auth == f"Bearer {api_key}" or token_param == api_key:
return None # Allow connection
return (401, [("Content-Type", "text/plain")], b"Unauthorized\n")
extra_ws_kwargs["process_request"] = _ws_auth
with serve(
functools.partial(
self.recv_audio,
@@ -779,7 +794,8 @@ class TranscriptionServer:
trt_py_session=trt_py_session,
),
host,
port
port,
**extra_ws_kwargs,
) as server:
server.serve_forever()