222 lines
7.8 KiB
Python
222 lines
7.8 KiB
Python
"""
|
|
Middleware for AI Service.
|
|
[AC-AISVC-10, AC-AISVC-12, AC-AISVC-50] X-Tenant-Id header validation, tenant context injection, and API Key authentication.
|
|
"""
|
|
|
|
import logging
|
|
import re
|
|
from collections.abc import Callable
|
|
|
|
from fastapi import Request, Response, status
|
|
from fastapi.responses import JSONResponse
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
|
|
from app.core.exceptions import ErrorCode, ErrorResponse
|
|
from app.core.tenant import clear_tenant_context, set_tenant_context
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TENANT_ID_HEADER = "X-Tenant-Id"
|
|
API_KEY_HEADER = "X-API-Key"
|
|
ACCEPT_HEADER = "Accept"
|
|
SSE_CONTENT_TYPE = "text/event-stream"
|
|
|
|
TENANT_ID_PATTERN = re.compile(r'^[^@]+@ash@\d{4}$')
|
|
|
|
PATHS_SKIP_API_KEY = {
|
|
"/health",
|
|
"/ai/health",
|
|
"/docs",
|
|
"/redoc",
|
|
"/openapi.json",
|
|
}
|
|
|
|
|
|
def validate_tenant_id_format(tenant_id: str) -> bool:
|
|
"""
|
|
[AC-AISVC-10] Validate tenant ID format: name@ash@year
|
|
Examples: szmp@ash@2026, abc123@ash@2025
|
|
"""
|
|
return bool(TENANT_ID_PATTERN.match(tenant_id))
|
|
|
|
|
|
def parse_tenant_id(tenant_id: str) -> tuple[str, str]:
|
|
"""
|
|
[AC-AISVC-10] Parse tenant ID into name and year.
|
|
Returns: (name, year)
|
|
"""
|
|
parts = tenant_id.split('@')
|
|
return parts[0], parts[2]
|
|
|
|
|
|
class ApiKeyMiddleware(BaseHTTPMiddleware):
|
|
"""
|
|
[AC-AISVC-50] Middleware to validate API Key for all requests.
|
|
|
|
Features:
|
|
- Validates X-API-Key header against in-memory cache
|
|
- Skips validation for health/docs endpoints
|
|
- Returns 401 for missing or invalid API key
|
|
"""
|
|
|
|
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
|
if self._should_skip_api_key(request.url.path):
|
|
return await call_next(request)
|
|
|
|
api_key = request.headers.get(API_KEY_HEADER)
|
|
|
|
if not api_key or not api_key.strip():
|
|
logger.warning(f"[AC-AISVC-50] Missing X-API-Key header for {request.url.path}")
|
|
return JSONResponse(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
content=ErrorResponse(
|
|
code=ErrorCode.UNAUTHORIZED.value,
|
|
message="Missing required header: X-API-Key",
|
|
).model_dump(exclude_none=True),
|
|
)
|
|
|
|
api_key = api_key.strip()
|
|
|
|
from app.services.api_key import get_api_key_service
|
|
service = get_api_key_service()
|
|
|
|
if not service._initialized:
|
|
logger.warning("[AC-AISVC-50] API key service not initialized, attempting lazy initialization...")
|
|
try:
|
|
from app.core.database import async_session_maker
|
|
async with async_session_maker() as session:
|
|
await service.initialize(session)
|
|
logger.info(f"[AC-AISVC-50] API key service lazy initialized with {len(service._keys_cache)} keys")
|
|
except Exception as e:
|
|
logger.error(f"[AC-AISVC-50] Failed to initialize API key service: {e}")
|
|
|
|
if not service.validate_key(api_key):
|
|
logger.warning(f"[AC-AISVC-50] Invalid API key for {request.url.path}")
|
|
return JSONResponse(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
content=ErrorResponse(
|
|
code=ErrorCode.UNAUTHORIZED.value,
|
|
message="Invalid API key",
|
|
).model_dump(exclude_none=True),
|
|
)
|
|
|
|
return await call_next(request)
|
|
|
|
def _should_skip_api_key(self, path: str) -> bool:
|
|
"""Check if the path should skip API key validation."""
|
|
if path in PATHS_SKIP_API_KEY:
|
|
return True
|
|
for skip_path in PATHS_SKIP_API_KEY:
|
|
if path.startswith(skip_path):
|
|
return True
|
|
return False
|
|
|
|
|
|
class TenantContextMiddleware(BaseHTTPMiddleware):
|
|
"""
|
|
[AC-AISVC-10, AC-AISVC-12] Middleware to extract and validate X-Tenant-Id header.
|
|
Injects tenant context into request state for downstream processing.
|
|
Validates tenant ID format and auto-creates tenant if not exists.
|
|
"""
|
|
|
|
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
|
clear_tenant_context()
|
|
|
|
if request.url.path in ("/health", "/ai/health"):
|
|
return await call_next(request)
|
|
|
|
tenant_id = request.headers.get(TENANT_ID_HEADER)
|
|
|
|
if not tenant_id or not tenant_id.strip():
|
|
logger.warning("[AC-AISVC-12] Missing or empty X-Tenant-Id header")
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content=ErrorResponse(
|
|
code=ErrorCode.MISSING_TENANT_ID.value,
|
|
message="Missing required header: X-Tenant-Id",
|
|
).model_dump(exclude_none=True),
|
|
)
|
|
|
|
tenant_id = tenant_id.strip()
|
|
|
|
if not validate_tenant_id_format(tenant_id):
|
|
logger.warning(f"[AC-AISVC-10] Invalid tenant ID format: {tenant_id}")
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content=ErrorResponse(
|
|
code=ErrorCode.INVALID_TENANT_ID.value,
|
|
message="Invalid tenant ID format. Expected: name@ash@year (e.g., szmp@ash@2026)",
|
|
).model_dump(exclude_none=True),
|
|
)
|
|
|
|
if request.url.path.startswith("/admin/") or request.url.path.startswith("/ai/"):
|
|
try:
|
|
await self._ensure_tenant_exists(request, tenant_id)
|
|
except Exception as e:
|
|
logger.error(f"[AC-AISVC-10] Failed to ensure tenant exists: {e}")
|
|
|
|
set_tenant_context(tenant_id)
|
|
request.state.tenant_id = tenant_id
|
|
|
|
logger.info(f"[AC-AISVC-10] Tenant context set: tenant_id={tenant_id}, path={request.url.path}")
|
|
|
|
try:
|
|
logger.info(f"[MIDDLEWARE] Calling next handler for path={request.url.path}")
|
|
response = await call_next(request)
|
|
logger.info(f"[MIDDLEWARE] Response received for path={request.url.path}, status={response.status_code}")
|
|
except Exception as e:
|
|
import traceback
|
|
logger.error(f"[MIDDLEWARE] Exception in call_next for path={request.url.path}: {type(e).__name__}: {e}\n{traceback.format_exc()}")
|
|
raise
|
|
finally:
|
|
clear_tenant_context()
|
|
|
|
return response
|
|
|
|
async def _ensure_tenant_exists(self, request: Request, tenant_id: str) -> None:
|
|
"""
|
|
[AC-AISVC-10] Ensure tenant exists in database, create if not.
|
|
"""
|
|
from sqlalchemy import select
|
|
|
|
from app.core.database import async_session_maker
|
|
from app.models.entities import Tenant
|
|
|
|
name, year = parse_tenant_id(tenant_id)
|
|
|
|
async with async_session_maker() as session:
|
|
stmt = select(Tenant).where(Tenant.tenant_id == tenant_id)
|
|
result = await session.execute(stmt)
|
|
existing_tenant = result.scalar_one_or_none()
|
|
|
|
if existing_tenant:
|
|
logger.debug(f"[AC-AISVC-10] Tenant already exists: {tenant_id}")
|
|
return
|
|
|
|
new_tenant = Tenant(
|
|
tenant_id=tenant_id,
|
|
name=name,
|
|
year=year,
|
|
)
|
|
session.add(new_tenant)
|
|
await session.commit()
|
|
|
|
logger.info(f"[AC-AISVC-10] Auto-created new tenant: {tenant_id} (name={name}, year={year})")
|
|
|
|
|
|
def is_sse_request(request: Request) -> bool:
|
|
"""
|
|
[AC-AISVC-06] Check if the request expects SSE streaming response.
|
|
Based on Accept header: text/event-stream indicates SSE mode.
|
|
"""
|
|
accept_header = request.headers.get(ACCEPT_HEADER, "")
|
|
return SSE_CONTENT_TYPE in accept_header
|
|
|
|
|
|
def get_response_mode(request: Request) -> str:
|
|
"""
|
|
[AC-AISVC-06] Determine response mode based on Accept header.
|
|
Returns 'streaming' for SSE, 'json' for regular JSON response.
|
|
"""
|
|
return "streaming" if is_sse_request(request) else "json"
|