svc-infra 0.1.506__py3-none-any.whl → 0.1.654__py3-none-any.whl

svc_infra/api/fastapi/middleware/idempotency_store.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import base64
+import json
+import time
+from dataclasses import dataclass
+from typing import Dict, Optional, Protocol
+
+
+@dataclass
+class IdempotencyEntry:
+    req_hash: str
+    exp: float
+    # Optional response fields when available
+    status: Optional[int] = None
+    body_b64: Optional[str] = None
+    headers: Optional[Dict[str, str]] = None
+    media_type: Optional[str] = None
+
+
+class IdempotencyStore(Protocol):
+    def get(self, key: str) -> Optional[IdempotencyEntry]:
+        pass
+
+    def set_initial(self, key: str, req_hash: str, exp: float) -> bool:
+        """Atomically create an entry if absent. Returns True if created, False if already exists."""
+        pass
+
+    def set_response(
+        self,
+        key: str,
+        *,
+        status: int,
+        body: bytes,
+        headers: Dict[str, str],
+        media_type: Optional[str],
+    ) -> None:
+        pass
+
+    def delete(self, key: str) -> None:
+        pass
+
+
+class InMemoryIdempotencyStore:
+    def __init__(self):
+        self._store: dict[str, IdempotencyEntry] = {}
+
+    def get(self, key: str) -> Optional[IdempotencyEntry]:
+        entry = self._store.get(key)
+        if not entry:
+            return None
+        # expire lazily
+        if entry.exp <= time.time():
+            self._store.pop(key, None)
+            return None
+        return entry
+
+    def set_initial(self, key: str, req_hash: str, exp: float) -> bool:
+        now = time.time()
+        existing = self._store.get(key)
+        if existing and existing.exp > now:
+            return False
+        self._store[key] = IdempotencyEntry(req_hash=req_hash, exp=exp)
+        return True
+
+    def set_response(
+        self,
+        key: str,
+        *,
+        status: int,
+        body: bytes,
+        headers: Dict[str, str],
+        media_type: Optional[str],
+    ) -> None:
+        entry = self._store.get(key)
+        if not entry:
+            # Create if missing to ensure replay works until exp
+            entry = IdempotencyEntry(req_hash="", exp=time.time() + 60)
+            self._store[key] = entry
+        entry.status = status
+        entry.body_b64 = base64.b64encode(body).decode()
+        entry.headers = dict(headers)
+        entry.media_type = media_type
+
+    def delete(self, key: str) -> None:
+        self._store.pop(key, None)
+
+
+class RedisIdempotencyStore:
+    """A simple Redis-backed store.
+
+    Notes:
+        - Uses GET/SET with JSON payload; initial claim uses SETNX semantics.
+        - Not fully atomic for response update; sufficient for basic dedupe.
+        - For strict guarantees, replace with a Lua script (future improvement).
+    """
+
+    def __init__(self, redis_client, *, prefix: str = "idmp"):
+        self.r = redis_client
+        self.prefix = prefix
+
+    def _k(self, key: str) -> str:
+        return f"{self.prefix}:{key}"
+
+    def get(self, key: str) -> Optional[IdempotencyEntry]:
+        raw = self.r.get(self._k(key))
+        if not raw:
+            return None
+        try:
+            data = json.loads(raw)
+        except Exception:
+            return None
+        entry = IdempotencyEntry(
+            req_hash=data.get("req_hash", ""),
+            exp=float(data.get("exp", 0)),
+            status=data.get("status"),
+            body_b64=data.get("body_b64"),
+            headers=data.get("headers"),
+            media_type=data.get("media_type"),
+        )
+        if entry.exp <= time.time():
+            try:
+                self.r.delete(self._k(key))
+            except Exception:
+                pass
+            return None
+        return entry
+
+    def set_initial(self, key: str, req_hash: str, exp: float) -> bool:
+        payload = json.dumps({"req_hash": req_hash, "exp": exp})
+        # Attempt NX set
+        ok = self.r.set(self._k(key), payload, nx=True)
+        # If set, also set TTL (expire at exp)
+        if ok:
+            ttl = max(1, int(exp - time.time()))
+            try:
+                self.r.expire(self._k(key), ttl)
+            except Exception:
+                pass
+            return True
+        # If exists but expired, overwrite
+        entry = self.get(key)
+        if not entry:
+            self.r.set(self._k(key), payload)
+            ttl = max(1, int(exp - time.time()))
+            try:
+                self.r.expire(self._k(key), ttl)
+            except Exception:
+                pass
+            return True
+        return False
+
+    def set_response(
+        self,
+        key: str,
+        *,
+        status: int,
+        body: bytes,
+        headers: Dict[str, str],
+        media_type: Optional[str],
+    ) -> None:
+        entry = self.get(key)
+        if not entry:
+            # default short ttl if missing; caller should have set initial
+            entry = IdempotencyEntry(req_hash="", exp=time.time() + 60)
+        entry.status = status
+        entry.body_b64 = base64.b64encode(body).decode()
+        entry.headers = dict(headers)
+        entry.media_type = media_type
+        ttl = max(1, int(entry.exp - time.time()))
+        payload = json.dumps(
+            {
+                "req_hash": entry.req_hash,
+                "exp": entry.exp,
+                "status": entry.status,
+                "body_b64": entry.body_b64,
+                "headers": entry.headers,
+                "media_type": entry.media_type,
+            }
+        )
+        self.r.set(self._k(key), payload, ex=ttl)
+
+    def delete(self, key: str) -> None:
+        try:
+            self.r.delete(self._k(key))
+        except Exception:
+            pass

svc_infra/api/fastapi/middleware/optimistic_lock.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import Annotated, Any, Callable, Optional
+
+from fastapi import Header, HTTPException
+
+
+async def require_if_match(
+    version: Annotated[Optional[str], Header(alias="If-Match")] = None
+) -> str:
+    """Require If-Match header for optimistic locking on mutating operations.
+
+    Returns the header value. Raises 428 if missing.
+    """
+    if not version:
+        raise HTTPException(
+            status_code=428, detail="Missing If-Match header for optimistic locking."
+        )
+    return version
+
+
+def check_version_or_409(get_current_version: Callable[[], Any], provided: str) -> None:
+    """Compare provided version with current version; raise 409 on mismatch.
+
+    - get_current_version: callable returning the resource's current version (int/str)
+    - provided: header value; attempts to coerce to int if current is int
+    """
+    current = get_current_version()
+    if isinstance(current, int):
+        try:
+            p = int(provided)
+        except Exception:
+            raise HTTPException(status_code=400, detail="Invalid If-Match value; expected integer.")
+    else:
+        p = provided
+    if p != current:
+        raise HTTPException(status_code=409, detail="Version mismatch (optimistic locking).")

svc_infra/api/fastapi/middleware/ratelimit.py

@@ -0,0 +1,119 @@
+import time
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import JSONResponse
+
+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
+
+
+class SimpleRateLimitMiddleware(BaseHTTPMiddleware):
+    def __init__(
+        self,
+        app,
+        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,
+    ):
+        super().__init__(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._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)
+
+    async def dispatch(self, request, call_next):
+        # 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:
+            # - If tenancy context is unavailable (minimal builds), accept header by default so
+            #   unit/integration tests can exercise per-tenant scoping without full auth state.
+            # - If tenancy is available, only trust the header when explicitly allowed.
+            if not tenant_id:
+                if _resolve_tenant_id is None:
+                    tenant_id = request.headers.get("X-Tenant-Id") or request.headers.get(
+                        "X-Tenant-ID"
+                    )
+                elif self._allow_untrusted_tenant_header:
+                    tenant_id = request.headers.get("X-Tenant-Id") or request.headers.get(
+                        "X-Tenant-ID"
+                    )
+
+        key = self.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
+
+        now = int(time.time())
+        # Increment counter in store
+        # Update store limit if it differs; stores capture configured limit internally
+        # For in-memory store, we can temporarily adjust per-request by swapping a new store instance
+        # but to keep API simple, we reuse store and clamp by eff_limit below.
+        count, store_limit, reset = self.store.incr(str(key), self.window)
+        # Enforce the effective limit selected for this request
+        limit = eff_limit
+        remaining = max(0, limit - count)
+
+        if remaining < 0:  # defensive clamp
+            remaining = 0
+
+        if count > limit:
+            retry = max(0, reset - now)
+            try:
+                emit_rate_limited(str(key), limit, retry)
+            except Exception:
+                pass
+            return JSONResponse(
+                status_code=429,
+                content={
+                    "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),
+                },
+            )
+
+        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

svc_infra/api/fastapi/middleware/ratelimit_store.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import time
+from typing import Optional, Protocol, Tuple
+
+
+class RateLimitStore(Protocol):
+    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.
+        """
+        ...
+
+
+class InMemoryRateLimitStore:
+    def __init__(self, limit: int = 120):
+        self.limit = limit
+        # 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
+
+
+class RedisRateLimitStore:
+    """Fixed-window counter store using Redis.
+
+    Keys are of the form: {prefix}:{key}:{windowStart}
+    Values are incremented and expire automatically at window end.
+
+    This implementation uses atomic INCR and EXPIRE semantics. To avoid race conditions
+    on first-set expiry, we set expiry when the counter is created.
+    """
+
+    def __init__(
+        self,
+        redis_client,
+        *,
+        limit: int = 120,
+        prefix: str = "ratelimit",
+        clock: Optional[callable] = 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]:
+        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]:
+        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()
+        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
+            expire_sec = (win + window) - now
+            if expire_sec <= 0:
+                expire_sec = window
+            try:
+                self.redis.expire(rkey, expire_sec)
+            except Exception:
+                pass
+        reset = win + window
+        return int(count), self.limit, reset
+
+
+__all__ = ["RateLimitStore", "InMemoryRateLimitStore", "RedisRateLimitStore"]

svc_infra/api/fastapi/middleware/request_id.py

@@ -0,0 +1,23 @@
+import contextvars
+from uuid import uuid4
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp
+
+request_id_ctx: contextvars.ContextVar[str] = contextvars.ContextVar("request_id", default="")
+
+
+class RequestIdMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: ASGIApp, header_name: str = "X-Request-Id"):
+        super().__init__(app)
+        self.header_name = header_name
+
+    async def dispatch(self, request, call_next):
+        rid = request.headers.get(self.header_name) or uuid4().hex
+        token = request_id_ctx.set(rid)
+        try:
+            resp = await call_next(request)
+            resp.headers[self.header_name] = rid
+            return resp
+        finally:
+            request_id_ctx.reset(token)

svc_infra/api/fastapi/middleware/request_size_limit.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import JSONResponse
+
+from svc_infra.obs.metrics import emit_suspect_payload
+
+
+class RequestSizeLimitMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app, max_bytes: int = 1_000_000):
+        super().__init__(app)
+        self.max_bytes = max_bytes
+
+    async def dispatch(self, request, call_next):
+        length = request.headers.get("content-length")
+        try:
+            size = int(length) if length is not None else None
+        except Exception:
+            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
+                )
+            except Exception:
+                pass
+            return JSONResponse(
+                status_code=413,
+                content={
+                    "title": "Payload Too Large",
+                    "status": 413,
+                    "detail": "Request body exceeds allowed size.",
+                    "code": "PAYLOAD_TOO_LARGE",
+                },
+            )
+        return await call_next(request)

svc_infra/api/fastapi/middleware/timeout.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import asyncio
+import os
+
+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.
+    """
+
+    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_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
+
+        async def _call_next() -> None:
+            await self.app(scope, receive, send)
+
+        try:
+            await asyncio.wait_for(_call_next(), timeout=self.timeout_seconds)
+        except asyncio.TimeoutError:
+            # Build a minimal Request to extract headers and URL for trace info
+            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=504,
+                title="Gateway Timeout",
+                detail="The request took too long to complete.",
+                code="GATEWAY_TIMEOUT",
+                instance=str(request.url),
+                trace_id=trace_id,
+            )
+            await resp(scope, receive, send)
+
+
+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.
+        sent = False
+
+        async def _replay_receive() -> dict:
+            nonlocal sent
+            if not sent:
+                sent = True
+                return {"type": "http.request", "body": bytes(buffered), "more_body": False}
+            # Subsequent calls return an empty terminal body event
+            return {"type": "http.request", "body": b"", "more_body": False}
+
+        await self.app(scope, _replay_receive, send)