interactive-story/backend/app/limits.py

"""Resource bounds on what a single request or a single story may cost.

Upstream carried three things here, and only one of them belongs in a local
single-user product. Per-IP and per-user **rate limiting**, the login-attempt
throttle, and the per-user **quotas** were hosted-service policy: they existed
to stop a hostile visitor exhausting a shared demo key or filling a shared
database. M2 removed all of it. There are no visitors, and throttling the one
person who started the application would be a bug rather than a guard.

What is left is defensive programming, and it applies whatever the deployment:

* a ceiling on the **request body**, so a malformed or hostile payload cannot
  be read into memory before anything looks at it;
* ceilings on how large **one adventure** may grow, in actions, memories,
  story cards and branches. These bound storage and the cost of the queries
  that walk them. They are per-story, not per-user: nothing here counts how
  many campaigns a person may have.

An import is checked against the same per-adventure ceilings that live creation
uses, so a bundle cannot carry a story past a limit that play could not reach.
"""

import json

from fastapi import HTTPException
from sqlalchemy import func
from sqlalchemy.orm import Session

from . import models

# ---------- Per-story row caps ----------

MAX_STORY_CARDS_PER_OWNER = 200   # Per scenario or per adventure.
MAX_MEMORIES_PER_ADVENTURE = 1000
MAX_ACTIONS_PER_ADVENTURE = 5000
# Phase 14, SP6. A tree holds one branch per divergence somebody built a story
# on, so a tree with more branches than the story has turns came from a file
# rather than from play. The cap applies to imports only. Forking is a POST that
# adds one row and has no cap of its own, and the cap that matters there is
# `MAX_ACTIONS_PER_ADVENTURE` above.
MAX_BRANCHES_PER_ADVENTURE = 1000


def check_row_cap(
    kind: str,
    db: Session,
    user: models.User,
    *,
    adventure: models.Adventure | None = None,
    scenario_id: int | None = None,
    adventure_id: int | None = None,
) -> None:
    """Raises a 409 when creating one more row of `kind` would exceed its cap.

    Only per-story kinds are capped. `adventures` and `scenarios` were per-user
    quotas and are no longer checked; the callers still pass them, and they are
    accepted and ignored so that adding a cap back is a change here rather than
    at every call site.
    """
    if kind in ("adventures", "scenarios"):
        return
    if kind == "story_cards":
        owner_filter = (
            models.StoryCard.scenario_id == scenario_id
            if scenario_id is not None
            else models.StoryCard.adventure_id == adventure_id
        )
        count = _count(db, models.StoryCard, owner_filter)
        cap, subject, hint = (
            MAX_STORY_CARDS_PER_OWNER, "story cards here", "delete one to make room"
        )
    elif kind == "memories":
        count = _count(db, models.Memory, models.Memory.adventure_id == adventure.id)
        cap, subject, hint = (
            MAX_MEMORIES_PER_ADVENTURE, "memories in this adventure",
            "delete some to make room",
        )
    elif kind == "actions":
        # Count every action in the adventure, which is the whole tree rather
        # than the path being played. That number is what costs storage, and
        # nothing is pruned automatically, so it is the right one to cap. It does
        # mean a heavily branched adventure reaches the cap while its story is
        # shorter than the cap, which is why the message counts "actions in this
        # adventure" rather than turns.
        count = _count(db, models.Action, models.Action.adventure_id == adventure.id)
        cap, subject, hint = (
            MAX_ACTIONS_PER_ADVENTURE, "actions in this adventure",
            "export it and continue in a new adventure",
        )
    else:  # pragma: no cover. This is a programming error, not user input.
        raise ValueError(f"Unknown row cap kind: {kind}")
    if count >= cap:
        raise HTTPException(409, f"You've reached the limit of {cap} {subject} — {hint}.")


def _count(db: Session, model, condition) -> int:
    return db.query(func.count(model.id)).filter(condition).scalar() or 0


_BUNDLE_LIST_CAPS = {
    "story_cards": MAX_STORY_CARDS_PER_OWNER,
    "memories": MAX_MEMORIES_PER_ADVENTURE,
    "actions": MAX_ACTIONS_PER_ADVENTURE,
    "branches": MAX_BRANCHES_PER_ADVENTURE,
}


def check_bundle_lists(**lists) -> None:
    """Raises a 409 when an import bundle's lists exceed the caps live creation uses.

    The keyword arguments are `story_cards`, `memories`, `actions`, and
    `branches`.
    """
    for name, value in lists.items():
        cap = _BUNDLE_LIST_CAPS[name]
        if isinstance(value, list) and len(value) > cap:
            noun = name.replace("_", " ")
            raise HTTPException(
                409, f"This file contains {len(value)} {noun} — the limit is {cap}."
            )


# ---------- Request body size ----------
# The limit is generous enough for the largest legitimate payload, which is an
# adventure export holding thousands of actions. No honest request approaches
# it.

MAX_BODY_BYTES = 2 * 1024 * 1024
MAX_IMPORT_BODY_BYTES = 20 * 1024 * 1024


class BodySizeLimitMiddleware:
    """Rejects oversized request bodies by their declared `Content-Length`.

    This is pure ASGI rather than `BaseHTTPMiddleware`, so SSE responses stream
    through unchanged. A chunked upload with no length is refused, because every
    real client of this API sends `Content-Length`, including browser fetch and
    curl with a file.
    """

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] == "http" and scope.get("method") in ("POST", "PUT", "PATCH"):
            headers = {k.decode("latin-1").lower(): v.decode("latin-1")
                       for k, v in scope.get("headers", [])}
            limit = (
                MAX_IMPORT_BODY_BYTES
                if scope.get("path", "").endswith("/import")
                else MAX_BODY_BYTES
            )
            length = headers.get("content-length")
            problem = None
            if length is None:
                if "chunked" in headers.get("transfer-encoding", "").lower():
                    problem = (411, "Content-Length is required.")
            else:
                try:
                    if int(length) > limit:
                        problem = (
                            413,
                            f"Request too large (limit {limit // (1024 * 1024)} MB).",
                        )
                except ValueError:
                    problem = (400, "Invalid Content-Length.")
            if problem:
                await _send_json_error(send, *problem)
                return
        await self.app(scope, receive, send)


async def _send_json_error(send, status: int, detail: str) -> None:
    body = json.dumps({"detail": detail}).encode()
    await send({
        "type": "http.response.start",
        "status": status,
        "headers": [(b"content-type", b"application/json"),
                    (b"content-length", str(len(body)).encode())],
    })
    await send({"type": "http.response.body", "body": body})