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

svc_infra/api/fastapi/middleware/ratelimit_store.py

@@ -1,11 +1,36 @@
 from __future__ import annotations
 
+import logging
+import os
 import time
-from typing import Optional, Protocol, Tuple
+import warnings
+from collections.abc import Callable
+from typing import Protocol
+
+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):
-    def incr(self, key: str, window: int) -> Tuple[int, int, int]:
+    def incr(self, key: str, window: int) -> tuple[int, int, int]:
         """Increment and return (count, limit, resetEpoch).
 
         Implementations should manage per-window buckets. The 'limit' is stored configuration.
@@ -15,15 +40,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] = {}
-
-    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
+        # 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 = 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,20 +75,20 @@ class RedisRateLimitStore:
         *,
         limit: int = 120,
         prefix: str = "ratelimit",
-        clock: Optional[callable] = None,
+        clock: Callable[[], float] | None = 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}"
         return redis_key, win, now
 
-    def incr(self, key: str, window: int) -> Tuple[int, int, int]:
+    def incr(self, key: str, window: int) -> tuple[int, int, int]:
         rkey, win, now = self._window_key(key, window)
         # Increment; if this is the first time we've seen this window key, set expiry to window end
         pipe = self.redis.pipeline()

svc_infra/api/fastapi/middleware/request_id.py

@@ -1,23 +1,37 @@
 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="")
 
 
-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,176 @@
+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).
+
+    Matching uses prefix matching: "/v1/chat" matches "/v1/chat", "/v1/chat/stream",
+    but not "/api/v1/chat" or "/v1/chatter".
+    """
+
+    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 using prefix matching (e.g., long-running endpoints)
+        if any(path.startswith(skip) 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 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 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)