b5f82fb48c
Some checks are pending
CI / ruff (push) Waiting to run
CI / mypy (push) Waiting to run
CI / pytest (push) Waiting to run
CI / package (push) Waiting to run
CI / Install smoke (${{ matrix.label }}) (linux, ubuntu-latest) (push) Blocked by required conditions
CI / Install smoke (${{ matrix.label }}) (macos-arm64, macos-14) (push) Blocked by required conditions
277 lines
11 KiB
Python
277 lines
11 KiB
Python
"""
|
|
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)
|