Spaces:
Running
Running
| from __future__ import annotations | |
| import time | |
| from typing import Any | |
| from uuid import uuid4 | |
| from fastapi import Request | |
| from fastapi.responses import JSONResponse | |
| from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint | |
| from starlette.responses import Response | |
| from app.core.config import Settings | |
| from app.core.logger import get_logger | |
| from app.security.audit import AuditService | |
| from app.security.context import AuthContext, auth_context, http_auth_applied | |
| from app.security.errors import ForbiddenError, RateLimitError, UnauthorizedError | |
| from app.security.policy import ScopePolicy | |
| from app.security.rate_limit import APIKeyRateLimiter, RateLimitLease | |
| from app.security.service import APIKeyService | |
| from app.security.studio_identity import constrain_context_to_membership, verify_studio_identity | |
| logger = get_logger(__name__) | |
| class APIKeyAuthenticationMiddleware(BaseHTTPMiddleware): | |
| """Authenticates, authorizes, rate-limits, and audits protected HTTP requests.""" | |
| def __init__( | |
| self, | |
| app: Any, | |
| *, | |
| settings: Settings, | |
| api_keys: APIKeyService, | |
| rate_limiter: APIKeyRateLimiter, | |
| audit: AuditService, | |
| ) -> None: | |
| super().__init__(app) | |
| self.settings = settings | |
| self.api_keys = api_keys | |
| self.rate_limiter = rate_limiter | |
| self.audit = audit | |
| self.policy = ScopePolicy() | |
| async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: | |
| if not getattr(request.state, "request_id", None): | |
| request.state.request_id = str(uuid4()) | |
| if not self.settings.auth_enabled or self.policy.is_public(request): | |
| return await call_next(request) | |
| started = time.monotonic() | |
| context: AuthContext | None = None | |
| lease: RateLimitLease | None = None | |
| context_token = None | |
| http_auth_token = None | |
| response_code = 500 | |
| required_scope: str | None = None | |
| bytes_uploaded = self._content_length(request.headers.get("content-length")) | |
| bytes_downloaded = 0 | |
| try: | |
| api_key = self._bearer_token(request.headers.get("authorization")) | |
| context = await self.api_keys.authenticate(api_key) | |
| studio_header_names = ( | |
| "x-mr-studio-subject", | |
| "x-mr-studio-email", | |
| "x-mr-studio-name", | |
| "x-mr-studio-timestamp", | |
| "x-mr-studio-signature", | |
| ) | |
| if any(request.headers.get(name) is not None for name in studio_header_names): | |
| secret = ( | |
| self.settings.studio_identity_secret.get_secret_value() | |
| if self.settings.studio_identity_secret | |
| else "" | |
| ) | |
| try: | |
| identity = verify_studio_identity(request.headers, secret) | |
| except ValueError as exc: | |
| logger.warning("rejected Studio identity envelope: %s", str(exc)) | |
| raise UnauthorizedError from exc | |
| principal = await self.api_keys.tenants.resolve_studio_subject(identity.subject) | |
| request.state.studio_identity_verified = True | |
| try: | |
| # The shared key authenticates transport only; persisted | |
| # membership capabilities constrain its effective scopes. | |
| context = constrain_context_to_membership(context, principal) | |
| except ValueError as exc: | |
| logger.warning("rejected Studio membership role: %s", str(exc)) | |
| raise ForbiddenError from exc | |
| request.state.auth = context | |
| context_token = auth_context.set(context) | |
| required_scope = await self.policy.required_scope(request) | |
| lease = await self.rate_limiter.acquire( | |
| context, | |
| is_job=self.policy.is_job(required_scope) | |
| or request.url.path.startswith("/v1/projects/") | |
| and "/renders" in request.url.path | |
| and request.method == "POST", | |
| is_upload=self.policy.is_upload(request, required_scope), | |
| uploaded_bytes=bytes_uploaded, | |
| ) | |
| if context.membership_role == "viewer" and required_scope not in { | |
| "templates:read", | |
| "operations:read", | |
| "jobs:read", | |
| "assets:read", | |
| "mcp:read", | |
| "system:read", | |
| "social:accounts:read", | |
| "social:posts:read", | |
| "social:schedules:read", | |
| "social:analytics:read", | |
| "analytics:read", | |
| "generation:providers:read", | |
| "generation:requests:read", | |
| "ai:read", | |
| "copilot:read", | |
| "projects:read", | |
| "members:read", | |
| "teams:read", | |
| "projects:collaborate", | |
| "comments:read", | |
| "approvals:read", | |
| }: | |
| raise ForbiddenError | |
| self.api_keys.authorize(context, required_scope) | |
| await self._apply_social_rate_limit(request, context) | |
| await self.api_keys.mark_used(context) | |
| http_auth_token = http_auth_applied.set(True) | |
| response = await call_next(request) | |
| response_code = response.status_code | |
| bytes_downloaded = self._content_length(response.headers.get("content-length")) | |
| response.headers.setdefault("X-Request-ID", request.state.request_id) | |
| return response | |
| except UnauthorizedError: | |
| response_code = 401 | |
| return self._error( | |
| 401, | |
| "Unauthorized", | |
| "Invalid or expired API key.", | |
| request, | |
| {"WWW-Authenticate": "Bearer"}, | |
| ) | |
| except ForbiddenError: | |
| response_code = 403 | |
| detail = f"Missing required scope: {required_scope}." if required_scope else "Workspace identity or membership authorization was denied." | |
| response = self._error(403, "Forbidden", detail, request) | |
| bytes_downloaded = len(response.body) | |
| return response | |
| except RateLimitError as exc: | |
| response_code = 429 | |
| response = self._error( | |
| 429, | |
| "Rate limit exceeded", | |
| "Retry later.", | |
| request, | |
| {"Retry-After": str(exc.retry_after)}, | |
| ) | |
| bytes_downloaded = len(response.body) | |
| return response | |
| finally: | |
| if lease is not None: | |
| await lease.release() | |
| if http_auth_token is not None: | |
| http_auth_applied.reset(http_auth_token) | |
| if context_token is not None: | |
| auth_context.reset(context_token) | |
| if context is not None: | |
| elapsed_ms = max(0, round((time.monotonic() - started) * 1000)) | |
| await self._audit_request( | |
| request, | |
| context, | |
| response_code, | |
| elapsed_ms, | |
| bytes_uploaded, | |
| bytes_downloaded, | |
| ) | |
| async def _apply_social_rate_limit(self, request: Request, context: AuthContext) -> None: | |
| path = request.url.path | |
| if path.startswith("/v1/analytics"): | |
| category = "analytics_sync" if path.endswith(("/sync", "/cancel")) else "analytics_read" | |
| limit = ( | |
| max(1, self.settings.social_analytics_requests_per_minute // 6) | |
| if category == "analytics_sync" | |
| else self.settings.social_analytics_requests_per_minute | |
| ) | |
| await self.rate_limiter.acquire_category( | |
| context, | |
| category, | |
| limit=limit, | |
| window_seconds=60, | |
| ) | |
| return | |
| if not path.startswith("/v1/social"): | |
| return | |
| if "/analytics" in path: | |
| await self.rate_limiter.acquire_category( | |
| context, | |
| "social_analytics", | |
| limit=self.settings.social_analytics_requests_per_minute, | |
| window_seconds=60, | |
| ) | |
| elif path.endswith("/connect") or path.endswith("/callback") or path.endswith("/refresh"): | |
| await self.rate_limiter.acquire_category( | |
| context, | |
| "social_oauth", | |
| limit=self.settings.social_oauth_requests_per_hour, | |
| window_seconds=3600, | |
| ) | |
| elif path.endswith("/publish"): | |
| await self.rate_limiter.acquire_category( | |
| context, | |
| "social_publish", | |
| limit=self.settings.social_publish_requests_per_minute, | |
| window_seconds=60, | |
| ) | |
| elif path.endswith("/schedule"): | |
| await self.rate_limiter.acquire_category( | |
| context, | |
| "social_schedule", | |
| limit=self.settings.social_schedule_requests_per_minute, | |
| window_seconds=60, | |
| ) | |
| elif "/batch" in path: | |
| await self.rate_limiter.acquire_category( | |
| context, | |
| "social_publish", | |
| limit=self.settings.social_publish_requests_per_minute, | |
| window_seconds=60, | |
| ) | |
| async def _audit_request( | |
| self, | |
| request: Request, | |
| context: AuthContext, | |
| response_code: int, | |
| elapsed_ms: int, | |
| bytes_uploaded: int, | |
| bytes_downloaded: int, | |
| ) -> None: | |
| try: | |
| await self.audit.record( | |
| request_id=str(getattr(request.state, "request_id", ""))[:64], | |
| api_key_id=str(context.api_key_id), | |
| key_name=str(getattr(context, "key_name", ""))[:128], | |
| ip_address=request.client.host if request.client else None, | |
| user_agent=request.headers.get("user-agent", "")[:512] or None, | |
| endpoint=request.url.path[:512], | |
| http_method=request.method[:16], | |
| response_code=response_code, | |
| processing_time_ms=elapsed_ms, | |
| bytes_uploaded=max(0, bytes_uploaded), | |
| bytes_downloaded=max(0, bytes_downloaded), | |
| ) | |
| except Exception: | |
| logger.exception("failed to persist API request audit record") | |
| def _error( | |
| status_code: int, | |
| error: str, | |
| message: str, | |
| request: Request, | |
| headers: dict[str, str] | None = None, | |
| ) -> JSONResponse: | |
| response_headers = dict(headers or {}) | |
| request_id = str(getattr(request.state, "request_id", "")) | |
| if request_id: | |
| response_headers.setdefault("X-Request-ID", request_id) | |
| return JSONResponse( | |
| status_code=status_code, | |
| headers=response_headers, | |
| content={ | |
| "error": { | |
| "code": error, | |
| "message": message, | |
| "request_id": request_id, | |
| } | |
| }, | |
| ) | |
| def _bearer_token(value: str | None) -> str: | |
| if not value: | |
| raise UnauthorizedError | |
| scheme, separator, token = value.partition(" ") | |
| if not separator or scheme.casefold() != "bearer" or not token.strip(): | |
| raise UnauthorizedError | |
| return token.strip() | |
| def _content_length(value: str | None) -> int: | |
| if not value: | |
| return 0 | |
| try: | |
| parsed = int(value) | |
| except (TypeError, ValueError): | |
| return 0 | |
| return max(0, parsed) | |