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