Add optional API key auth and rate limiting for REST API

- api_key param: requires 'Authorization: Bearer <key>' 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
This commit is contained in:
Aaron Boxer
2026-04-17 09:37:06 -04:00
committed by Aaron Boxer
parent e86d98dd80
commit b648bcb2a4
3 changed files with 133 additions and 1 deletions
+15
View File
@@ -96,6 +96,19 @@ if __name__ == "__main__":
default=0, default=0,
help='Port for Prometheus /metrics endpoint. 0 = disabled (default). Requires prometheus_client.' 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 <key>" 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() args = parser.parse_args()
if args.backend == "tensorrt": if args.backend == "tensorrt":
@@ -127,4 +140,6 @@ if __name__ == "__main__":
batch_window_ms=args.batch_window_ms, batch_window_ms=args.batch_window_ms,
raw_pcm_input=args.raw_pcm_input, raw_pcm_input=args.raw_pcm_input,
metrics_port=args.metrics_port, metrics_port=args.metrics_port,
api_key=args.api_key,
rate_limit_rpm=args.rate_limit_rpm,
) )
+86
View File
@@ -1,6 +1,7 @@
import json import json
import time import time
import threading import threading
import collections
import unittest import unittest
from unittest import mock from unittest import mock
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
@@ -566,5 +567,90 @@ class TestRESTAPIParamWarnings(unittest.TestCase):
self.assertGreaterEqual(len(ignored), 2) 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__": if __name__ == "__main__":
unittest.main() unittest.main()
+32 -1
View File
@@ -1,6 +1,7 @@
import os import os
import time import time
import threading import threading
import collections
import queue import queue
import json import json
import functools import functools
@@ -8,7 +9,7 @@ import logging
import shutil import shutil
import tempfile import tempfile
from typing import Optional, List 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.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse
from starlette.responses import PlainTextResponse, JSONResponse, StreamingResponse from starlette.responses import PlainTextResponse, JSONResponse, StreamingResponse
@@ -533,6 +534,8 @@ class TranscriptionServer:
batch_window_ms=50, batch_window_ms=50,
raw_pcm_input=False, raw_pcm_input=False,
metrics_port: int = 0, metrics_port: int = 0,
api_key: Optional[str] = None,
rate_limit_rpm: int = 0,
segment_post_processor=None): segment_post_processor=None):
""" """
Run the transcription server. Run the transcription server.
@@ -612,6 +615,34 @@ class TranscriptionServer:
allow_headers=["*"], # Allows all headers 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") @app.post("/v1/audio/transcriptions")
async def transcribe( async def transcribe(