interactive-story/backend/app/limits.py

"""Phase 9 — abuse guards for hosted (multi-user) deployments.

Rate limits and row caps are no-ops in local mode: a single local player
should never be throttled by their own app. Values are hardcoded on purpose —
generous enough that a legitimate player never notices, tight enough that a
hostile visitor can't burn the demo key, peg the CPU, or bloat the database.
"""

import json
import threading
import time
from collections import defaultdict, deque

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

from . import auth, models

# ---------- Rate limiting ----------
# Fixed windows per (scope, caller). In-memory: fine for the single-process
# deployment this app targets (and the worst case after a restart is a brief
# extra allowance).

# scope -> (max requests, window seconds)
RATE_LIMITS: dict[str, tuple[int, int]] = {
    "turn": (10, 60),             # AI turn generation (demo key also has a daily cap)
    "chat": (30, 60),             # AI Chat scratchpad (power users only)
    "script-test": (30, 60),      # sandboxed, but each run costs up to 2s CPU
    "connection-test": (10, 60),  # outbound HTTP to a user-supplied URL
    "import": (30, 60),           # large writes
    "auth": (10, 300),            # register/login attempts, per IP
    "guest": (30, 300),           # new guest users, per IP (each is a DB row)
}

_windows: dict[tuple[str, str], deque] = defaultdict(deque)
_windows_guard = threading.Lock()


def _client_ip(request: Request) -> str:
    return request.client.host if request.client else "unknown"


def rate_limit(scope: str, request: Request, user: models.User | None = None) -> None:
    """429 when the caller exceeds the scope's window. Keyed per user when one
    is known (accounts survive IP changes), per IP otherwise."""
    if not auth.MULTI_USER:
        return
    limit, window_seconds = RATE_LIMITS[scope]
    key = (scope, f"u{user.id}" if user else f"ip{_client_ip(request)}")
    now = time.time()
    with _windows_guard:
        window = _windows[key]
        while window and window[0] < now - window_seconds:
            window.popleft()
        if len(window) >= limit:
            raise HTTPException(
                429, "You're doing that too fast — wait a minute and try again."
            )
        window.append(now)
        if len(_windows) > 10_000:
            _prune(now)


def _prune(now: float) -> None:
    """Drop callers whose whole window has expired (call with guard held) so
    the per-IP dict can't grow without bound."""
    longest = max(seconds for _, seconds in RATE_LIMITS.values())
    stale = [key for key, window in _windows.items()
             if not window or window[-1] < now - longest]
    for key in stale:
        del _windows[key]


# ---------- Per-user row caps ----------

MAX_ADVENTURES_PER_USER = 100
MAX_SCENARIOS_PER_USER = 200
MAX_SCRIPTS_PER_USER = 200
MAX_STORY_CARDS_PER_OWNER = 200   # per scenario or adventure
MAX_MEMORIES_PER_ADVENTURE = 1000
MAX_ACTIONS_PER_ADVENTURE = 5000


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:
    """409 with a friendly message when creating one more row of `kind` would
    exceed its cap. Ownership of the passed scenario/adventure has already
    been checked by the caller."""
    if not auth.MULTI_USER:
        return
    if kind == "adventures":
        count = _count(db, models.Adventure, models.Adventure.user_id == user.id)
        cap, subject, hint = (
            MAX_ADVENTURES_PER_USER, "adventures",
            "delete one you no longer play to make room",
        )
    elif kind == "scenarios":
        count = _count(db, models.Scenario, models.Scenario.user_id == user.id)
        cap, subject, hint = (
            MAX_SCENARIOS_PER_USER, "scenarios", "delete one to make room"
        )
    elif kind == "scripts":
        count = _count(db, models.Script, models.Script.user_id == user.id)
        cap, subject, hint = (
            MAX_SCRIPTS_PER_USER, "scripts", "delete one to make room"
        )
    elif 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 = _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 — 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,
}


def check_bundle_lists(**lists) -> None:
    """409 when an import bundle's lists exceed the same caps live creation
    enforces (kwargs: story_cards=, memories=, actions=)."""
    if not auth.MULTI_USER:
        return
    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 ----------
# Generous enough for the biggest legitimate payload (an adventure export with
# thousands of actions), applied in every mode — no honest request comes close.

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


class BodySizeLimitMiddleware:
    """Rejects oversized request bodies by declared Content-Length. Pure ASGI
    (not BaseHTTPMiddleware) so SSE responses stream through untouched.
    Chunked uploads without a length are refused — every real client of this
    API (browser fetch, curl with a file) sends Content-Length."""

    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})