10xgraph-api 0.7.0__py3-none-any.whl

tenxgraph_api/src/app/core/middleware/rate_limit/keying.py

@@ -0,0 +1,133 @@
+"""Shared client-key derivation for rate limiting.
+
+Used by both the HTTP ``RateLimitMiddleware`` and the WebSocket connection guard so that
+WebSocket handshakes are counted against the *same* rate-limit bucket as REST requests.
+Works on any ``HTTPConnection`` (both ``Request`` and ``WebSocket`` subclass it).
+"""
+
+import ipaddress
+import logging
+
+from starlette.requests import HTTPConnection
+from starlette.responses import Response
+
+from tenxgraph_api.src.app.core.config.graph_config import RateLimitConfig
+
+
+logger = logging.getLogger("tenxgraph_api")
+
+
+def _client_ip(connection: HTTPConnection, config: RateLimitConfig) -> str:
+    """Resolve the client IP, honouring X-Forwarded-For only where it is trustworthy.
+
+    ``X-Forwarded-For`` is a list that each proxy APPENDS to. Whatever the caller
+    sent arrives at the left; only the entries your own proxies appended, on the
+    right, are trustworthy.
+
+    Reading ``split(",")[0]`` -- the previous behaviour -- therefore took a value
+    the caller fully controls. An attacker could send a different
+    ``X-Forwarded-For`` on every request, land in a fresh bucket each time, and
+    never be rate limited at all.
+
+    We instead count ``trusted_proxy_hops`` back from the RIGHT. With one proxy in
+    front (the default), the last entry is the address that proxy actually observed
+    -- the real peer.
+
+    All ``X-Forwarded-For`` header lines are combined in order: a proxy that adds its own
+    line instead of appending to the caller's must not let the caller's line decide. When
+    ``trusted_proxies`` is configured, the header is only honoured for connections that come
+    from one of those networks; anyone reaching the app directly is keyed by their peer.
+    """
+    peer = connection.client.host if connection.client else "unknown"
+    if config.trusted_proxy_headers and _from_trusted_proxy(peer, config):
+        lines = connection.headers.getlist("X-Forwarded-For")
+        if lines:
+            parts = [p.strip() for line in lines for p in line.split(",") if p.strip()]
+            hops = max(1, getattr(config, "trusted_proxy_hops", 1))
+            if len(parts) >= hops:
+                return parts[-hops]
+            # Fewer entries than our own infrastructure should have appended: the
+            # header is not shaped the way we expect, so trust none of it.
+            logger.warning(
+                "X-Forwarded-For has %d entries but %d trusted proxy hop(s) are "
+                "configured; ignoring the header and using the peer address.",
+                len(parts),
+                hops,
+            )
+
+    return peer
+
+
+def _from_trusted_proxy(peer: str, config: RateLimitConfig) -> bool:
+    trusted = getattr(config, "trusted_proxies", ())
+    if not trusted:
+        return True  # no list configured: trust the header (the hop count still applies)
+    try:
+        address = ipaddress.ip_address(peer)
+    except ValueError:
+        return False
+    return any(address in ipaddress.ip_network(net) for net in trusted)
+
+
+def verified_user_id(connection: HTTPConnection) -> str | None:
+    """Identify the caller from their credential, exactly as ``RequirePermission`` does.
+
+    The rate limiter runs before any route dependency, so it cannot read an authenticated
+    user from request state. It verifies the credential itself with the configured auth
+    backend instead. Returns ``None`` when auth is not configured, no credential was sent, or
+    the credential does not verify; the caller is then limited by address, so forged tokens
+    cannot buy fresh buckets.
+    """
+    from injectq import InjectQ
+
+    from tenxgraph_api.src.app.core.auth.base_auth import BaseAuth
+    from tenxgraph_api.src.app.core.auth.permissions import _extract_credential
+    from tenxgraph_api.src.app.core.config.graph_config import GraphConfig
+
+    credential = _extract_credential(connection)
+    if credential is None:
+        return None
+    try:
+        container = InjectQ.get_instance()
+        graph_config = container.try_get(GraphConfig)
+        if graph_config is None or not graph_config.auth_config():
+            return None
+        auth_backend = container.try_get(BaseAuth)
+        if auth_backend is None:
+            return None
+        # A throwaway Response: anything the backend sets on it belongs to the real
+        # authentication in RequirePermission, which runs again for the route.
+        user = auth_backend.authenticate(connection, Response(), credential)
+    except Exception as exc:  # invalid token, backend error: fall back to the address
+        logger.debug("Rate limit could not verify caller identity: %s", exc)
+        return None
+    user_id = user.get("user_id") if isinstance(user, dict) else None
+    return str(user_id) if user_id else None
+
+
+def client_key_for(connection: HTTPConnection, config: RateLimitConfig) -> str:
+    """Derive the rate-limit bucket key for a connection.
+
+    Supports ``by``:
+
+    - ``"global"``: a single bucket for the whole service.
+    - ``"user"``: one bucket per authenticated user. This is what you want once auth
+      is enabled: limiting purely by IP gives one user roaming across addresses an
+      unlimited budget, while a NAT'd office sharing one address gets throttled as
+      though it were a single caller.
+    - ``"ip"``: one bucket per client address (the default).
+
+    ``"user"`` falls back to the peer address when there is no authenticated user,
+    so anonymous traffic is still limited per-caller rather than sharing one bucket
+    that any single client could exhaust for everybody.
+    """
+    if config.by == "global":
+        return "__global__"
+
+    if config.by == "user":
+        user_id = verified_user_id(connection)
+        if user_id:
+            return f"user:{user_id}"
+        return f"ip:{_client_ip(connection, config)}"
+
+    return _client_ip(connection, config)

tenxgraph_api/src/app/core/middleware/rate_limit/memory.py

@@ -0,0 +1,98 @@
+import asyncio
+import logging
+import time
+from collections import deque
+
+from .base import BaseRateLimitBackend, RateLimitDecision
+
+
+logger = logging.getLogger("tenxgraph_api.rate_limit")
+
+_NUM_STRIPES = 64
+
+
+class MemoryRateLimitBackend(BaseRateLimitBackend):
+    """In-process sliding-window rate limiter.
+
+    This is useful for development and single-process deployments. It is not
+    shared across workers or containers; use Redis for distributed enforcement.
+    """
+
+    _SWEEP_INTERVAL = 2000
+
+    def __init__(self, max_unique_keys: int = 50_000) -> None:
+        self._buckets: dict[str, deque[float]] = {}
+        self._last_seen: dict[str, float] = {}
+        self._locks = [asyncio.Lock() for _ in range(_NUM_STRIPES)]
+        self._sweep_lock = asyncio.Lock()
+        self._bg_tasks: set[asyncio.Task] = set()
+        self._check_count = 0
+        self._max_unique_keys = max_unique_keys
+
+    def _stripe_lock(self, key: str) -> asyncio.Lock:
+        return self._locks[hash(key) % _NUM_STRIPES]
+
+    def _schedule_sweep(self, window: int) -> None:
+        task = asyncio.ensure_future(self._sweep(window))
+        self._bg_tasks.add(task)
+        task.add_done_callback(self._bg_tasks.discard)
+
+    def _evict_oldest_bucket(self) -> None:
+        if not self._last_seen:
+            return
+        oldest_key = min(self._last_seen, key=self._last_seen.__getitem__)
+        self._buckets.pop(oldest_key, None)
+        self._last_seen.pop(oldest_key, None)
+
+    async def _sweep(self, window: int) -> None:
+        async with self._sweep_lock:
+            cutoff = time.monotonic() - window * 2
+            stale = [k for k, ts in self._last_seen.items() if ts < cutoff]
+            for key in stale:
+                self._buckets.pop(key, None)
+                self._last_seen.pop(key, None)
+            if stale:
+                logger.debug("Rate-limit memory sweep removed %d stale buckets", len(stale))
+
+    async def check(self, key: str, *, limit: int, window: int) -> RateLimitDecision:
+        now = time.monotonic()
+        window_start = now - window
+
+        self._check_count += 1
+        if self._check_count % self._SWEEP_INTERVAL == 0:
+            self._schedule_sweep(window)
+
+        async with self._stripe_lock(key):
+            if key not in self._buckets and len(self._buckets) >= self._max_unique_keys:
+                logger.warning(
+                    "Rate-limit bucket cap (%d) reached; running emergency sweep",
+                    self._max_unique_keys,
+                )
+                await self._sweep(window)
+                if len(self._buckets) >= self._max_unique_keys:
+                    self._evict_oldest_bucket()
+
+            bucket = self._buckets.setdefault(key, deque())
+            self._last_seen[key] = now
+
+            while bucket and bucket[0] < window_start:
+                bucket.popleft()
+
+            count = len(bucket)
+            if count >= limit:
+                reset_after = max(1, int(bucket[0] + window - now) + 1)
+                return RateLimitDecision(allowed=False, remaining=0, reset_after=reset_after)
+
+            bucket.append(now)
+            return RateLimitDecision(
+                allowed=True,
+                remaining=limit - len(bucket),
+                reset_after=window,
+            )
+
+    async def close(self) -> None:
+        for task in self._bg_tasks:
+            task.cancel()
+        if self._bg_tasks:
+            await asyncio.gather(*self._bg_tasks, return_exceptions=True)
+        self._bg_tasks.clear()

tenxgraph_api/src/app/core/middleware/rate_limit/middleware.py

@@ -0,0 +1,90 @@
+import time
+
+from fastapi import status
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+
+from tenxgraph_api.src.app.core import logger
+from tenxgraph_api.src.app.core.config.graph_config import RateLimitConfig
+
+from .base import BaseRateLimitBackend
+from .keying import client_key_for
+
+
+# Liveness/readiness probes hit /ping. Limiting it (or failing closed on a backend outage)
+# would make the orchestrator restart healthy pods.
+HEALTH_CHECK_PATHS = frozenset({"/ping"})
+
+
+class RateLimitMiddleware(BaseHTTPMiddleware):
+    """Backend-agnostic rate-limit middleware."""
+
+    def __init__(
+        self,
+        app,
+        config: RateLimitConfig,
+        backend: BaseRateLimitBackend,
+    ) -> None:
+        super().__init__(app)
+        self.config = config
+        self.backend = backend
+        self._exclude = frozenset(config.exclude_paths) | HEALTH_CHECK_PATHS
+
+    def _client_key(self, request: Request) -> str:
+        return client_key_for(request, self.config)
+
+    async def dispatch(self, request: Request, call_next):
+        if request.url.path in self._exclude:
+            return await call_next(request)
+
+        key = self._client_key(request)
+        decision = await self.backend.check(
+            key,
+            limit=self.config.requests,
+            window=self.config.window,
+        )
+        reset_at_epoch = int(time.time()) + decision.reset_after
+
+        if not decision.allowed:
+            request_id = getattr(request.state, "request_id", "unknown")
+            logger.warning(
+                "Rate limit exceeded for %s on %s %s",
+                key,
+                request.method,
+                request.url.path,
+            )
+            return JSONResponse(
+                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+                content={
+                    "error": {
+                        "code": "RATE_LIMIT_EXCEEDED",
+                        "message": (
+                            f"Too many requests. Limit: {self.config.requests} "
+                            f"per {self.config.window}s. "
+                            f"Retry after {decision.reset_after}s."
+                        ),
+                        "limit": self.config.requests,
+                        "window_seconds": self.config.window,
+                        "retry_after_seconds": decision.reset_after,
+                    },
+                    "metadata": {
+                        "request_id": request_id,
+                        "status": "error",
+                    },
+                },
+                headers={
+                    "Retry-After": str(decision.reset_after),
+                    "X-RateLimit-Limit": str(self.config.requests),
+                    "X-RateLimit-Remaining": "0",
+                    "X-RateLimit-Reset": str(reset_at_epoch),
+                    "X-RateLimit-Reset-After": str(decision.reset_after),
+                },
+            )
+
+        response = await call_next(request)
+        response.headers["X-RateLimit-Limit"] = str(self.config.requests)
+        response.headers["X-RateLimit-Remaining"] = str(decision.remaining)
+        response.headers["X-RateLimit-Reset"] = str(reset_at_epoch)
+        response.headers["X-RateLimit-Reset-After"] = str(decision.reset_after)
+        return response

tenxgraph_api/src/app/core/middleware/rate_limit/redis.py

@@ -0,0 +1,115 @@
+import logging
+import time
+import uuid
+from typing import Any
+
+from .base import BaseRateLimitBackend, RateLimitDecision
+
+
+try:
+    from redis.asyncio import Redis as AsyncRedis  # type: ignore[import]
+
+    _REDIS_AVAILABLE = True
+except ImportError:
+    _REDIS_AVAILABLE = False
+    AsyncRedis = None  # type: ignore[assignment,misc]
+
+
+logger = logging.getLogger("tenxgraph_api.rate_limit")
+_REDIS_EXTRA_INSTALL = 'pip install "10xgraph-api[redis]"'
+
+# Atomic sliding-window check using a Redis sorted set.
+#
+# The script:
+# 1. Removes timestamps older than the current window.
+# 2. Counts the remaining requests for the key.
+# 3. If under the limit, adds the current request as a unique sorted-set member.
+# 4. Sets an expiry so idle keys clean themselves up.
+# 5. Returns whether the request is allowed, remaining quota, and reset time.
+_SLIDING_WINDOW_LUA = """
+local key        = KEYS[1]
+local now_ms     = tonumber(ARGV[1])
+local window_ms  = tonumber(ARGV[2])
+local limit      = tonumber(ARGV[3])
+local member     = ARGV[4]
+
+local window_start = now_ms - window_ms
+
+redis.call('ZREMRANGEBYSCORE', key, '-inf', tostring(window_start))
+
+local count = tonumber(redis.call('ZCARD', key))
+
+if count < limit then
+    redis.call('ZADD', key, tostring(now_ms), member)
+    redis.call('EXPIRE', key, math.ceil(window_ms / 1000) + 1)
+    return {1, limit - count - 1, math.ceil(window_ms / 1000)}
+else
+    local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')
+    local reset_after
+    if #oldest >= 2 then
+        local oldest_ms = tonumber(oldest[2])
+        reset_after = math.ceil((oldest_ms + window_ms - now_ms) / 1000) + 1
+    else
+        reset_after = math.ceil(window_ms / 1000)
+    end
+    return {0, 0, math.max(reset_after, 1)}
+end
+"""
+
+
+class RedisRateLimitBackend(BaseRateLimitBackend):
+    """Distributed sliding-window rate limiter backed by Redis."""
+
+    def __init__(
+        self,
+        redis: Any,
+        prefix: str,
+        fail_open: bool = True,
+        close_redis: bool = False,
+    ) -> None:
+        self._redis = redis
+        self._prefix = prefix
+        self._fail_open = fail_open
+        self._close_redis = close_redis
+        self._script = self._redis.register_script(_SLIDING_WINDOW_LUA)
+
+    @classmethod
+    def from_url(
+        cls,
+        redis_url: str,
+        prefix: str,
+        fail_open: bool = True,
+    ) -> "RedisRateLimitBackend":
+        if not _REDIS_AVAILABLE or AsyncRedis is None:
+            raise ImportError(
+                "Redis backend requires the 'redis' package. "
+                f"Install the optional Redis extra with: {_REDIS_EXTRA_INSTALL}"
+            )
+        redis = AsyncRedis.from_url(redis_url, decode_responses=False)
+        return cls(redis=redis, prefix=prefix, fail_open=fail_open, close_redis=True)
+
+    async def check(self, key: str, *, limit: int, window: int) -> RateLimitDecision:
+        redis_key = f"{self._prefix}:{key}"
+        now_ms = int(time.time() * 1000)
+        window_ms = window * 1000
+        member = f"{now_ms}:{uuid.uuid4().hex}"
+
+        try:
+            result = await self._script(
+                keys=[redis_key],
+                args=[now_ms, window_ms, limit, member],
+            )
+            return RateLimitDecision(
+                allowed=bool(result[0]),
+                remaining=int(result[1]),
+                reset_after=int(result[2]),
+            )
+        except Exception:
+            logger.exception("Redis rate-limit backend error")
+            if self._fail_open:
+                return RateLimitDecision(allowed=True, remaining=0, reset_after=window)
+            return RateLimitDecision(allowed=False, remaining=0, reset_after=window)
+
+    async def close(self) -> None:
+        if self._close_redis:
+            await self._redis.aclose()

tenxgraph_api/src/app/core/middleware/request_limits.py

@@ -0,0 +1,154 @@
+"""Request size limit middleware for DoS protection."""
+
+from typing import Any
+
+from fastapi import status
+from fastapi.responses import JSONResponse
+from starlette.datastructures import Headers
+from starlette.exceptions import HTTPException
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from tenxgraph_api.src.app.core import logger
+
+
+class RequestBodyTooLarge(HTTPException):
+    """Raised from ``receive`` once a streamed body passes the limit.
+
+    It is an ``HTTPException`` because FastAPI re-raises those when they come from reading
+    the body, instead of turning them into a generic 400, so the client gets a 413.
+    """
+
+    def __init__(self, max_size: int) -> None:
+        super().__init__(
+            status_code=status.HTTP_413_CONTENT_TOO_LARGE,
+            detail=f"Request body too large. Maximum size is {max_size / (1024 * 1024):.1f}MB",
+        )
+
+
+class RequestSizeLimitMiddleware:
+    """Enforce a maximum request body size, whether or not ``Content-Length`` is sent.
+
+    A declared ``Content-Length`` over the limit is rejected before any body is read. Bodies
+    without one (``Transfer-Encoding: chunked``) are counted as they stream in and cut off as
+    soon as they pass the limit. That matters because FastAPI reads and parses the whole body,
+    including spooling multipart uploads to disk, before auth dependencies run.
+
+    This is a pure ASGI middleware: ``BaseHTTPMiddleware`` cannot see body chunks.
+
+    Args:
+        app: The ASGI application
+        max_size: Maximum request body size in bytes (default: 10MB)
+        path_limits: Per-path limits that replace ``max_size``, e.g. a larger one for the
+            file upload route so ``MEDIA_MAX_SIZE_MB`` is reachable.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        max_size: int = 10 * 1024 * 1024,
+        path_limits: dict[str, int] | None = None,
+    ) -> None:
+        self.app = app
+        self.max_size = max_size
+        self.max_size_mb = max_size / (1024 * 1024)
+        self.path_limits = dict(path_limits or {})
+
+    def _limit_for(self, scope: Scope) -> int:
+        path = scope.get("path", "")
+        root_path = (scope.get("root_path") or "").rstrip("/")
+        if root_path and path.startswith(root_path):
+            path = path[len(root_path) :] or "/"
+        return self.path_limits.get(path, self.max_size)
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        max_size = self._limit_for(scope)
+        max_size_mb = max_size / (1024 * 1024)
+
+        content_length = Headers(scope=scope).get("content-length")
+        if content_length is not None:
+            try:
+                declared = int(content_length)
+            except ValueError:
+                # A malformed Content-Length is a client error, not a server 500.
+                await self._error(
+                    scope,
+                    send,
+                    status.HTTP_400_BAD_REQUEST,
+                    {"code": "INVALID_CONTENT_LENGTH", "message": "Invalid Content-Length header."},
+                )
+                return
+            if declared > max_size:
+                logger.warning(
+                    f"Request rejected: size {declared} bytes exceeds limit of "
+                    f"{max_size} bytes ({max_size_mb:.1f}MB)"
+                )
+                await self._too_large(scope, send, max_size)
+                return
+
+        received = 0
+        response_started = False
+
+        async def limited_receive() -> Message:
+            nonlocal received
+            message = await receive()
+            if message["type"] == "http.request":
+                received += len(message.get("body", b""))
+                if received > max_size:
+                    logger.warning(
+                        f"Request rejected: streamed body passed the limit of "
+                        f"{max_size} bytes ({max_size_mb:.1f}MB)"
+                    )
+                    raise RequestBodyTooLarge(max_size)
+            return message
+
+        async def tracking_send(message: Message) -> None:
+            nonlocal response_started
+            if message["type"] == "http.response.start":
+                response_started = True
+            await send(message)
+
+        try:
+            await self.app(scope, limited_receive, tracking_send)
+        except RequestBodyTooLarge:
+            # Normally the app's HTTPException handling answers with the 413. This covers a
+            # body read outside FastAPI's request handling (e.g. by another middleware).
+            if response_started:
+                raise
+            await self._too_large(scope, send, max_size)
+
+    async def _too_large(self, scope: Scope, send: Send, max_size: int) -> None:
+        max_size_mb = max_size / (1024 * 1024)
+        await self._error(
+            scope,
+            send,
+            status.HTTP_413_CONTENT_TOO_LARGE,
+            {
+                "code": "REQUEST_TOO_LARGE",
+                "message": f"Request body too large. Maximum size is {max_size_mb:.1f}MB",
+                "max_size_bytes": max_size,
+                "max_size_mb": max_size_mb,
+            },
+        )
+
+    @staticmethod
+    async def _error(scope: Scope, send: Send, status_code: int, error: dict[str, Any]) -> None:
+        state = scope.get("state") or {}
+        response = JSONResponse(
+            status_code=status_code,
+            content={
+                "error": error,
+                "metadata": {
+                    "request_id": state.get("request_id", "unknown"),
+                    "status": "error",
+                },
+            },
+        )
+        await response(scope, _drain_receive, send)
+
+
+async def _drain_receive() -> Message:
+    return {"type": "http.disconnect"}