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:
+17
-1
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user