File size: 5,010 Bytes
860cdea 344b3b7 860cdea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """
API Middleware — Authentication & Rate Limiting
Validates X-API-Key header on every request.
Applies rate limiting via Redis.
"""
import logging
from typing import Optional
from fastapi import Request, HTTPException, Depends
from fastapi.security import APIKeyHeader
from src.db import mongo, redis as redis_db
from src.utils.hashing import hash_api_key
logger = logging.getLogger("llm_firewall.middleware")
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
async def validate_api_key(
request: Request,
api_key: Optional[str] = Depends(api_key_header),
) -> dict:
"""
Validate the API key from X-API-Key header.
Returns the key document from MongoDB.
Raises HTTPException on auth failure.
"""
# Check if key is in URL params (security violation)
client_host = request.client.host if request.client else "unknown"
if request.query_params.get("api_key") or request.query_params.get("apikey"):
logger.warning(f"API key in URL params from {client_host}")
raise HTTPException(
status_code=401,
detail="API key must be in X-API-Key header, not URL parameters",
)
if not api_key:
raise HTTPException(status_code=401, detail="API key required")
if not api_key.startswith("fw_live_") or len(api_key) != 72:
raise HTTPException(status_code=401, detail="Invalid API key format")
key_hash = hash_api_key(api_key)
keys_collection = mongo.get_keys_collection()
key_doc = await keys_collection.find_one({"key_hash": key_hash})
if not key_doc:
raise HTTPException(status_code=401, detail="API key not found")
if not key_doc.get("is_active", False):
raise HTTPException(status_code=401, detail="API key has been revoked")
rate_status = await redis_db.check_rate_limit(
api_key_id=str(key_doc["_id"])
)
if not rate_status.allowed:
raise HTTPException(
status_code=429,
detail={
"error": "rate_limit_exceeded",
"limit_type": rate_status.limit_type,
"retry_after_seconds": rate_status.retry_after_seconds,
},
headers={
"Retry-After": str(rate_status.retry_after_seconds),
"X-RateLimit-Limit": str(rate_status.limit),
"X-RateLimit-Remaining": "0",
"X-RateLimit-Reset": str(rate_status.reset_at),
},
)
request.state.rate_limit = rate_status
request.state.api_key_doc = key_doc
return key_doc
async def resolve_check_auth(request: Request) -> dict:
"""
Dual-auth resolver for /v1/check endpoints.
External flow: X-API-Key header → standard hash validation.
Dashboard flow: Authorization: Bearer <jwt> + X-Dashboard-Key-ID: <key_id>
→ validates JWT, looks up key by ID, asserts ownership.
ponytail: no new files, reuses validate_api_key and auth_middleware's JWT logic.
"""
api_key = request.headers.get("X-API-Key")
if api_key:
# External flow — existing path unchanged
return await validate_api_key(request, api_key)
# Dashboard flow
dashboard_key_id = request.headers.get("X-Dashboard-Key-ID")
auth_header = request.headers.get("Authorization", "")
if not dashboard_key_id or not auth_header.startswith("Bearer "):
raise HTTPException(status_code=401, detail="API key or dashboard session required")
token = auth_header.split(" ", 1)[1]
# Reuse JWT logic from auth_middleware inline to avoid circular imports
import os
from jose import JWTError, jwt as _jwt
from bson import ObjectId
SECRET_KEY = os.getenv("JWT_SECRET_KEY", "d9f933e574ce4f4f41da2de40e9b3f2e72452ee35e906d8639f36da8033bb982")
ALGORITHM = os.getenv("JWT_ALGORITHM", "HS256")
try:
payload = _jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
user_id = ObjectId(payload.get("sub"))
except (JWTError, Exception):
raise HTTPException(status_code=401, detail="Invalid session token")
try:
key_obj_id = ObjectId(dashboard_key_id)
except Exception:
raise HTTPException(status_code=400, detail="Invalid key ID")
keys_collection = mongo.get_keys_collection()
key_doc = await keys_collection.find_one({"_id": key_obj_id, "user_id": user_id, "is_active": True})
if not key_doc:
raise HTTPException(status_code=404, detail="API key not found or access denied")
# Rate limiting still applies
rate_status = await redis_db.check_rate_limit(api_key_id=str(key_doc["_id"]))
if not rate_status.allowed:
raise HTTPException(status_code=429, detail={"error": "rate_limit_exceeded"})
request.state.rate_limit = rate_status
request.state.api_key_doc = key_doc
return key_doc
|