Delete limiter.py

limiter.py

@@ -1,276 +0,0 @@
-"""
-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)