lorebase-keeper 0.2.3__py3-none-any.whl

lorebase_keeper/middleware.py

@@ -0,0 +1,285 @@
+"""ASGI middleware: request context (id, tracing, timing, access log), body size limit and
+security headers.
+
+All are plain ASGI classes so that they see every request, including those rejected by the
+framework, and so that streaming bodies are never buffered.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import time
+import uuid
+from collections.abc import Sequence
+from typing import Any
+
+from opentelemetry import metrics, propagate, trace
+from opentelemetry.metrics import MeterProvider
+from opentelemetry.trace import SpanKind, Status, StatusCode, TracerProvider
+from starlette.datastructures import Headers, MutableHeaders
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from lorebase_keeper.client_address import ClientAddresses, parse_address
+from lorebase_keeper.errors import (
+    INTERNAL_ERROR,
+    PAYLOAD_TOO_LARGE,
+    ProblemError,
+    ProblemRenderer,
+    request_id_of,
+)
+from lorebase_keeper.settings import SecurityHeadersSettings
+from lorebase_platform.logging import request_context
+from lorebase_platform.telemetry import LATENCY_BUCKETS_S
+
+logger = logging.getLogger("lorebase.keeper")
+access_logger = logging.getLogger("lorebase.keeper.access")
+
+_SAFE_REQUEST_ID = re.compile(r"^[A-Za-z0-9._:-]{1,128}$")
+UNMATCHED_ROUTE = "unmatched"
+
+
+def new_request_id() -> str:
+    """A fresh request id."""
+    return uuid.uuid4().hex
+
+
+class RequestContextMiddleware:
+    """Assign the request id, open the server span, time the request and log it.
+
+    The request id comes from the configured header when it is well formed, otherwise a new
+    one is generated; it is echoed in the response header, bound to every log record of the
+    request and attached to the span. Incoming W3C trace context is continued. Unhandled
+    exceptions are logged with their stack trace and answered with a generic problem detail.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        *,
+        header: str,
+        renderer: ProblemRenderer,
+        tracer_provider: TracerProvider | None = None,
+        meter_provider: MeterProvider | None = None,
+    ) -> None:
+        self.app = app
+        self._header = header
+        self._renderer = renderer
+        self._tracer = (tracer_provider or trace.get_tracer_provider()).get_tracer(
+            "lorebase.keeper"
+        )
+        self._duration = (
+            (meter_provider or metrics.get_meter_provider())
+            .get_meter("lorebase.keeper")
+            .create_histogram(
+                "http.server.request.duration",
+                unit="s",
+                description="Duration of HTTP requests.",
+                explicit_bucket_boundaries_advisory=LATENCY_BUCKETS_S,
+            )
+        )
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        headers = Headers(scope=scope)
+        candidate = headers.get(self._header)
+        request_id = candidate if candidate and _SAFE_REQUEST_ID.match(candidate) else None
+        request_id = request_id or new_request_id()
+        scope.setdefault("state", {})["request_id"] = request_id
+        method = scope["method"]
+        status = 500
+        started = False
+
+        async def send_with_id(message: Message) -> None:
+            nonlocal status, started
+            if message["type"] == "http.response.start":
+                started = True
+                status = message["status"]
+                MutableHeaders(scope=message)[self._header] = request_id
+            await send(message)
+
+        begin = time.perf_counter()
+        with (
+            request_context(request_id),
+            self._tracer.start_as_current_span(
+                method,
+                context=propagate.extract(dict(headers)),
+                kind=SpanKind.SERVER,
+                attributes={"http.request.method": method, "url.path": scope["path"]},
+            ) as span,
+        ):
+            span.set_attribute("lorebase.request_id", request_id)
+            try:
+                await self.app(scope, receive, send_with_id)
+            except Exception as exc:
+                logger.exception("unhandled error while serving request", extra={"method": method})
+                span.record_exception(exc)
+                if not started:
+                    response = self._renderer.response(
+                        INTERNAL_ERROR,
+                        detail=f"Unexpected error. Quote request id {request_id} when reporting.",
+                        instance=scope["path"],
+                        request_id=request_id,
+                    )
+                    await response(scope, receive, send_with_id)
+                else:
+                    raise  # the status line is gone: let the server abort the connection
+            finally:
+                elapsed = time.perf_counter() - begin
+                route = _route_template(scope)
+                span.update_name(f"{method} {route}")
+                span.set_attribute("http.route", route)
+                span.set_attribute("http.response.status_code", status)
+                if status >= 500:
+                    span.set_status(Status(StatusCode.ERROR))
+                self._duration.record(
+                    elapsed,
+                    {
+                        "http.request.method": method,
+                        "http.route": route,
+                        "http.response.status_code": status,
+                    },
+                )
+                access_logger.info(
+                    "request completed",
+                    extra={
+                        "method": method,
+                        "route": route,
+                        "status": status,
+                        "duration_ms": round(elapsed * 1000, 3),
+                    },
+                )
+
+
+def _route_template(scope: Scope) -> str:
+    route: Any = scope.get("route")
+    path = getattr(route, "path", None)
+    return path if isinstance(path, str) else UNMATCHED_ROUTE
+
+
+class BodyLimitMiddleware:
+    """Reject request bodies above the configured size with 413.
+
+    A declared `Content-Length` over the limit is refused before the handler runs; chunked
+    bodies are counted as they are read. `overrides` give other limits to the requests whose
+    method and path match (for example document uploads); the first match wins.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        *,
+        max_bytes: int,
+        renderer: ProblemRenderer,
+        overrides: Sequence[tuple[str, re.Pattern[str], int]] = (),
+    ) -> None:
+        self.app = app
+        self._max_bytes = max_bytes
+        self._renderer = renderer
+        self._overrides = tuple(overrides)
+
+    def _limit(self, scope: Scope) -> int:
+        for method, path, limit in self._overrides:
+            if scope.get("method") == method and path.fullmatch(scope["path"]):
+                return limit
+        return self._max_bytes
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        limit = self._limit(scope)
+        declared = Headers(scope=scope).get("content-length")
+        if declared is not None and declared.isdecimal() and int(declared) > limit:
+            response = self._renderer.response(
+                PAYLOAD_TOO_LARGE,
+                detail=self._detail(limit),
+                instance=scope["path"],
+                request_id=request_id_of(scope.get("state")),
+            )
+            await response(scope, receive, send)
+            return
+
+        received = 0
+
+        async def counted_receive() -> Message:
+            nonlocal received
+            message = await receive()
+            if message["type"] == "http.request":
+                received += len(message.get("body", b""))
+                if received > limit:
+                    raise ProblemError(PAYLOAD_TOO_LARGE, self._detail(limit))
+            return message
+
+        await self.app(scope, counted_receive, send)
+
+    def _detail(self, limit: int) -> str:
+        return f"The request body exceeds {limit} bytes."
+
+
+class SecurityHeadersMiddleware:
+    """Add the configured security headers to every response, keeping those a route set.
+
+    `Strict-Transport-Security` follows `hsts.enabled`: in `auto` it is sent when the request
+    came in over https, either directly or through a trusted proxy (`keeper.trusted_proxies`)
+    whose last `X-Forwarded-Proto` value is https.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        *,
+        settings: SecurityHeadersSettings,
+        proxies: ClientAddresses,
+    ) -> None:
+        self.app = app
+        self._hsts = settings.hsts
+        self._proxies = proxies
+        static: list[tuple[str, str]] = []
+        if settings.content_type_nosniff:
+            static.append(("X-Content-Type-Options", "nosniff"))
+        if settings.frame_options != "off":
+            ancestors = "'none'" if settings.frame_options == "DENY" else "'self'"
+            static.append(("X-Frame-Options", settings.frame_options))
+            static.append(("Content-Security-Policy", f"frame-ancestors {ancestors}"))
+        if settings.referrer_policy:
+            static.append(("Referrer-Policy", settings.referrer_policy))
+        self._static = tuple(static)
+
+    def _secure(self, scope: Scope) -> bool:
+        if scope["scheme"] in ("https", "wss"):
+            return True
+        client = scope.get("client")
+        peer = parse_address(str(client[0])) if client else None
+        if peer is None or not self._proxies.trusted(peer):
+            return False
+        forwarded = Headers(scope=scope).getlist("x-forwarded-proto")
+        hops = [hop.strip().lower() for value in forwarded for hop in value.split(",")]
+        return bool(hops) and hops[-1] == "https"
+
+    def _send_hsts(self, scope: Scope) -> bool:
+        mode = self._hsts.enabled
+        return mode == "always" or (mode == "auto" and self._secure(scope))
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        added = list(self._static)
+        if self._send_hsts(scope):
+            added.append(("Strict-Transport-Security", self._hsts.header_value))
+
+        async def send_with_headers(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                headers = MutableHeaders(scope=message)
+                for name, value in added:
+                    if name not in headers:
+                        headers[name] = value
+            await send(message)
+
+        await self.app(scope, receive, send_with_headers)

lorebase_keeper/pagination.py

@@ -0,0 +1,87 @@
+"""Cursor pagination shared by every list endpoint.
+
+A list request takes `limit` (bounded by settings) and an opaque `cursor`; the response is a
+`Page` with `items` and `next_cursor` (null on the last page). Cursors are base64url JSON so
+that routes can keep a keyset position in them; clients must treat them as opaque.
+"""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Annotated, Any
+
+from fastapi import Depends, Query
+from pydantic import BaseModel
+
+from lorebase_keeper.deps import KeeperSettingsDep
+from lorebase_keeper.errors import VALIDATION_FAILED, ProblemError
+
+MAX_CURSOR_LENGTH = 2048
+
+
+class Page[T](BaseModel):
+    """One page of results."""
+
+    items: list[T]
+    next_cursor: str | None = None
+
+
+def encode_cursor(position: Mapping[str, Any]) -> str:
+    """Opaque cursor carrying a keyset position."""
+    raw = json.dumps(position, separators=(",", ":"), sort_keys=True).encode()
+    return base64.urlsafe_b64encode(raw).decode().rstrip("=")
+
+
+def decode_cursor(cursor: str) -> dict[str, Any]:
+    """Position stored in a cursor; `ValueError` when the cursor is not one we issued."""
+    try:
+        decoded = json.loads(base64.urlsafe_b64decode(cursor + "=" * (-len(cursor) % 4)))
+    except (binascii.Error, ValueError):
+        raise ValueError("malformed cursor") from None
+    if not isinstance(decoded, dict):
+        raise ValueError("malformed cursor")
+    return decoded
+
+
+@dataclass(frozen=True)
+class PageParams:
+    """Validated pagination request: effective limit and decoded cursor position."""
+
+    limit: int
+    position: dict[str, Any] | None
+
+
+def _invalid(pointer: str, message: str) -> ProblemError:
+    return ProblemError(
+        VALIDATION_FAILED,
+        "One or more request members are invalid.",
+        extensions={"errors": [{"pointer": pointer, "message": message, "code": "invalid"}]},
+    )
+
+
+def get_page_params(
+    keeper: KeeperSettingsDep,
+    limit: Annotated[int | None, Query(ge=1, description="Page size.")] = None,
+    cursor: Annotated[
+        str | None, Query(max_length=MAX_CURSOR_LENGTH, description="Opaque page cursor.")
+    ] = None,
+) -> PageParams:
+    """Dependency resolving `limit` and `cursor` against the configured bounds."""
+    bounds = keeper.pagination
+    effective = bounds.default_limit if limit is None else limit
+    if effective > bounds.max_limit:
+        raise _invalid("/query/limit", f"must not exceed {bounds.max_limit}")
+    position = None
+    if cursor:
+        try:
+            position = decode_cursor(cursor)
+        except ValueError:
+            raise _invalid("/query/cursor", "is not a valid cursor") from None
+    return PageParams(limit=effective, position=position)
+
+
+PageParamsDep = Annotated[PageParams, Depends(get_page_params)]

lorebase_keeper/ratelimit/__init__.py

@@ -0,0 +1,10 @@
+"""Rate limits and daily quotas (design 13.6).
+
+`config` holds the settings, `bucket` the token bucket math and the `RateStore` protocol,
+`limiter` the per-principal limiter, `quotas` the daily counters of API keys and `enforce` the
+dependency, middleware and helpers that apply them to requests.
+"""
+
+from lorebase_keeper.ratelimit.bucket import BucketSpec, Decision, InMemoryRateStore, RateStore
+
+__all__ = ["BucketSpec", "Decision", "InMemoryRateStore", "RateStore"]

lorebase_keeper/ratelimit/bucket.py

@@ -0,0 +1,118 @@
+"""Token bucket arithmetic and the storage protocol behind it.
+
+A bucket holds up to `capacity` tokens and gains `refill_per_s` tokens each second; a request
+takes one. Time comes from the caller, so the math is deterministic and a shared store (Valkey,
+Redis) can run the same function server side.
+"""
+
+from __future__ import annotations
+
+import math
+import time
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Protocol
+
+from lorebase_keeper.ratelimit.config import BucketSettings
+
+
+@dataclass(frozen=True, slots=True)
+class BucketSpec:
+    """Size and refill rate of a bucket."""
+
+    capacity: int
+    refill_per_s: float
+
+    @classmethod
+    def from_settings(cls, settings: BucketSettings) -> BucketSpec:
+        return cls(capacity=settings.burst, refill_per_s=settings.requests_per_minute / 60)
+
+
+@dataclass(slots=True)
+class BucketState:
+    """Tokens a bucket held at `updated` (a reading of the store's clock, in seconds)."""
+
+    tokens: float
+    updated: float
+
+
+@dataclass(frozen=True, slots=True)
+class Decision:
+    """The outcome of taking from a bucket, in the terms of the `RateLimit-*` headers.
+
+    `reset_s` is the time until the bucket is full again; `retry_after_s` the time until the
+    request would be allowed (0 when it was).
+    """
+
+    allowed: bool
+    limit: int
+    remaining: int
+    reset_s: int
+    retry_after_s: int = 0
+
+
+def take(
+    state: BucketState | None, spec: BucketSpec, now: float, cost: float = 1.0
+) -> tuple[BucketState, Decision]:
+    """Refill the bucket to `now`, then take `cost` tokens if there are enough.
+
+    A bucket that was never used is full. Tokens are only taken from a bucket that holds at
+    least `cost`, so a refused request costs nothing.
+    """
+    if state is None:
+        tokens = float(spec.capacity)
+    else:
+        elapsed = max(0.0, now - state.updated)
+        tokens = min(float(spec.capacity), state.tokens + elapsed * spec.refill_per_s)
+    allowed = tokens >= cost
+    if allowed:
+        tokens -= cost
+    reset = math.ceil((spec.capacity - tokens) / spec.refill_per_s)
+    retry = 0 if allowed else max(1, math.ceil((cost - tokens) / spec.refill_per_s))
+    decision = Decision(
+        allowed=allowed,
+        limit=spec.capacity,
+        remaining=max(0, math.floor(tokens)),
+        reset_s=reset,
+        retry_after_s=retry,
+    )
+    return BucketState(tokens, now), decision
+
+
+class RateStore(Protocol):
+    """Keeps bucket state. `take` is atomic per key: concurrent requests of one principal see
+    each other's tokens."""
+
+    async def take(self, key: str, spec: BucketSpec, cost: float = 1.0) -> Decision: ...
+
+    async def peek(self, key: str, spec: BucketSpec, cost: float = 1.0) -> Decision:
+        """Whether `take` would allow `cost` now, without taking anything."""
+        ...
+
+
+class InMemoryRateStore:
+    """Buckets in this process, the least recently used dropped beyond `max_buckets`.
+
+    Nothing awaits between reading and writing a bucket, so on one event loop `take` is atomic.
+    """
+
+    def __init__(self, max_buckets: int, *, clock: Callable[[], float] = time.monotonic) -> None:
+        self._max = max_buckets
+        self._clock = clock
+        self._buckets: OrderedDict[str, BucketState] = OrderedDict()
+
+    def __len__(self) -> int:
+        return len(self._buckets)
+
+    async def take(self, key: str, spec: BucketSpec, cost: float = 1.0) -> Decision:
+        state, decision = take(self._buckets.get(key), spec, self._clock(), cost)
+        self._buckets[key] = state
+        self._buckets.move_to_end(key)
+        while len(self._buckets) > self._max:
+            self._buckets.popitem(last=False)
+        return decision
+
+    async def peek(self, key: str, spec: BucketSpec, cost: float = 1.0) -> Decision:
+        _, decision = take(self._buckets.get(key), spec, self._clock(), cost)
+        return decision