From b648bcb2a445e2e132119222aa6a4d248a5c65ae Mon Sep 17 00:00:00 2001 From: Aaron Boxer Date: Fri, 17 Apr 2026 09:37:06 -0400 Subject: [PATCH] Add optional API key auth and rate limiting for REST API - api_key param: requires 'Authorization: Bearer ' header - rate_limit_rpm param: per-IP sliding-window rate limit (requests/min) - Both are off by default (backward compatible) - CLI flags: --api_key, --rate_limit_rpm - Added 5 unit tests for auth and rate limiting --- run_server.py | 15 ++++++ tests/test_server_extended.py | 86 +++++++++++++++++++++++++++++++++++ whisper_live/server.py | 33 +++++++++++++- 3 files changed, 133 insertions(+), 1 deletion(-) diff --git a/run_server.py b/run_server.py index d48eb70..858b74e 100644 --- a/run_server.py +++ b/run_server.py @@ -96,6 +96,19 @@ if __name__ == "__main__": default=0, help='Port for Prometheus /metrics endpoint. 0 = disabled (default). Requires prometheus_client.' ) + parser.add_argument( + '--api_key', + type=str, + default=None, + help='Optional API key for authenticating REST API requests. ' + 'Clients must send "Authorization: Bearer " header.' + ) + parser.add_argument( + '--rate_limit_rpm', + type=int, + default=0, + help='Maximum REST API requests per minute per client IP. 0 = unlimited (default).' + ) args = parser.parse_args() if args.backend == "tensorrt": @@ -127,4 +140,6 @@ if __name__ == "__main__": batch_window_ms=args.batch_window_ms, raw_pcm_input=args.raw_pcm_input, metrics_port=args.metrics_port, + api_key=args.api_key, + rate_limit_rpm=args.rate_limit_rpm, ) \ No newline at end of file diff --git a/tests/test_server_extended.py b/tests/test_server_extended.py index 47b608f..744b87c 100644 --- a/tests/test_server_extended.py +++ b/tests/test_server_extended.py @@ -1,6 +1,7 @@ import json import time import threading +import collections import unittest from unittest import mock from unittest.mock import MagicMock, patch @@ -566,5 +567,90 @@ class TestRESTAPIParamWarnings(unittest.TestCase): self.assertGreaterEqual(len(ignored), 2) +class TestAPIKeyAuth(unittest.TestCase): + """Test optional API key authentication middleware.""" + + @classmethod + def setUpClass(cls): + from fastapi import FastAPI, Request + from fastapi.testclient import TestClient + from fastapi.responses import JSONResponse as JSONR + + app = FastAPI() + + @app.middleware("http") + async def _check_api_key(request: Request, call_next): + auth = request.headers.get("Authorization", "") + if auth != "Bearer test-secret": + return JSONR({"error": "Invalid or missing API key"}, status_code=401) + return await call_next(request) + + @app.get("/ping") + async def ping(): + return {"status": "ok"} + + cls.test_client = TestClient(app) + + def test_missing_key_returns_401(self): + resp = self.test_client.get("/ping") + self.assertEqual(resp.status_code, 401) + + def test_wrong_key_returns_401(self): + resp = self.test_client.get("/ping", headers={"Authorization": "Bearer wrong"}) + self.assertEqual(resp.status_code, 401) + + def test_correct_key_returns_200(self): + resp = self.test_client.get("/ping", headers={"Authorization": "Bearer test-secret"}) + self.assertEqual(resp.status_code, 200) + self.assertEqual(resp.json()["status"], "ok") + + +class TestRateLimiting(unittest.TestCase): + """Test per-IP rate limiting middleware.""" + + def _make_app(self, rpm_limit=3): + from fastapi import FastAPI, Request + from fastapi.testclient import TestClient + from fastapi.responses import JSONResponse as JSONR + + _rate_lock = threading.Lock() + _rate_buckets: dict = {} + + app = FastAPI() + + @app.middleware("http") + async def _rate_limit(request: Request, call_next): + client_ip = request.client.host if request.client else "unknown" + now = time.time() + with _rate_lock: + bucket = _rate_buckets.setdefault(client_ip, collections.deque()) + while bucket and bucket[0] < now - 60: + bucket.popleft() + if len(bucket) >= rpm_limit: + return JSONR({"error": "Rate limit exceeded"}, status_code=429) + bucket.append(now) + return await call_next(request) + + @app.get("/ping") + async def ping(): + return {"status": "ok"} + + return TestClient(app) + + def test_within_limit_succeeds(self): + client = self._make_app(rpm_limit=3) + for _ in range(3): + resp = client.get("/ping") + self.assertEqual(resp.status_code, 200) + + def test_exceeding_limit_returns_429(self): + client = self._make_app(rpm_limit=3) + for _ in range(3): + client.get("/ping") + resp = client.get("/ping") + self.assertEqual(resp.status_code, 429) + self.assertIn("Rate limit", resp.json()["error"]) + + if __name__ == "__main__": unittest.main() diff --git a/whisper_live/server.py b/whisper_live/server.py index 1a7a36b..c3e8a70 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -1,6 +1,7 @@ import os import time import threading +import collections import queue import json import functools @@ -8,7 +9,7 @@ import logging import shutil import tempfile from typing import Optional, List -from fastapi import FastAPI, UploadFile, Form +from fastapi import FastAPI, UploadFile, Form, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from starlette.responses import PlainTextResponse, JSONResponse, StreamingResponse @@ -533,6 +534,8 @@ class TranscriptionServer: batch_window_ms=50, raw_pcm_input=False, metrics_port: int = 0, + api_key: Optional[str] = None, + rate_limit_rpm: int = 0, segment_post_processor=None): """ Run the transcription server. @@ -612,6 +615,34 @@ class TranscriptionServer: allow_headers=["*"], # Allows all headers ) + # Optional API key authentication + if api_key: + @app.middleware("http") + async def _check_api_key(request: Request, call_next): + auth = request.headers.get("Authorization", "") + if auth != f"Bearer {api_key}": + return JSONResponse({"error": "Invalid or missing API key"}, status_code=401) + return await call_next(request) + + # Optional rate limiting (requests per minute per client IP) + if rate_limit_rpm > 0: + _rate_lock = threading.Lock() + _rate_buckets: dict = {} # ip -> deque of timestamps + + @app.middleware("http") + async def _rate_limit(request: Request, call_next): + client_ip = request.client.host if request.client else "unknown" + now = time.time() + with _rate_lock: + bucket = _rate_buckets.setdefault(client_ip, collections.deque()) + # Discard entries older than 60s + while bucket and bucket[0] < now - 60: + bucket.popleft() + if len(bucket) >= rate_limit_rpm: + return JSONResponse({"error": "Rate limit exceeded"}, status_code=429) + bucket.append(now) + return await call_next(request) + @app.post("/v1/audio/transcriptions") async def transcribe(