diff --git a/passiveworkers/net/coordinator_app.py b/passiveworkers/net/coordinator_app.py index fe204a2..e54bbb6 100644 --- a/passiveworkers/net/coordinator_app.py +++ b/passiveworkers/net/coordinator_app.py @@ -23,6 +23,11 @@ GET /dashboard → live operator map GET /healthz +Request bodies are capped at PW_MAX_BODY_BYTES (default 4 MiB) BEFORE any parsing or auth: a larger +declared Content-Length, or a chunked body that streams past the cap, gets a 413 (F11 review). The +cap sits above the largest legitimate body (a 200k-char /deliver JSON-escaped is ≤ ~2.4 MB); blob +uploads keep their own tighter 512 KiB cap. + Run: PW_TOKEN=… python -m passiveworkers.net.coordinator_app """ @@ -34,7 +39,7 @@ import threading from fastapi import FastAPI, Header, HTTPException, Request, Response -from fastapi.responses import HTMLResponse +from fastapi.responses import HTMLResponse, JSONResponse from pydantic import BaseModel, Field from passiveworkers.net import geoip @@ -70,7 +75,81 @@ def _startup_guard() -> None: print("[coordinator] PW_TRUST_XFF enabled — rate-limit keys trust X-Forwarded-For. Ensure your " "reverse proxy sets AND sanitizes that header (drops client-supplied values), or attackers " "can rotate it to defeat per-client limits.", flush=True) + + +_MAX_BODY_DEFAULT = 4 * 1024 * 1024 # see the module docstring for why 4 MiB + + +def _max_body_bytes() -> int: + """The request-body cap in bytes (PW_MAX_BODY_BYTES). Invalid or non-positive → the default: + there is deliberately no 'unbounded' setting.""" + try: + cap = int(os.environ.get("PW_MAX_BODY_BYTES", _MAX_BODY_DEFAULT)) + except (TypeError, ValueError): + return _MAX_BODY_DEFAULT + return cap if cap > 0 else _MAX_BODY_DEFAULT + + +class _BodyTooLarge(HTTPException): + def __init__(self) -> None: + super().__init__(status_code=413, detail="request body too large") + + +class BodySizeLimit: + """Pure-ASGI request-body cap (F11 review). A declared Content-Length over the cap is refused + before the app runs at all; otherwise the body is counted as it streams (chunked or lying + clients), and the first message that crosses the cap raises 413 — the app never buffers more + than cap + one server-sized chunk. Raised as an HTTPException so FastAPI's body parser re-raises + it unchanged (a generic error there would become a 400).""" + + def __init__(self, app) -> None: + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + cap = _max_body_bytes() + for name, value in scope.get("headers", []): + if name == b"content-length": + try: + declared = int(value) + except ValueError: + break # the server rejects malformed lengths; the stream count still applies + if declared > cap: + await JSONResponse({"detail": "request body too large"}, status_code=413)( + scope, receive, send) + return + break + seen = 0 + started = False + + async def capped_receive(): + nonlocal seen + message = await receive() + if message["type"] == "http.request": + seen += len(message.get("body", b"")) + if seen > cap: + raise _BodyTooLarge() + return message + + async def tracking_send(message): + nonlocal started + if message["type"] == "http.response.start": + started = True + await send(message) + + try: + await self.app(scope, capped_receive, tracking_send) + except _BodyTooLarge: + if started: + raise + await JSONResponse({"detail": "request body too large"}, status_code=413)( + scope, receive, send) + + app = FastAPI(title="Passive Workers — Council Coordinator") +app.add_middleware(BodySizeLimit) store = Store() _rl = RateLimiter() # per-process sliding-window rate limiter (D36) @@ -419,6 +498,44 @@ def _baseline_async(job_id: str, question: str) -> None: print(f"[baseline] store failed for job {job_id[:8]}: {type(e).__name__}: {e}", flush=True) +def _baseline_workers() -> int: + try: + return max(1, int(os.environ.get("PW_BASELINE_WORKERS", "4"))) + except (TypeError, ValueError): + return 4 + + +# F12 review: at most `_BASELINE_MAX` baseline threads at once instead of one per job. Each can live +# up to 900 s (waiting for a local job to settle), so a job submitted while every slot is busy simply +# gets no baseline — never a queued backlog. Daemon threads, not a ThreadPoolExecutor: executor +# workers are joined at interpreter exit, which would stall a coordinator shutdown for 15 minutes. +_BASELINE_MAX = _baseline_workers() +_baseline_slots = threading.BoundedSemaphore(_BASELINE_MAX) + + +def _start_baseline(job_id: str, question: str) -> bool: + """Start the baseline if a slot is free. False (and a log line) when all slots are busy.""" + slots = _baseline_slots # release the semaphore we acquired, even if the global is rebound + if not slots.acquire(blocking=False): + print(f"[baseline] skipped for job {job_id[:8]}: {_BASELINE_MAX} already running", flush=True) + return False + + def run() -> None: + try: + _baseline_async(job_id, question) + except Exception as e: # best-effort, like _baseline_async's own store failure + print(f"[baseline] failed for job {job_id[:8]}: {type(e).__name__}: {e}", flush=True) + finally: + slots.release() + + try: + threading.Thread(target=run, daemon=True, name="pw-baseline").start() + except RuntimeError: # can't start a thread (resource limit / interpreter shutdown) + slots.release() + return False + return True + + @app.get("/job-types") def job_types(): """The marketplace catalog — what kinds of work computers can request here.""" @@ -442,8 +559,7 @@ def submit_job(body: JobBody, x_user_secret: str | None = Header(default=None)): if out.get("status") == "pending_answers" and not task_behavior(body.type).sharded: # the honest single-model compare only makes sense for answer/report jobs — # a one-shot model can't process a sharded item batch (shard_map/download_extract/code_generation) - threading.Thread(target=_baseline_async, args=(out["job_id"], body.question), - daemon=True, name="pw-baseline").start() + _start_baseline(out["job_id"], body.question) return out diff --git a/passiveworkers/net/ratelimit.py b/passiveworkers/net/ratelimit.py index 0dfc980..4514261 100644 --- a/passiveworkers/net/ratelimit.py +++ b/passiveworkers/net/ratelimit.py @@ -9,17 +9,21 @@ Deliberately tiny + dependency-free (one coordinator process today; swap for Redis behind the same `allow()` seam when horizontally scaled). Per-key sliding window: at most `limit` events per `window_s`. `limit <= 0` disables that key (operator escape hatch). Memory is bounded — a denied -event is NOT recorded (so a key's deque never exceeds `limit`), and idle keys are swept when the -table grows past `max_keys`. +event is NOT recorded (so a key's deque never exceeds `limit`), and once the table grows past +`max_keys` it is swept — at most once per `_SWEEP_EVERY_S` (F16 review: an unthrottled O(n) sweep +ran on EVERY call while the table stayed over the cap) — dropping idle keys first, then the +least-recently-hit keys until it is back at `max_keys`. """ from __future__ import annotations +import heapq import threading import time from collections import deque _SWEEP_HORIZON_S = 3600.0 # keys idle longer than this are dropped on overflow (>> any window) +_SWEEP_EVERY_S = 10.0 # at most one sweep per this many seconds, however far over the cap class RateLimiter: @@ -27,6 +31,7 @@ def __init__(self, max_keys: int = 50_000): self._hits: dict[str, deque] = {} self._lock = threading.Lock() self._max_keys = max_keys + self._last_sweep = float("-inf") def allow(self, key: str, limit: int, window_s: float, now: float | None = None) -> bool: """True if an event for `key` is within `limit` per `window_s` (and records it); False if it @@ -46,13 +51,21 @@ def allow(self, key: str, limit: int, window_s: float, now: float | None = None) if len(dq) >= limit: return False # denied — do not record, so the deque stays bounded dq.append(t) - if len(self._hits) > self._max_keys: + # abs(): a clock that stepped backwards must not suppress sweeping until it catches up + if len(self._hits) > self._max_keys and abs(t - self._last_sweep) >= _SWEEP_EVERY_S: + self._last_sweep = t self._sweep(t) return True def _sweep(self, now: float) -> None: """Drop keys that are empty or idle past the horizon (>> any real window, so a swept key's - limit window has certainly elapsed) — bounds memory under churny key spaces.""" + limit window has certainly elapsed). If that still leaves the table over `max_keys` (key + rotation keeps every key fresh), evict the least-recently-hit keys down to the cap: memory + stays bounded, and an evicted key only loses its own history, so it is never MORE limited.""" horizon = now - _SWEEP_HORIZON_S for k in [k for k, d in self._hits.items() if not d or d[-1] < horizon]: del self._hits[k] + excess = len(self._hits) - self._max_keys + if excess > 0: + for k in heapq.nsmallest(excess, self._hits, key=lambda k: self._hits[k][-1]): + del self._hits[k] diff --git a/passiveworkers/net/store.py b/passiveworkers/net/store.py index c62c031..86f23df 100644 --- a/passiveworkers/net/store.py +++ b/passiveworkers/net/store.py @@ -152,4 +152,17 @@ def _init_schema(self) -> None: self.conn.execute("ALTER TABLE nodes ADD COLUMN ok_count INTEGER DEFAULT 0") if "fail_count" not in ncols: self.conn.execute("ALTER TABLE nodes ADD COLUMN fail_count INTEGER DEFAULT 0") + # F13 review: the hot lookups (secret auth on every heartbeat/poll, a job's tasks on every + # job_view, a node's queue, an asker's jobs, the reaper's status scans) were full-table scans + # under the global lock. IF NOT EXISTS → idempotent on every open, existing DBs included. + self.conn.executescript( + """ + CREATE INDEX IF NOT EXISTS idx_nodes_secret_hash ON nodes(secret_hash); + CREATE INDEX IF NOT EXISTS idx_users_secret_hash ON users(secret_hash); + CREATE INDEX IF NOT EXISTS idx_tasks_job_id ON tasks(job_id); + CREATE INDEX IF NOT EXISTS idx_tasks_node_status ON tasks(node_id, status); + CREATE INDEX IF NOT EXISTS idx_jobs_asker_created ON jobs(asker, created); + CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status); + """ + ) self.conn.commit() diff --git a/tests/test_coordinator_limits.py b/tests/test_coordinator_limits.py new file mode 100644 index 0000000..8d8128d --- /dev/null +++ b/tests/test_coordinator_limits.py @@ -0,0 +1,217 @@ +"""Coordinator resource bounds (review F11, F12, F13): + • F11 — request bodies over PW_MAX_BODY_BYTES get a 413 BEFORE parsing or auth, whether the + size is declared (Content-Length) or only discovered while streaming (chunked / lying client). + • F12 — baseline generation runs on at most PW_BASELINE_WORKERS threads, not one per job. + • F13 — the hot lookups have SQLite indexes, created idempotently on every open. +""" +import asyncio +import importlib +import sqlite3 +import threading + +import pytest + + +@pytest.fixture() +def capp(tmp_path, monkeypatch): + monkeypatch.setenv("PW_DB", str(tmp_path / "limits.db")) + monkeypatch.setenv("PW_TOKEN", "tok") + monkeypatch.setenv("PW_STARTER_CREDITS", "100000") + monkeypatch.setenv("PW_BASELINE_WORKERS", "2") + for m in ("passiveworkers.ledger", "passiveworkers.net.config", "passiveworkers.net.store", + "passiveworkers.net.coordinator_app"): + importlib.reload(importlib.import_module(m)) + import passiveworkers.net.coordinator_app as mod + return mod + + +@pytest.fixture() +def client(capp): + from fastapi.testclient import TestClient + return TestClient(capp.app) + + +TOK = {"X-PW-Token": "tok"} + + +# --------------------------------------------------------------------------------- F11 body cap +def _spy_complete(capp, monkeypatch): + calls = [] + monkeypatch.setattr(capp.store, "complete_task", lambda *a, **k: calls.append(a) or True) + return calls + + +def test_declared_oversize_body_is_413_before_auth_or_parsing(capp, client, monkeypatch): + calls = _spy_complete(capp, monkeypatch) + body = b'{"text": "' + b"x" * (4 * 1024 * 1024) + b'"}' # default cap is 4 MiB + r = client.post("/tasks/t1/result", content=body, + headers={"content-type": "application/json", "X-Node-Secret": "nope"}) + assert r.status_code == 413 + assert r.json()["detail"] == "request body too large" + assert calls == [] + + +def test_chunked_oversize_body_is_413(capp, client, monkeypatch): + monkeypatch.setenv("PW_MAX_BODY_BYTES", "4096") + calls = _spy_complete(capp, monkeypatch) + + def gen(): # no Content-Length → Transfer-Encoding: chunked + yield b'{"text": "' + yield b"x" * 8192 + yield b'"}' + + r = client.post("/tasks/t1/result", content=gen(), + headers={"content-type": "application/json", "X-Node-Secret": "nope"}) + assert "content-length" not in {k.lower() for k in r.request.headers} + assert r.status_code == 413 + assert calls == [] + + +def test_streamed_body_is_cut_off_at_the_cap(capp, monkeypatch): + """Drive the ASGI app with a body that arrives in many 64 KiB messages: the read must stop + at the first message past the cap instead of buffering the whole thing.""" + monkeypatch.setenv("PW_MAX_BODY_BYTES", str(256 * 1024)) + chunk, total = b"x" * 65536, 100 + consumed = 0 + sent = [] + + async def receive(): + nonlocal consumed + if consumed >= total: + return {"type": "http.disconnect"} + consumed += 1 + return {"type": "http.request", "body": chunk, "more_body": consumed < total} + + async def send(message): + sent.append(message) + + scope = {"type": "http", "asgi": {"version": "3.0"}, "http_version": "1.1", "method": "POST", + "scheme": "http", "path": "/tasks/t1/result", "raw_path": b"/tasks/t1/result", + "root_path": "", "query_string": b"", "client": ("127.0.0.1", 1), "server": ("t", 80), + "headers": [(b"content-type", b"application/json"), (b"x-node-secret", b"nope"), + (b"transfer-encoding", b"chunked")]} + asyncio.run(capp.app(scope, receive, send)) + start = next(m for m in sent if m["type"] == "http.response.start") + assert start["status"] == 413 + assert consumed <= 5 # 256 KiB cap / 64 KiB chunks → stops on the 5th + + +def test_bodies_under_the_cap_still_reach_the_endpoint(client, monkeypatch): + monkeypatch.setenv("PW_MAX_BODY_BYTES", "4096") + r = client.post("/users", json={"handle": "alice"}) + assert r.status_code == 200 and r.json()["user_secret"] + r = client.post("/tasks/t1/result", json={"text": "x" * 1000}, headers={"X-Node-Secret": "nope"}) + assert r.status_code == 401 # normal auth path, not the size cap + + +@pytest.mark.parametrize("raw", ["0", "-5", "lots"]) +def test_bad_cap_setting_falls_back_to_the_default(capp, monkeypatch, raw): + monkeypatch.setenv("PW_MAX_BODY_BYTES", raw) + assert capp._max_body_bytes() == 4 * 1024 * 1024 + + +# --------------------------------------------------------------------------- F12 baseline bound +def _new_baseline_threads(before): + """pw-baseline threads started since `before` (other modules' jobs may leave their own running).""" + return [t for t in threading.enumerate() if t.name == "pw-baseline" and t not in before] + + +def _join(threads): + for t in threads: + t.join(5) + assert not t.is_alive() + + +def test_baseline_threads_are_bounded(capp, client, monkeypatch): + release = threading.Event() + lock = threading.Lock() + running = {"now": 0, "peak": 0, "calls": 0} + + def fake_baseline(job_id, question): + with lock: + running["now"] += 1 + running["calls"] += 1 + running["peak"] = max(running["peak"], running["now"]) + release.wait(10) + with lock: + running["now"] -= 1 + + monkeypatch.setattr(capp, "_baseline_async", fake_baseline) + before = set(threading.enumerate()) + client.post("/nodes/register", headers=TOK, + json={"owner": "op", "answer_model": "m", "profile": {"models": ["m"]}}) + usec = client.post("/users", json={"handle": "alice"}).json()["user_secret"] + for i in range(6): + r = client.post("/jobs", json={"question": f"q{i}"}, headers={"X-User-Secret": usec}) + assert r.json()["status"] == "pending_answers" + try: + live = [t for t in _new_baseline_threads(before) if t.is_alive()] + assert len(live) <= 2 # PW_BASELINE_WORKERS=2, not one thread per job + assert running["calls"] <= 2 + finally: + release.set() + _join(_new_baseline_threads(before)) + assert running["peak"] == 2 + # slots are released when a baseline finishes, so the next job gets one again + client.post("/jobs", json={"question": "after"}, headers={"X-User-Secret": usec}) + _join(_new_baseline_threads(before)) + assert running["calls"] == 3 + + +def test_baseline_failure_releases_its_slot(capp, monkeypatch): + def boom(job_id, question): + raise RuntimeError("model down") + + monkeypatch.setattr(capp, "_baseline_async", boom) + before = set(threading.enumerate()) + for i in range(5): # more than the 2 slots, run one after another + assert capp._start_baseline(f"job{i}", "q") is True + _join(_new_baseline_threads(before)) + + +# ------------------------------------------------------------------------------ F13 indexes +_EXPECTED = { + "idx_nodes_secret_hash": ("nodes", ["secret_hash"]), + "idx_users_secret_hash": ("users", ["secret_hash"]), + "idx_tasks_job_id": ("tasks", ["job_id"]), + "idx_tasks_node_status": ("tasks", ["node_id", "status"]), + "idx_jobs_asker_created": ("jobs", ["asker", "created"]), + "idx_jobs_status": ("jobs", ["status"]), +} + + +def _indexes(conn): + out = {} + for name, table in conn.execute("SELECT name, tbl_name FROM sqlite_master WHERE type='index' " + "AND name LIKE 'idx_%'"): + cols = [r[2] for r in conn.execute(f"PRAGMA index_info({name})")] + out[name] = (table, cols) + return out + + +def test_store_creates_the_hot_path_indexes(tmp_path): + from passiveworkers.net.store import Store + s = Store(str(tmp_path / "a.db")) + assert _indexes(s.conn) == _EXPECTED + plan = " ".join(r[3] for r in s.conn.execute( + "EXPLAIN QUERY PLAN SELECT node_id FROM nodes WHERE secret_hash=?", ("h",))) + assert "idx_nodes_secret_hash" in plan + plan = " ".join(r[3] for r in s.conn.execute( + "EXPLAIN QUERY PLAN SELECT * FROM tasks WHERE node_id=? AND status='queued' " + "ORDER BY created LIMIT 1", ("n",))) + assert "idx_tasks_node_status" in plan + + +def test_indexes_are_added_to_an_existing_db_and_reopen_is_idempotent(tmp_path): + from passiveworkers.net.store import Store + path = str(tmp_path / "old.db") + Store(path) + raw = sqlite3.connect(path) # simulate a pre-index database + for name in _EXPECTED: + raw.execute(f"DROP INDEX {name}") + raw.commit() + assert _indexes(raw) == {} + raw.close() + Store(path) # open → indexes created + s = Store(path) # open again → no error, no duplicates + assert _indexes(s.conn) == _EXPECTED diff --git a/tests/test_ratelimit.py b/tests/test_ratelimit.py index 2049f6b..c68a301 100644 --- a/tests/test_ratelimit.py +++ b/tests/test_ratelimit.py @@ -39,6 +39,35 @@ def test_sweep_drops_idle_keys_on_overflow(): rl = RateLimiter(max_keys=3) for i in range(5): rl.allow(f"old{i}", 5, 60, now=0) # 5 keys, all hit at t=0 → over capacity - assert len(rl._hits) == 5 + # the first overflow (old3) swept and evicted the LRU key back to the cap; old4 arrived inside + # the same throttle interval, so no second sweep yet (F16) + assert len(rl._hits) == 4 and "old0" not in rl._hits rl.allow("new", 5, 60, now=10_000) # far future → triggers sweep of idle keys assert set(rl._hits) == {"new"} # idle 'old*' keys swept; only 'new' remains + + +def test_sweep_is_throttled_while_over_capacity(): + """F16: once over max_keys, the O(n) sweep used to run on EVERY allow(); now at most once per + _SWEEP_EVERY_S, however many calls arrive in between.""" + from passiveworkers.net import ratelimit as RL + rl = RateLimiter(max_keys=10) + sweeps = [] + real = rl._sweep + rl._sweep = lambda now: (sweeps.append(now), real(now)) + for i in range(1000): # key rotation: every call is a fresh key + rl.allow(f"k{i}", 5, 60, now=100 + i * 0.001) # 1000 calls inside one second + assert len(sweeps) == 1 + rl.allow("later", 5, 60, now=100 + RL._SWEEP_EVERY_S + 1) + assert len(sweeps) == 2 # the throttle re-opens after the interval + + +def test_sweep_evicts_oldest_keys_when_none_are_idle(): + """F16: rotating keys keep every entry fresh, so the idle-horizon sweep alone never shrank the + table. A sweep now evicts the least-recently-hit keys down to max_keys.""" + rl = RateLimiter(max_keys=3) + for i in range(3): + rl.allow(f"k{i}", 5, 60, now=i) + rl.allow("k0", 5, 60, now=3) # k0 is now the MOST recently hit + rl.allow("k3", 5, 60, now=4) # over the cap → sweep, nothing idle + assert set(rl._hits) == {"k0", "k2", "k3"} # k1 (least recently hit) evicted + assert rl.allow("k0", 2, 60, now=5) is False # survivors keep their windows (k0 has 2 hits)