svc-infra 0.1.595__py3-none-any.whl → 0.1.706__py3-none-any.whl

svc_infra/api/fastapi/middleware/ratelimit.py

@@ -1,61 +1,158 @@
+import json
 import time
 
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.responses import JSONResponse
+from fastapi import Request
+from starlette.types import ASGIApp, Receive, Scope, Send
 
 from svc_infra.obs.metrics import emit_rate_limited
 
 from .ratelimit_store import InMemoryRateLimitStore, RateLimitStore
 
+try:
+    # Optional import: tenancy may not be enabled in all apps
+    from svc_infra.api.fastapi.tenancy.context import (
+        resolve_tenant_id as _resolve_tenant_id,
+    )
+except Exception:  # pragma: no cover - fallback for minimal builds
+    _resolve_tenant_id = None  # type: ignore[assignment]
+
+
+class SimpleRateLimitMiddleware:
+    """
+    Pure ASGI rate limiting middleware.
+
+    Applies per-key rate limits with configurable windows. Use skip_paths for
+    endpoints that should bypass rate limiting (e.g., health checks, webhooks).
+    """
 
-class SimpleRateLimitMiddleware(BaseHTTPMiddleware):
     def __init__(
         self,
-        app,
+        app: ASGIApp,
         limit: int = 120,
         window: int = 60,
         key_fn=None,
+        *,
+        # When provided, dynamically computes a limit for the current request (e.g. per-tenant quotas)
+        # Signature: (request: Request, tenant_id: Optional[str]) -> int | None
+        limit_resolver=None,
+        # If True, automatically scopes the bucket key by tenant id when available
+        scope_by_tenant: bool = False,
+        # When True, allows unresolved tenant IDs to fall back to an "X-Tenant-Id" header value.
+        # Disabled by default to avoid trusting arbitrary client-provided headers which could
+        # otherwise be used to evade per-tenant limits when authentication fails.
+        allow_untrusted_tenant_header: bool = False,
         store: RateLimitStore | None = None,
+        skip_paths: list[str] | None = None,
     ):
-        super().__init__(app)
+        self.app = app
         self.limit, self.window = limit, window
-        self.key_fn = key_fn or (lambda r: r.headers.get("X-API-Key") or r.client.host)
+        self.key_fn = key_fn
+        self._limit_resolver = limit_resolver
+        self.scope_by_tenant = scope_by_tenant
+        self._allow_untrusted_tenant_header = allow_untrusted_tenant_header
         self.store = store or InMemoryRateLimitStore(limit=limit)
+        self.skip_paths = skip_paths or []
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+
+        path = scope.get("path", "")
+
+        # Skip specified paths
+        if any(skip in path for skip in self.skip_paths):
+            await self.app(scope, receive, send)
+            return
+
+        # Create a Request object for key extraction and tenant resolution
+        request = Request(scope, receive)
+
+        # Default key function
+        key_fn = self.key_fn or (
+            lambda r: r.headers.get("X-API-Key")
+            or (r.client.host if r.client else "unknown")
+        )
+
+        # Resolve tenant when possible
+        tenant_id = None
+        if self.scope_by_tenant or self._limit_resolver:
+            try:
+                if _resolve_tenant_id is not None:
+                    tenant_id = await _resolve_tenant_id(request)
+            except Exception:
+                tenant_id = None
+            # Fallback header behavior - ONLY if explicitly allowed
+            # Never trust untrusted headers by default to prevent rate limit evasion
+            if not tenant_id and self._allow_untrusted_tenant_header:
+                tenant_id = request.headers.get("X-Tenant-Id") or request.headers.get(
+                    "X-Tenant-ID"
+                )
+
+        key = key_fn(request)
+        if self.scope_by_tenant and tenant_id:
+            key = f"{key}:tenant:{tenant_id}"
+
+        # Allow dynamic limit overrides
+        eff_limit = self.limit
+        if self._limit_resolver:
+            try:
+                v = self._limit_resolver(request, tenant_id)
+                eff_limit = int(v) if v is not None else self.limit
+            except Exception:
+                eff_limit = self.limit
 
-    async def dispatch(self, request, call_next):
-        key = self.key_fn(request)
         now = int(time.time())
-        # Increment counter in store
-        count, limit, reset = self.store.incr(str(key), self.window)
+        count, store_limit, reset = self.store.incr(str(key), self.window)
+        limit = eff_limit
         remaining = max(0, limit - count)
 
-        if remaining < 0:  # defensive clamp
-            remaining = 0
-
         if count > limit:
+            # Rate limited - return 429
             retry = max(0, reset - now)
             try:
                 emit_rate_limited(str(key), limit, retry)
             except Exception:
                 pass
-            return JSONResponse(
-                status_code=429,
-                content={
+
+            body = json.dumps(
+                {
                     "title": "Too Many Requests",
                     "status": 429,
                     "detail": "Rate limit exceeded.",
                     "code": "RATE_LIMITED",
-                },
-                headers={
-                    "X-RateLimit-Limit": str(limit),
-                    "X-RateLimit-Remaining": "0",
-                    "X-RateLimit-Reset": str(reset),
-                    "Retry-After": str(retry),
-                },
+                }
+            ).encode("utf-8")
+
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 429,
+                    "headers": [
+                        (b"content-type", b"application/json"),
+                        (b"x-ratelimit-limit", str(limit).encode()),
+                        (b"x-ratelimit-remaining", b"0"),
+                        (b"x-ratelimit-reset", str(reset).encode()),
+                        (b"retry-after", str(retry).encode()),
+                    ],
+                }
             )
+            await send({"type": "http.response.body", "body": body, "more_body": False})
+            return
+
+        # Not rate limited - add headers to response
+        async def send_with_headers(message):
+            if message["type"] == "http.response.start":
+                headers = list(message.get("headers", []))
+                # Add rate limit headers if not already present
+                header_names = {h[0].lower() for h in headers}
+                if b"x-ratelimit-limit" not in header_names:
+                    headers.append((b"x-ratelimit-limit", str(limit).encode()))
+                if b"x-ratelimit-remaining" not in header_names:
+                    headers.append((b"x-ratelimit-remaining", str(remaining).encode()))
+                if b"x-ratelimit-reset" not in header_names:
+                    headers.append((b"x-ratelimit-reset", str(reset).encode()))
+                message = {**message, "headers": headers}
+            await send(message)
 
-        resp = await call_next(request)
-        resp.headers.setdefault("X-RateLimit-Limit", str(limit))
-        resp.headers.setdefault("X-RateLimit-Remaining", str(remaining))
-        resp.headers.setdefault("X-RateLimit-Reset", str(reset))
-        return resp
+        await self.app(scope, receive, send_with_headers)

svc_infra/api/fastapi/middleware/ratelimit_store.py

@@ -1,7 +1,31 @@
 from __future__ import annotations
 
+import logging
+import os
 import time
-from typing import Optional, Protocol, Tuple
+import warnings
+from typing import Callable, Optional, Protocol, Tuple
+
+logger = logging.getLogger(__name__)
+
+_INMEMORY_WARNED = False
+
+
+def _check_inmemory_production_warning(class_name: str) -> None:
+    """Warn if in-memory store is used in production."""
+    global _INMEMORY_WARNED
+    if _INMEMORY_WARNED:
+        return
+    env = os.getenv("ENV", "development").lower()
+    if env in ("production", "staging", "prod"):
+        _INMEMORY_WARNED = True
+        msg = (
+            f"{class_name} is being used in {env} environment. "
+            "This is NOT suitable for production - data will be lost on restart. "
+            "Use RedisRateLimitStore instead."
+        )
+        warnings.warn(msg, RuntimeWarning, stacklevel=3)
+        logger.critical(msg)
 
 
 class RateLimitStore(Protocol):
@@ -15,15 +39,22 @@ class RateLimitStore(Protocol):
 
 class InMemoryRateLimitStore:
     def __init__(self, limit: int = 120):
+        _check_inmemory_production_warning("InMemoryRateLimitStore")
         self.limit = limit
-        self._buckets: dict[tuple[str, int], int] = {}
+        # Track per-key rolling windows: key -> (count, window_start_epoch)
+        self._state: dict[str, tuple[int, float]] = {}
 
     def incr(self, key: str, window: int) -> Tuple[int, int, int]:
-        now = int(time.time())
-        win = now - (now % window)
-        count = self._buckets.get((key, win), 0) + 1
-        self._buckets[(key, win)] = count
-        reset = win + window
+        now = time.time()
+        count, window_start = self._state.get(key, (0, now))
+        # If outside the rolling window, reset
+        if now >= window_start + window:
+            count = 1
+            window_start = now
+        else:
+            count += 1
+        self._state[key] = (count, window_start)
+        reset = int(window_start + window)
         return count, self.limit, reset
 
 
@@ -43,14 +74,14 @@ class RedisRateLimitStore:
         *,
         limit: int = 120,
         prefix: str = "ratelimit",
-        clock: Optional[callable] = None,
+        clock: Optional[Callable[[], float]] = None,
     ):
         self.redis = redis_client
         self.limit = limit
         self.prefix = prefix
         self._clock = clock or time.time
 
-    def _window_key(self, key: str, window: int) -> tuple[str, int, str]:
+    def _window_key(self, key: str, window: int) -> tuple[str, int, int]:
         now = int(self._clock())
         win = now - (now % window)
         redis_key = f"{self.prefix}:{key}:{win}"
@@ -63,7 +94,9 @@ class RedisRateLimitStore:
         pipe.incr(rkey)
         pipe.ttl(rkey)
         count, ttl = pipe.execute()
-        if ttl == -1:  # key exists without expire or just created; set expire to end of window
+        if (
+            ttl == -1
+        ):  # key exists without expire or just created; set expire to end of window
             expire_sec = (win + window) - now
             if expire_sec <= 0:
                 expire_sec = window

svc_infra/api/fastapi/middleware/request_id.py

@@ -1,23 +1,39 @@
 import contextvars
 from uuid import uuid4
 
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.types import ASGIApp
+from starlette.datastructures import Headers, MutableHeaders
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
-request_id_ctx: contextvars.ContextVar[str] = contextvars.ContextVar("request_id", default="")
+request_id_ctx: contextvars.ContextVar[str] = contextvars.ContextVar(
+    "request_id", default=""
+)
 
 
-class RequestIdMiddleware(BaseHTTPMiddleware):
+class RequestIdMiddleware:
+    """Pure ASGI middleware that adds request IDs. Compatible with streaming responses."""
+
     def __init__(self, app: ASGIApp, header_name: str = "X-Request-Id"):
-        super().__init__(app)
-        self.header_name = header_name
+        self.app = app
+        self.header_name = header_name.lower()
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
 
-    async def dispatch(self, request, call_next):
-        rid = request.headers.get(self.header_name) or uuid4().hex
+        # Extract or generate request ID
+        headers = Headers(scope=scope)
+        rid = headers.get(self.header_name) or uuid4().hex
         token = request_id_ctx.set(rid)
+
+        async def send_with_request_id(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                # Add request ID to response headers
+                response_headers = MutableHeaders(scope=message)
+                response_headers.append(self.header_name, rid)
+            await send(message)
+
         try:
-            resp = await call_next(request)
-            resp.headers[self.header_name] = rid
-            return resp
+            await self.app(scope, receive, send_with_request_id)
         finally:
             request_id_ctx.reset(token)

svc_infra/api/fastapi/middleware/request_size_limit.py

@@ -19,9 +19,9 @@ class RequestSizeLimitMiddleware(BaseHTTPMiddleware):
             size = None
         if size is not None and size > self.max_bytes:
             try:
-                emit_suspect_payload(
-                    getattr(request, "url", None).path if hasattr(request, "url") else None, size
-                )
+                url = getattr(request, "url", None)
+                path = url.path if url is not None else None
+                emit_suspect_payload(path, size)
             except Exception:
                 pass
             return JSONResponse(

svc_infra/api/fastapi/middleware/timeout.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import Any
+
+from fastapi import Request
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from svc_infra.api.fastapi.middleware.errors.handlers import problem_response
+from svc_infra.app.env import pick
+
+
+def _env_int(name: str, default: int) -> int:
+    v = os.getenv(name)
+    if v is None:
+        return default
+    try:
+        return int(v)
+    except Exception:
+        return default
+
+
+REQUEST_BODY_TIMEOUT_SECONDS: int = pick(
+    prod=_env_int("REQUEST_BODY_TIMEOUT_SECONDS", 15),
+    nonprod=_env_int("REQUEST_BODY_TIMEOUT_SECONDS", 30),
+)
+REQUEST_TIMEOUT_SECONDS: int = pick(
+    prod=_env_int("REQUEST_TIMEOUT_SECONDS", 30),
+    nonprod=_env_int("REQUEST_TIMEOUT_SECONDS", 15),
+)
+
+
+class HandlerTimeoutMiddleware:
+    """
+    Caps total handler execution time. If exceeded, returns 504 Problem+JSON.
+
+    Use skip_paths for endpoints that may run longer than the timeout
+    (e.g., streaming responses, long-polling, file uploads).
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        timeout_seconds: int | None = None,
+        skip_paths: list[str] | None = None,
+    ) -> None:
+        self.app = app
+        self.timeout_seconds = (
+            timeout_seconds if timeout_seconds is not None else REQUEST_TIMEOUT_SECONDS
+        )
+        self.skip_paths = skip_paths or []
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+
+        path = scope.get("path", "")
+
+        # Skip specified paths (e.g., long-running endpoints)
+        if any(skip in path for skip in self.skip_paths):
+            await self.app(scope, receive, send)
+            return
+
+        # Track if response has started (headers sent)
+        response_started = False
+
+        async def send_wrapper(message: dict) -> None:
+            nonlocal response_started
+            if message.get("type") == "http.response.start":
+                response_started = True
+            await send(message)
+
+        try:
+            await asyncio.wait_for(
+                self.app(scope, receive, send_wrapper),  # type: ignore[arg-type]  # ASGI send signature
+                timeout=self.timeout_seconds,
+            )
+        except asyncio.TimeoutError:
+            # Only send 504 if response hasn't started yet
+            if not response_started:
+                response = problem_response(
+                    status=504,
+                    title="Gateway Timeout",
+                    detail=f"Handler did not complete within {self.timeout_seconds}s",
+                )
+                await response(scope, receive, send)
+            # If response already started, we can't change it - just let it fail
+
+
+class BodyReadTimeoutMiddleware:
+    """
+    Enforces a timeout while reading the request body to mitigate slowloris.
+    If body read does not make progress within the timeout, returns 408 Problem+JSON.
+    """
+
+    def __init__(self, app: ASGIApp, timeout_seconds: int | None = None) -> None:
+        self.app = app
+        self.timeout_seconds = (
+            timeout_seconds
+            if timeout_seconds is not None
+            else REQUEST_BODY_TIMEOUT_SECONDS
+        )
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+
+        # Strategy: greedily drain the incoming request body here while enforcing
+        # per-receive timeout, then replay it to the downstream app from a buffer.
+        # This ensures we can detect slowloris-style uploads even if the app only
+        # reads the body later (after the server has finished buffering).
+        buffered = bytearray()
+
+        try:
+            while True:
+                message = await asyncio.wait_for(
+                    receive(), timeout=self.timeout_seconds
+                )
+
+                mtype = message.get("type")
+                if mtype == "http.request":
+                    chunk = message.get("body", b"") or b""
+                    if chunk:
+                        buffered.extend(chunk)
+                    # Stop when server indicates no more body
+                    if not message.get("more_body", False):
+                        break
+                    # else: continue reading remaining chunks with timeout
+                    continue
+
+                if mtype == "http.disconnect":  # client disconnected mid-upload
+                    # Treat as end of body for the purposes of replay; downstream
+                    # will see an empty body. No timeout response needed here.
+                    break
+                # Ignore other message types and continue
+        except asyncio.TimeoutError:
+            # Timed out while waiting for the next body chunk → return 408
+            request = Request(scope, receive=receive)
+            trace_id = None
+            for h in ("x-request-id", "x-correlation-id", "x-trace-id"):
+                v = request.headers.get(h)
+                if v:
+                    trace_id = v
+                    break
+            resp = problem_response(
+                status=408,
+                title="Request Timeout",
+                detail="Timed out while reading request body.",
+                code="REQUEST_TIMEOUT",
+                instance=str(request.url),
+                trace_id=trace_id,
+            )
+            await resp(scope, receive, send)
+            return
+
+        # Replay the drained body to the app as a single http.request message.
+        # IMPORTANT: After replaying the body, we must forward the original receive()
+        # so that Starlette's listen_for_disconnect can properly detect client disconnects.
+        # This is required for streaming responses on ASGI spec < 2.4.
+        body_sent = False
+
+        async def _replay_receive() -> dict[str, Any]:
+            nonlocal body_sent
+            if not body_sent:
+                body_sent = True
+                return {
+                    "type": "http.request",
+                    "body": bytes(buffered),
+                    "more_body": False,
+                }
+            # After body is sent, forward to original receive for disconnect detection
+            return dict(await receive())
+
+        await self.app(scope, _replay_receive, send)

svc_infra/api/fastapi/openapi/apply.py

@@ -5,7 +5,9 @@ from typing import Any, Callable
 from fastapi import APIRouter
 
 
-def apply_default_security(router: APIRouter, *, default_security: list[dict] | None) -> None:
+def apply_default_security(
+    router: APIRouter, *, default_security: list[dict] | None
+) -> None:
     if default_security is None:
         return
     original_add = router.add_api_route
@@ -17,7 +19,7 @@ def apply_default_security(router: APIRouter, *, default_security: list[dict] |
             kwargs["openapi_extra"] = ox
         return original_add(path, endpoint, **kwargs)
 
-    router.add_api_route = _wrapped_add_api_route  # type: ignore[attr-defined]
+    setattr(router, "add_api_route", _wrapped_add_api_route)
 
 
 def apply_default_responses(router: APIRouter, defaults: dict[int, dict]) -> None:
@@ -38,4 +40,4 @@ def apply_default_responses(router: APIRouter, defaults: dict[int, dict]) -> Non
         kwargs["responses"] = responses
         return original_add(path, endpoint, **kwargs)
 
-    router.add_api_route = _wrapped_add_api_route  # type: ignore[attr-defined]
+    setattr(router, "add_api_route", _wrapped_add_api_route)

svc_infra/api/fastapi/openapi/conventions.py

@@ -16,7 +16,11 @@ PROBLEM_SCHEMA: Dict[str, Any] = {
             "description": "URI identifying the error type",
         },
         "title": {"type": "string", "description": "Short, human-readable summary"},
-        "status": {"type": "integer", "format": "int32", "description": "HTTP status code"},
+        "status": {
+            "type": "integer",
+            "format": "int32",
+            "description": "HTTP status code",
+        },
         "detail": {"type": "string", "description": "Human-readable explanation"},
         "instance": {
             "type": "string",
@@ -36,7 +40,10 @@ PROBLEM_SCHEMA: Dict[str, Any] = {
                 },
             },
         },
-        "trace_id": {"type": "string", "description": "Correlation/trace id (if available)"},
+        "trace_id": {
+            "type": "string",
+            "description": "Correlation/trace id (if available)",
+        },
     },
     "required": ["title", "status"],
 }