92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
"""Edge hardening: run throttling, security headers.
|
|
|
|
Body caps come from Starlette's RequestBodyLimitMiddleware (wired in
|
|
main.py). All limits are settings-driven so tests can shrink them. State
|
|
is in-process memory: correct for the single-replica SQLite deployment;
|
|
revisit with Redis if the API ever scales horizontally.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from collections import defaultdict
|
|
|
|
from starlette.datastructures import MutableHeaders
|
|
from starlette.requests import Request
|
|
from starlette.responses import JSONResponse
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
from app.settings import settings
|
|
|
|
# POST .../run is the expensive endpoint (LFX execution, LLM cost):
|
|
# per client IP and minute.
|
|
RUN_PATH_SUFFIX = "/run"
|
|
THROTTLE_WINDOW_SECONDS = 60.0
|
|
|
|
_hits: dict[str, list[float]] = defaultdict(list)
|
|
|
|
|
|
def reset_throttle() -> None:
|
|
"""Drop all counters (tests)."""
|
|
_hits.clear()
|
|
|
|
|
|
def _client_ip(request: Request) -> str:
|
|
# Behind the TLS proxy ProxyHeadersMiddleware has already rewritten
|
|
# scope["client"] from X-Forwarded-For.
|
|
if request.client is not None:
|
|
return request.client.host
|
|
forwarded = request.headers.get("x-forwarded-for", "")
|
|
return forwarded.split(",")[0].strip() or "unknown"
|
|
|
|
|
|
class RunThrottleMiddleware:
|
|
"""Per-IP sliding window on POST .../run (429 + Retry-After)."""
|
|
|
|
def __init__(self, app: ASGIApp) -> None:
|
|
self.app = app
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if (
|
|
scope["type"] != "http"
|
|
or scope.get("method") != "POST"
|
|
or not scope.get("path", "").endswith(RUN_PATH_SUFFIX)
|
|
):
|
|
await self.app(scope, receive, send)
|
|
return
|
|
request = Request(scope)
|
|
now = time.monotonic()
|
|
window_start = now - THROTTLE_WINDOW_SECONDS
|
|
bucket = [t for t in _hits[_client_ip(request)] if t > window_start]
|
|
if len(bucket) >= settings.run_throttle_per_minute:
|
|
retry_after = max(1, int(bucket[0] + THROTTLE_WINDOW_SECONDS - now))
|
|
response = JSONResponse({"detail": "too many runs, try again later"}, status_code=429)
|
|
response.headers["Retry-After"] = str(retry_after)
|
|
await response(scope, receive, send)
|
|
return
|
|
bucket.append(now)
|
|
_hits[_client_ip(request)] = bucket
|
|
await self.app(scope, receive, send)
|
|
|
|
|
|
class SecurityHeadersMiddleware:
|
|
"""Baseline response headers (TLS/HSTS live on the proxy)."""
|
|
|
|
def __init__(self, app: ASGIApp) -> None:
|
|
self.app = app
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
async def secured(message: Message) -> None:
|
|
if message["type"] == "http.response.start":
|
|
headers = MutableHeaders(raw=message.setdefault("headers", []))
|
|
headers["X-Content-Type-Options"] = "nosniff"
|
|
headers["X-Frame-Options"] = "SAMEORIGIN"
|
|
headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
|
|
headers["Content-Security-Policy"] = "frame-ancestors 'self'"
|
|
await send(message)
|
|
|
|
await self.app(scope, receive, secured)
|