limiter.py

"""
Per-day talk-time budget for the speech-to-speech demo.

Our server isn't in the audio path (the browser dials the compute WebSocket
directly), so it can't cut a live stream. What it *can* do is meter time with a
server-clock, chunked reservation:

  - At grant we reserve the first chunk (CHUNK_SEC) and debit it from the day's
    budget. A parallel grant therefore sees the budget already spent.
  - The client heartbeats; each heartbeat extends the reservation one chunk at a
    time until the daily budget runs out, then we report `expired` so the client
    tears down.
  - On a clean end (sendBeacon) we reconcile to the real elapsed time and refund
    the unused chunk. A crash (no end, no heartbeats) is reaped by a sweep and
    forfeits at most one chunk.

All time is the server's clock. Budgets are per UTC day; a new day is simply a
new row (no explicit reset). Logged-in users are keyed by a hashed HF `sub`;
anonymous users by BOTH a hashed IP and a hashed signed-cookie id, OR-matched
(spent = max of the two) so clearing one identifier doesn't reset the budget.

Storage is SQLite at $USAGE_DB_PATH, else /data (persistent Spaces storage),
else a /tmp fallback. On /tmp the budget is only per-uptime — flagged in logs.
"""

import hashlib
import hmac
import logging
import math
import os
import sqlite3
import tempfile
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

logger = logging.getLogger("s2s.limiter")

# ── Tunables (env-overridable) ───────────────────────────────────────────────
ANON_SEC = int(os.environ.get("LIMIT_ANON_SEC", "300"))   # 5 min/day, not signed in
FREE_SEC = int(os.environ.get("LIMIT_FREE_SEC", "600"))   # 10 min/day, signed in, no PRO
CHUNK_SEC = int(os.environ.get("RESERVE_CHUNK_SEC", "10"))  # reservation granularity
HEARTBEAT_SEC = int(os.environ.get("HEARTBEAT_SEC", "5"))   # advertised client cadence
REAP_AFTER_SEC = int(os.environ.get("SESSION_REAP_SEC", "15"))  # silence before sweep

# Stable across restarts or the hashed keys (and signed cookies) rotate and the
# budget effectively resets. Set it as a Space secret. Falls back to a per-boot
# random value (keys then only hold within one uptime).
_HASH_SECRET = (os.environ.get("USAGE_HASH_SECRET", "").strip() or os.urandom(32).hex()).encode()

_lock = threading.Lock()
_db_path: "Path | None" = None


def budget_for(tier: str) -> "int | None":
    """Daily second-budget for a tier, or None for unlimited.

    Unlimited tiers: 'pro' (paying PRO members) and 'org' (members of an
    allow-listed organisation, see UNLIMITED_ORGS in auth.py)."""
    if tier in ("pro", "org"):
        return None
    if tier == "free":
        return FREE_SEC
    return ANON_SEC


def hash_key(raw: str) -> str:
    """HMAC a raw identifier (sub / ip / cookie id) into an opaque storage key."""
    digest = hmac.new(_HASH_SECRET, raw.encode("utf-8"), hashlib.sha256).hexdigest()
    return f"k_{digest}"


def sign_cookie(value: str) -> str:
    """`<id>.<sig>` so a forged anon-cookie id is rejected on read."""
    sig = hmac.new(_HASH_SECRET, value.encode("utf-8"), hashlib.sha256).hexdigest()[:32]
    return f"{value}.{sig}"


def verify_cookie(signed: str) -> "str | None":
    """Return the id if the signature checks out, else None."""
    if not signed or "." not in signed:
        return None
    value, _, sig = signed.rpartition(".")
    want = hmac.new(_HASH_SECRET, value.encode("utf-8"), hashlib.sha256).hexdigest()[:32]
    return value if hmac.compare_digest(sig, want) else None


def _today() -> str:
    return datetime.now(timezone.utc).date().isoformat()


def _resolve_db_path() -> Path:
    explicit = os.environ.get("USAGE_DB_PATH", "").strip()
    if explicit:
        return Path(explicit)
    data = Path("/data")
    if data.is_dir() and os.access(data, os.W_OK):
        return data / "s2s-usage.sqlite3"
    logger.warning("No persistent /data — usage budget falls back to /tmp (per-uptime only).")
    return Path(tempfile.gettempdir()) / "s2s-usage.sqlite3"


def _connect() -> sqlite3.Connection:
    con = sqlite3.connect(_db_path, timeout=5.0)
    con.execute("PRAGMA journal_mode=WAL")
    con.execute("PRAGMA busy_timeout=5000")
    return con


def init() -> None:
    """Create the schema. Call once at startup."""
    global _db_path
    _db_path = _resolve_db_path()
    with _lock, _connect() as con:
        con.execute(
            """CREATE TABLE IF NOT EXISTS usage_daily (
                   user_key   TEXT NOT NULL,
                   day        TEXT NOT NULL,
                   spent_sec  INTEGER NOT NULL DEFAULT 0,
                   updated_at INTEGER NOT NULL,
                   PRIMARY KEY (user_key, day)
               )"""
        )
        con.execute(
            """CREATE TABLE IF NOT EXISTS sessions (
                   session_id   TEXT PRIMARY KEY,
                   keys         TEXT NOT NULL,   -- comma-joined usage_daily keys to debit
                   day          TEXT NOT NULL,
                   tier         TEXT NOT NULL,
                   grant_ts     REAL NOT NULL,
                   last_seen_ts REAL NOT NULL,
                   reserved_sec INTEGER NOT NULL,
                   ended        INTEGER NOT NULL DEFAULT 0
               )"""
        )
    logger.info("Usage limiter ready at %s (anon=%ss free=%ss chunk=%ss)", _db_path, ANON_SEC, FREE_SEC, CHUNK_SEC)


# ── Internal helpers (call under _lock) ───────────────────────────────────────

def _spent(con, key: str, day: str) -> int:
    row = con.execute(
        "SELECT spent_sec FROM usage_daily WHERE user_key=? AND day=?", (key, day)
    ).fetchone()
    return int(row[0]) if row else 0


def _spent_max(con, keys, day: str) -> int:
    """OR-match: the most-spent identifier governs."""
    return max((_spent(con, k, day) for k in keys), default=0)


def _add(con, keys, day: str, delta: int) -> None:
    now = int(time.time())
    for k in keys:
        cur = _spent(con, k, day)
        nxt = max(0, cur + delta)
        con.execute(
            """INSERT INTO usage_daily (user_key, day, spent_sec, updated_at)
               VALUES (?, ?, ?, ?)
               ON CONFLICT(user_key, day) DO UPDATE SET
                   spent_sec = excluded.spent_sec, updated_at = excluded.updated_at""",
            (k, day, nxt, now),
        )


# ── Public API ────────────────────────────────────────────────────────────────

def remaining(keys, tier: str) -> "int | None":
    """Seconds left today for these keys (None = unlimited). No mutation."""
    budget = budget_for(tier)
    if budget is None:
        return None
    with _lock, _connect() as con:
        return max(0, budget - _spent_max(con, keys, _today()))


def begin(session_id: str, keys, tier: str) -> int:
    """Reserve the first chunk for a new session and record it. Returns the
    chunk reserved (0 if the budget is already exhausted — the first heartbeat
    will then expire it). PRO (unlimited) is never tracked; don't call it here."""
    day = _today()
    budget = budget_for(tier)
    now = time.time()
    with _lock, _connect() as con:
        avail = budget - _spent_max(con, keys, day) if budget is not None else CHUNK_SEC
        chunk = max(0, min(CHUNK_SEC, avail))
        if chunk:
            _add(con, keys, day, chunk)
        con.execute(
            """INSERT OR REPLACE INTO sessions
               (session_id, keys, day, tier, grant_ts, last_seen_ts, reserved_sec, ended)
               VALUES (?, ?, ?, ?, ?, ?, ?, 0)""",
            (session_id, ",".join(keys), day, tier, now, now, chunk),
        )
        return chunk


def heartbeat(session_id: str) -> bool:
    """Keep a session alive: extend the reservation toward `elapsed + 1 chunk`,
    debiting the budget chunk by chunk. Returns True while alive, False once the
    budget is spent (caller should tear down) or the session is unknown/ended."""
    now = time.time()
    with _lock, _connect() as con:
        row = con.execute(
            "SELECT keys, day, tier, grant_ts, reserved_sec, ended FROM sessions WHERE session_id=?",
            (session_id,),
        ).fetchone()
        if not row or row[5]:
            return False
        keys = row[0].split(",")
        day, tier, grant_ts, reserved = row[1], row[2], row[3], int(row[4])
        budget = budget_for(tier)
        elapsed = now - grant_ts

        # Grow the reservation one chunk at a time until it covers elapsed + a
        # one-chunk lookahead, or the budget runs dry.
        while reserved < elapsed + CHUNK_SEC:
            if budget is not None and _spent_max(con, keys, day) >= budget:
                break
            _add(con, keys, day, CHUNK_SEC)
            reserved += CHUNK_SEC

        alive = reserved > elapsed  # could we cover the time already elapsed?
        con.execute(
            "UPDATE sessions SET last_seen_ts=?, reserved_sec=?, ended=? WHERE session_id=?",
            (now, reserved, 0 if alive else 1, session_id),
        )
        if not alive:
            _reconcile(con, keys, day, grant_ts, reserved, end_ts=now)
        return alive


def end(session_id: str) -> None:
    """Clean teardown: reconcile to actual elapsed time and refund the unused
    reservation. Idempotent."""
    now = time.time()
    with _lock, _connect() as con:
        row = con.execute(
            "SELECT keys, day, grant_ts, reserved_sec, ended FROM sessions WHERE session_id=?",
            (session_id,),
        ).fetchone()
        if not row or row[4]:
            return
        keys, day, grant_ts, reserved = row[0].split(","), row[1], row[2], int(row[3])
        _reconcile(con, keys, day, grant_ts, reserved, end_ts=now)
        con.execute("UPDATE sessions SET ended=1, last_seen_ts=? WHERE session_id=?", (now, session_id))


def sweep() -> None:
    """Reap sessions that went silent (crash / closed without a beacon): bill
    their elapsed time, refund the rest, mark ended. Forfeits ≤ one chunk."""
    now = time.time()
    cutoff = now - REAP_AFTER_SEC
    with _lock, _connect() as con:
        stale = con.execute(
            "SELECT session_id, keys, day, grant_ts, last_seen_ts, reserved_sec FROM sessions "
            "WHERE ended=0 AND last_seen_ts < ?",
            (cutoff,),
        ).fetchall()
        for session_id, keys_s, day, grant_ts, last_seen, reserved in stale:
            _reconcile(con, keys_s.split(","), day, grant_ts, int(reserved), end_ts=last_seen)
            con.execute("UPDATE sessions SET ended=1 WHERE session_id=?", (session_id,))
        if stale:
            logger.debug("swept %d stale session(s)", len(stale))


def _reconcile(con, keys, day: str, grant_ts: float, reserved: int, end_ts: float) -> None:
    """Refund reserved-but-unused time. Bill elapsed rounded up to a chunk,
    capped at what was reserved."""
    elapsed = max(0.0, end_ts - grant_ts)
    billed = min(reserved, int(math.ceil(elapsed / CHUNK_SEC) * CHUNK_SEC))
    refund = reserved - billed
    if refund > 0:
        _add(con, keys, day, -refund)