diff --git a/.env.example b/.env.example index 7ed86d8..5432d8d 100644 --- a/.env.example +++ b/.env.example @@ -6,8 +6,9 @@ FUSIONLAB_DEVICE= # User-provided compute over SSH — see the README section "User-provided compute (SSH)" and # docs/COMPUTE.md. These override (or replace) the matching [compute] keys in fusionlab.toml at the -# repo root; the env var wins. No key, password, or token belongs here: SSH auth comes from the -# user's own config, agent, and keys. +# repo root; the env var wins. No key, password, or SSH credential belongs here: SSH auth comes +# from the user's own config, agent, and keys. The one app-level exception is the worker token +# below, whose documented home is exactly here (same precedent as ANTHROPIC_API_KEY). # FUSIONLAB_COMPUTE_PROVIDER=local # FUSIONLAB_COMPUTE_HOST=gpu-lab # FUSIONLAB_SSH_HOST= @@ -16,6 +17,12 @@ FUSIONLAB_DEVICE= # FUSIONLAB_SSH_PORT=8420 # FUSIONLAB_SSH_REMOTE_PORT=8420 +# Optional shared token for the remote worker (#45): set the same value on the worker host's +# environment and here, and every /v1 route on the worker demands `Authorization: Bearer` while +# /health stays open for the tunnel probe. Unset (the default), the worker is loopback-first and +# open, exactly as before the token existed. +# FUSIONLAB_WORKER_TOKEN= + # --- Optional in-app agent: all three are unnecessary for the offline demo ------------------- # Master switch for the agent panel; optional, off unless exactly 1, true, or yes. FUSIONLAB_AGENT_ENABLED= diff --git a/README.md b/README.md index 9029757..8fa04b4 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,11 @@ autostart = true # start the worker over SSH when none is answering More than one GPU box? Define a `[compute.hosts.]` profile per host (same keys) and switch with `FUSIONLAB_COMPUTE_HOST` — see [docs/COMPUTE.md](docs/COMPUTE.md). +By default the worker binds loopback on the remote host and trusts the tunnel. If it might be +reachable by more than your own tunnel, set `FUSIONLAB_WORKER_TOKEN` to the same value on both ends +(`.env` here, the worker's environment there) and every `/v1` route on the worker demands a Bearer +token — same file, no new services; details in [docs/COMPUTE.md](docs/COMPUTE.md). + Full guide — prerequisites, what the chip and banner mean, and a first-run checklist against a real host: [docs/COMPUTE.md](docs/COMPUTE.md). diff --git a/docs/COMPUTE.md b/docs/COMPUTE.md index b7f032f..318934b 100644 --- a/docs/COMPUTE.md +++ b/docs/COMPUTE.md @@ -9,7 +9,8 @@ streams results back. The compute side of the connection is your machine, your account, your SSH setup. FusionLab adds no credentials of its own: it uses your keys, agent, and `~/.ssh/config` exactly as they are, and -neither `fusionlab.toml` nor `.env` ever holds a key, password, or token. +neither `fusionlab.toml` nor `.env` ever holds a key, password, or token. (The one deliberate +exception is the optional worker token below — env-only, on both ends; never in `fusionlab.toml`.) ## Quick start @@ -129,7 +130,8 @@ checklist is the real-host path that remains. remote jobs only. - **The worker executes a fixed registry of named tasks** — field-line tracing, the two surrogates, the model-metrics readout, benchmarks. It is not a shell: schema-checked inputs, no - command execution, payload-size caps. It binds loopback only, reachable through the tunnel. + command execution, payload-size caps. It binds loopback only, reachable through the tunnel — and + if you expose it further, an optional shared token covers the whole `/v1` surface (below). - **Shot data** is read by the worker through the same FAIR-MAST path as everywhere else, from its own clone (which ships the same small showcase cache; other shots fetch on demand). Compact inputs — EFIT equilibria, signal arrays — ride inside the job itself. So the compute host needs @@ -137,6 +139,15 @@ checklist is the real-host path that remains. - **Version drift** between your laptop's clone and the worker's is surfaced, not hidden: the handshake reports the worker's version and the chip warns on mismatch. The fix is the usual one on the remote side — `git pull` and `uv sync`. -- **Not in this version**: SLURM/PBS submission, multi-user queues or auth, training jobs, Windows - remotes, and anything that opens a port on your network. If you cannot install software on the - compute target, this design does not fit it — that is its one real prerequisite. +- **An optional shared token locks the worker's `/v1` surface** (issue #45): if the worker might be + reachable by more than your own tunnel, set `FUSIONLAB_WORKER_TOKEN` to the same value on both + ends (the laptop's `.env` and the worker host's environment) and every `/v1` route on the worker + demands `Authorization: Bearer` — constant-time compare, 401 before any request is parsed or + queued — while `/health` stays open so the tunnel probe and autostart keep working. Unset (the + default), the worker is exactly as open as before. When the two ends disagree, the fallback + banner names the token mismatch instead of a generic "unreachable", and the handshake advertises + the mode (`"auth": "token" | "open"`) so the two sides can see the disagreement coming. +- **Not in this version**: SLURM/PBS submission, multi-user queues or per-user accounts (the token + is one shared secret, not an account system), training jobs, Windows remotes, and anything that + opens a port on your network. If you cannot install software on the compute target, this design + does not fit it — that is its one real prerequisite. diff --git a/fusionlab.toml.example b/fusionlab.toml.example index 508d493..5bfed5f 100644 --- a/fusionlab.toml.example +++ b/fusionlab.toml.example @@ -1,7 +1,9 @@ # FusionLab compute configuration (optional). Copy to fusionlab.toml at the repo root, beside .env. # Every value can also be set with a FUSIONLAB_* environment variable (see .env.example); the env # var wins over this file. Neither place ever holds a key, password, or token — the connection uses -# the user's own SSH config, agent, and keys (design §8). +# the user's own SSH config, agent, and keys (design §8). The worker's shared token is env-only by +# design (FUSIONLAB_WORKER_TOKEN in .env, both ends); a credential-shaped key in this file is a +# loud config error, never a silently-carried option. [compute] # "local" (default): compute runs in this process — zero config, works offline. diff --git a/fusionlab/compute/__init__.py b/fusionlab/compute/__init__.py index 0d7aa09..948c02f 100644 --- a/fusionlab/compute/__init__.py +++ b/fusionlab/compute/__init__.py @@ -36,6 +36,10 @@ "remote_port": "FUSIONLAB_SSH_REMOTE_PORT", } +# The worker's shared token (#45) is env-only: FUSIONLAB_WORKER_TOKEN is deliberately absent from the +# override map above, and a credential-shaped key in the toml is a loud error, never a carried option. +_CREDENTIAL_KEYS = ("token", "worker_token", "password", "secret", "api_key") + def _ssh_table(cfg: dict[str, Any]) -> dict[str, Any]: """The base ssh table for the selected host — raw [compute.ssh]-shaped keys, no env merging yet. @@ -70,7 +74,8 @@ def ssh_options() -> dict[str, Any]: from fusionlab.toml, merged with FUSIONLAB_SSH_* env overrides (env wins). Returns kwargs ready for SshProvider(**...). Never contains credentials — the connection uses the - user's own SSH agent, keys, and ~/.ssh/config (design §8). + user's own SSH agent, keys, and ~/.ssh/config (design §8), and a credential-shaped key in the + toml is a loud ValueError, never a silently-carried option (the worker token is env-only, #45). """ opts: dict[str, Any] = {} if _TOML_PATH.exists(): @@ -88,6 +93,11 @@ def ssh_options() -> dict[str, Any]: opts[key] = int(raw) else: opts[key] = raw + leaked = sorted(set(opts) & set(_CREDENTIAL_KEYS)) + if leaked: + raise ValueError(f"fusionlab.toml holds credential key(s) {leaked} — config never holds a key, " + "password, or token (see fusionlab.toml.example): the worker token is the " + "FUSIONLAB_WORKER_TOKEN environment variable") return opts diff --git a/fusionlab/compute/ssh.py b/fusionlab/compute/ssh.py index 354d1b2..ebe41ec 100644 --- a/fusionlab/compute/ssh.py +++ b/fusionlab/compute/ssh.py @@ -36,6 +36,7 @@ import importlib.metadata import json import logging +import os import threading import time from collections.abc import AsyncIterator @@ -58,6 +59,14 @@ _TERMINAL_EVENTS = ("completed", "error", "cancelled") +def _worker_auth_headers() -> dict[str, str]: + """Authorization headers for every worker HTTP call, from FUSIONLAB_WORKER_TOKEN — empty when no + token is configured (opt-in by absence: a tokenless backend against an open worker is today's + behavior, byte for byte; a token the worker doesn't expect is a plain header it ignores).""" + token = os.environ.get("FUSIONLAB_WORKER_TOKEN", "") + return {"Authorization": f"Bearer {token}"} if token else {} + + class RemoteComputeError(RuntimeError): """Any SshProvider failure the caller should handle by falling back to LocalProvider. @@ -71,6 +80,20 @@ def __init__(self, phase: str, message: str): self.message = message +def _raise_for_auth(r: httpx.Response) -> None: + """Map a worker 401 to the one actionable auth error, at every HTTP call site. + + A 401 means the two ends disagree about FUSIONLAB_WORKER_TOKEN — say exactly that, so the + fallback banner reads "token differs" instead of a generic unreachable (run_task records this + message and GET /compute renders it verbatim).""" + if r.status_code == 401: + raise RemoteComputeError( + "auth", + "worker rejected the token — FUSIONLAB_WORKER_TOKEN differs between this machine and " + "the worker host (or is set on one side only)", + ) + + class WorkerSession(Protocol): """What SshProvider needs from a transport: a loopback HTTP base_url that answers after ensure().""" @@ -174,7 +197,7 @@ async def _wait_worker(self) -> None: started = False while True: try: - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=_worker_auth_headers()) as client: r = await client.get(f"{self.base_url}/health", timeout=2.0) if r.status_code == 200 and r.json().get("ok"): return @@ -255,7 +278,8 @@ def __init__(self, session: WorkerSession | None = None, *, self._session = _SshSession(host, **session_kw) self._result_timeout_s = result_timeout_s self._remote: dict[str, Any] | None = None # the worker's last handshake - self._http = httpx.Client(base_url=self._session.base_url, timeout=30.0) + self._http = httpx.Client(base_url=self._session.base_url, timeout=30.0, + headers=_worker_auth_headers()) # -- ComputeProvider protocol def capabilities(self) -> dict: @@ -281,6 +305,7 @@ def submit(self, task: str, inputs: dict) -> JobHandle: payload = {"task": task, "inputs": frames.encode_inputs(inputs), "artifact_refs": self._ensure_artifacts(task)} r = self._http.post("/v1/jobs", json=payload) + _raise_for_auth(r) if r.status_code == 409: raise RemoteComputeError("submit", r.json().get("detail", "worker busy")) if r.status_code == 422: @@ -301,10 +326,12 @@ async def stream(self, job: JobHandle) -> AsyncIterator[dict]: """ sse_timeouts = httpx.Timeout(connect=10.0, read=None, write=10.0, pool=10.0) try: - async with httpx.AsyncClient(base_url=self._session.base_url, timeout=sse_timeouts) as client, \ + async with httpx.AsyncClient(base_url=self._session.base_url, timeout=sse_timeouts, + headers=_worker_auth_headers()) as client, \ client.stream("GET", f"/v1/jobs/{job.job_id}/events") as resp: if resp.status_code == 404: raise RemoteComputeError("job", f"unknown job {job.job_id} (worker restarted?)") + _raise_for_auth(resp) resp.raise_for_status() event: dict = {} async for line in resp.aiter_lines(): @@ -329,6 +356,7 @@ def cancel(self, job: JobHandle) -> None: r = self._http.post(f"/v1/jobs/{job.job_id}/cancel") except httpx.HTTPError as e: raise RemoteComputeError("job", f"cancel failed: {type(e).__name__}: {e}") from None + _raise_for_auth(r) if r.status_code == 409: return if r.status_code >= 400: @@ -343,6 +371,7 @@ def result(self, job: JobHandle) -> Any: r = self._http.get(f"/v1/jobs/{job.job_id}") except httpx.HTTPError as e: raise RemoteComputeError("collect", f"{type(e).__name__}: {e}") from None + _raise_for_auth(r) if r.status_code == 404: raise RemoteComputeError("collect", f"unknown job {job.job_id} (worker restarted?)") r.raise_for_status() @@ -361,6 +390,7 @@ def result(self, job: JobHandle) -> Any: rr = self._http.get(f"/v1/jobs/{job.job_id}/result") except httpx.HTTPError as e: raise RemoteComputeError("collect", f"{type(e).__name__}: {e}") from None + _raise_for_auth(rr) if rr.status_code != 200: raise RemoteComputeError("collect", f"result fetch failed: HTTP {rr.status_code}") try: @@ -380,6 +410,7 @@ def health(self) -> None: RemoteComputeError, the one typed signal GET /compute renders as "unreachable".""" try: r = self._http.get("/health", timeout=2.0) + _raise_for_auth(r) # /health is open today; if a future worker ever gates it, say "auth" r.raise_for_status() except httpx.HTTPError as e: try: @@ -394,6 +425,7 @@ def _handshake(self) -> dict: try: self._session.ensure() r = self._http.get("/v1/capabilities") + _raise_for_auth(r) r.raise_for_status() self._remote = r.json() except httpx.HTTPError as e: @@ -422,6 +454,7 @@ def _ensure_artifacts(self, task: str) -> list[dict[str, str]]: def _upload(self, path: Path, digest: str) -> None: try: r = self._http.post(f"/v1/artifacts/{digest}", content=path.read_bytes()) + _raise_for_auth(r) if r.status_code != 201: raise RemoteComputeError("submit", f"weight upload failed: HTTP {r.status_code} {r.text[:200]}") logger.info("uploaded models/%s (%d bytes) to the worker cache — one-time per hash", diff --git a/fusionlab/worker.py b/fusionlab/worker.py index a72fbc8..8f63eab 100644 --- a/fusionlab/worker.py +++ b/fusionlab/worker.py @@ -9,13 +9,20 @@ an artifact cache for trained weights: no shell, no eval, no filesystem passthrough. Worst case if an exposed worker were found by a stranger is compute theft, not host compromise (design §8). +For a worker that is not loopback-only, an optional shared token locks the whole /v1 surface: set +FUSIONLAB_WORKER_TOKEN to the same value on both ends (the worker's environment and the backend's) +and every /v1 route demands `Authorization: Bearer` — constant-time compare, 401 before any body +parsing or queue interaction — while /health stays open, because a token-gated probe would be +misread as "worker absent" and trigger a doomed autostart. Unset (the default) nothing changes. + POST /v1/jobs {task, inputs, artifact_refs} -> {job_id, state}; 409 when busy GET /v1/jobs/{id}/events SSE: queued / running / progress / completed / error / cancelled + GET /v1/jobs/{id} state without the payload — the poll target while a job runs GET /v1/jobs/{id}/result application/x-fusionlab-frames (fusionlab.frames), consumed once POST /v1/jobs/{id}/cancel cooperative: immediate while queued; benchmark tasks check between sections POST /v1/artifacts/{sha256} upload one trained weight into the cache (hash-verified, capped) - GET /v1/capabilities handshake: worker/torch/warp versions, CUDA, tasks, weight hashes - GET /health {ok: true} — the tunnel/autostart probe + GET /v1/capabilities handshake: worker/torch/warp versions, CUDA, tasks, weight hashes, auth mode + GET /health {ok: true} — the tunnel/autostart probe, always open One job at a time (a single-user GPU box is honest about itself); a second submission gets a clean 409 instead of a hidden queue. Jobs live in memory only — a worker restart forgets them, and the @@ -28,6 +35,7 @@ import asyncio import contextlib import hashlib +import hmac import importlib.metadata import json import os @@ -35,7 +43,7 @@ import re import threading import time -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass, field from pathlib import Path from typing import Any, Literal @@ -369,6 +377,19 @@ def execute() -> bytes: _JOBS.append(job, {"event": "error", "error": str(e), "error_type": type(e).__name__}) +# ---------------------------------------------------------------- auth (opt-in shared token) +def worker_token() -> str: + """The shared bearer token, read per request (env-only by design — config files never hold one; + tests monkeypatch the environment). Empty means the worker is open: every route behaves exactly + as before the token existed.""" + return os.environ.get("FUSIONLAB_WORKER_TOKEN", "") + + +def auth_mode() -> str: + """What the handshake advertises: 'token' when a shared token is set, else 'open'.""" + return "token" if worker_token() else "open" + + # ---------------------------------------------------------------- handshake def worker_version() -> str: try: @@ -383,7 +404,8 @@ def worker_version() -> str: def handshake() -> dict[str, Any]: """Version/device identity (computed once — a running box does not change its GPU), then the - per-call parts: tasks, shipped weight hashes, and what the artifact cache already holds.""" + per-call parts: tasks, shipped weight hashes, what the artifact cache already holds, and the + auth mode.""" global _HANDSHAKE with _HANDSHAKE_LOCK: if _HANDSHAKE is None: @@ -397,13 +419,34 @@ def handshake() -> dict[str, Any]: _HANDSHAKE = {"worker_protocol_version": WORKER_PROTOCOL_VERSION, "worker_version": worker_version(), "torch": torch.__version__, "warp": warp_version, "cuda_available": cuda, "device_name": torch.cuda.get_device_name(0) if cuda else (platform.processor() or "CPU")} - return {**_HANDSHAKE, "tasks": sorted(TASKS), "models": shipped_models(), "artifacts": cached_artifacts()} + return {**_HANDSHAKE, "tasks": sorted(TASKS), "models": shipped_models(), + "artifacts": cached_artifacts(), "auth": auth_mode()} # ---------------------------------------------------------------- the app app = FastAPI(title="fusionlab worker", version=worker_version(), docs_url=None, redoc_url=None) +@app.middleware("http") +async def require_bearer_token(request: Request, + call_next: Callable[[Request], Awaitable[Response]]) -> Response: + """Opt-in shared-token gate (issue #45): with FUSIONLAB_WORKER_TOKEN set, every /v1 route demands + `Authorization: Bearer ` (constant-time compare) and rejects anything else with 401 + + WWW-Authenticate BEFORE any body parsing or queue interaction. /health stays outside the check: + a token-gated probe would be misread as "worker absent" and trigger a doomed autostart.""" + token = worker_token() + if token and request.url.path.startswith("/v1/"): + supplied = request.headers.get("authorization", "") + if not hmac.compare_digest(supplied, f"Bearer {token}"): + return JSONResponse( + status_code=401, + content={"detail": "missing or mismatched bearer token — set FUSIONLAB_WORKER_TOKEN, " + "the same value on both ends of the tunnel"}, + headers={"WWW-Authenticate": "Bearer"}, + ) + return await call_next(request) + + @app.get("/health") async def health() -> dict: return {"ok": True} diff --git a/tests/test_compute_config.py b/tests/test_compute_config.py index 8c6dc2f..22f51d0 100644 --- a/tests/test_compute_config.py +++ b/tests/test_compute_config.py @@ -169,3 +169,36 @@ def test_unknown_host_fails_at_get_provider(no_env, monkeypatch, tmp_path): monkeypatch.setenv("FUSIONLAB_COMPUTE_HOST", "lab-c") with pytest.raises(ValueError, match="unknown compute host 'lab-c'"): compute.get_provider() + + +# -- worker token is env-only (#45) ------------------------------------------------------------ + + +def test_a_token_in_the_toml_fails_loudly_and_is_never_carried(no_env, monkeypatch, tmp_path): + """The twice-documented no-secrets rule, enforced at runtime: a credential-shaped key in the + toml — either spelling someone might try, in [compute.ssh] or a host profile — is a loud + ValueError naming FUSIONLAB_WORKER_TOKEN, never a silently-carried ssh option.""" + toml = tmp_path / "fusionlab.toml" + toml.write_text('[compute.ssh]\nhost = "gpu-lab"\ntoken = "from-the-toml"\n') + monkeypatch.setattr(compute, "_TOML_PATH", toml) + with pytest.raises(ValueError, match=r"credential key\(s\) \['token'\].*FUSIONLAB_WORKER_TOKEN"): + compute.ssh_options() + toml.write_text('[compute.hosts.lab]\nhost = "gpu-lab"\nworker_token = "from-the-toml"\n') + monkeypatch.setenv("FUSIONLAB_COMPUTE_HOST", "lab") + with pytest.raises(ValueError, match="FUSIONLAB_WORKER_TOKEN"): + compute.ssh_options() + + +def test_worker_token_env_var_is_never_an_ssh_option(no_env, monkeypatch, tmp_path): + """FUSIONLAB_WORKER_TOKEN is deliberately absent from the FUSIONLAB_SSH_* override map and never + rides ssh_options() — the token reaches the worker through the provider's own header builder, + not the ssh kwargs.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "s3cret-env-token") + monkeypatch.setattr(compute, "_TOML_PATH", tmp_path / "absent.toml") + assert compute.ssh_options() == {} + toml = tmp_path / "fusionlab.toml" + toml.write_text('[compute.ssh]\nhost = "gpu-lab"\n') + monkeypatch.setattr(compute, "_TOML_PATH", toml) + opts = compute.ssh_options() + assert opts == {"host": "gpu-lab"} + assert "s3cret-env-token" not in str(opts) diff --git a/tests/test_compute_status.py b/tests/test_compute_status.py index d13d24c..56d0a01 100644 --- a/tests/test_compute_status.py +++ b/tests/test_compute_status.py @@ -209,6 +209,23 @@ def test_run_fallback_survives_a_healthy_probe_until_a_remote_run_succeeds(clean assert status.snapshot(0.0)["fallback"] is False +def test_token_mismatch_falls_back_with_the_token_banner(clean): + """A worker that rejects the token (401 → the auth error): run_task falls back locally and the + banner reason names the token — the operator reads 'the two ends disagree', not 'unreachable'.""" + class TokenMismatchRemote(StubRemote): + def submit(self, task: str, inputs: dict) -> JobHandle: + raise RemoteComputeError( + "auth", "worker rejected the token — FUSIONLAB_WORKER_TOKEN differs between this " + "machine and the worker host (or is set on one side only)") + + compute._provider = TokenMismatchRemote() + out = status.run_task("model_metrics", {"model": "surrogate"}) # the local re-run answers + assert "block_size" in out + fb = status.fallback_state() + assert fb["fallback"] is True + assert "FUSIONLAB_WORKER_TOKEN" in fb["reason"] and "(auth)" in fb["reason"] + + # ---------------------------------------------------------------- job progress (HARN-1) def test_run_task_exposes_the_remote_job_while_it_runs(clean, monkeypatch): """run_task's remote job is visible to active_job() from submit until result returns — the seam diff --git a/tests/test_ssh_provider.py b/tests/test_ssh_provider.py index 579c421..67a9ac1 100644 --- a/tests/test_ssh_provider.py +++ b/tests/test_ssh_provider.py @@ -508,3 +508,124 @@ def slow(_inputs, progress): with pytest.raises(RemoteComputeError, match="cancelled"): p.result(nxt) p.close() + + +# ---------------------------------------------------------------- auth (opt-in shared token, #45) +def test_long_lived_client_carries_the_configured_token(monkeypatch): + """The provider's one long-lived httpx.Client is built with the env token's Authorization header — + and with none at all when the token is unset (opt-in by absence, not an empty header).""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "unit-token") + p = SshProvider(session=FakeSession("http://worker.test")) + assert p._http.headers.get("authorization") == "Bearer unit-token" + p.close() + monkeypatch.delenv("FUSIONLAB_WORKER_TOKEN") + p2 = SshProvider(session=FakeSession("http://worker.test")) + assert "authorization" not in p2._http.headers + p2.close() + + +def test_401_on_submit_becomes_the_actionable_auth_error(): + p = make_provider(json_handler({"GET /v1/capabilities": handshake_response(), + "POST /v1/jobs": httpx.Response(401, headers={"WWW-Authenticate": "Bearer"})})) + with pytest.raises(RemoteComputeError) as ei: + p.submit("benchmark", {"profile": "quick"}) + assert ei.value.phase == "auth" + assert "FUSIONLAB_WORKER_TOKEN" in ei.value.message and "differs" in ei.value.message + p.close() + + +def test_401_on_the_handshake_becomes_the_auth_error(): + """A tokenless backend against a token-gated worker: the handshake is the first /v1 call to + fail, and it fails with the token message — the banner's wording, not a generic connect error.""" + p = make_provider(json_handler({"GET /v1/capabilities": httpx.Response(401, headers={"WWW-Authenticate": "Bearer"})})) + with pytest.raises(RemoteComputeError) as ei: + p.capabilities() + assert ei.value.phase == "auth" and "FUSIONLAB_WORKER_TOKEN" in ei.value.message + p.close() + + +def test_401_on_the_stream_becomes_the_auth_error(monkeypatch): + """stream() builds its own AsyncClient, so the fake transport rides that construction (the same + patch the probe test uses) — then a 401 on the events route maps to the auth error, not the + generic 'event stream broke'.""" + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/events"): + return httpx.Response(401, headers={"WWW-Authenticate": "Bearer"}) + return httpx.Response(200, json=handshake_response()) + + real_cls = httpx.AsyncClient + + def with_mock_transport(*a, **kw): + kw["transport"] = httpx.MockTransport(handler) + return real_cls(*a, **kw) + + monkeypatch.setattr(ssh_mod.httpx, "AsyncClient", with_mock_transport) + p = make_provider(handler) + with pytest.raises(RemoteComputeError) as ei: + asyncio.run(collect(p.stream(JobHandle(job_id="j1", task="benchmark", provider="ssh")))) + assert ei.value.phase == "auth" + p.close() + + +def test_health_probe_client_rides_the_token_too(fake_asyncssh, worker_url, monkeypatch): + """The autostart probe is the third httpx client: it carries the header as well — harmless on + the always-open /health, correct if a future worker ever gates it.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "probe-token") + captured: dict = {} + real = ssh_mod.httpx.AsyncClient + + def spy(*a, **kw): + captured["headers"] = kw.get("headers") + return real(*a, **kw) + + monkeypatch.setattr(ssh_mod.httpx, "AsyncClient", spy) + session = _session(local_port=_free_url_port(worker_url)) + session.ensure() + assert captured["headers"] == {"Authorization": "Bearer probe-token"} + session.close() + + +def test_loopback_token_end_to_end(worker_url, monkeypatch, tmp_path): + """The loopback acceptance with a shared token: handshake, submit, stream, result, upload and + cancel all green through the real token-gated worker, every call bearing the Bearer header.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "loopback-token") + p = SshProvider(session=FakeSession(worker_url)) + caps = p.capabilities() # GET /v1/capabilities + assert caps["remote"]["auth"] == "token" + weight = tmp_path / "loopback-weight.pt" # POST /v1/artifacts/{sha} + weight.write_bytes(b"loopback-token-e2e") + p._upload(weight, hashlib.sha256(weight.read_bytes()).hexdigest()) + job = p.submit("model_metrics", {"model": "surrogate"}) # POST /v1/jobs + events = asyncio.run(collect(p.stream(job))) # the SSE client + assert events[-1]["event"] == "completed" + out = p.result(job) # poll + result fetch + assert "block_size" in out + + from dataclasses import replace as dc_replace + + def slow(_inputs, progress): + for _ in range(40): + progress({"section": "waiting"}) + time.sleep(0.05) + return {} + + monkeypatch.setitem(worker.TASKS, "model_metrics", dc_replace(worker.TASKS["model_metrics"], run=slow)) + slow_job = p.submit("model_metrics", {"model": "surrogate"}) + p.cancel(slow_job) # POST cancel + with pytest.raises(RemoteComputeError, match="cancelled"): + p.result(slow_job) + p.close() + + +def test_loopback_token_mismatch_is_an_actionable_auth_error(worker_url, monkeypatch): + """Two ends, two environments — simulated by re-pointing the shared env after the provider + captured its header: the worker rejects the token with 401 and the provider raises the + actionable auth error run_task's fallback banner renders.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "the-worker-expects") + p = SshProvider(session=FakeSession(worker_url)) # sends Bearer the-worker-expects + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "the-provider-sends") # the worker now demands this + with pytest.raises(RemoteComputeError) as ei: + p.capabilities() + assert ei.value.phase == "auth" + assert "FUSIONLAB_WORKER_TOKEN" in ei.value.message and "differs" in ei.value.message + p.close() diff --git a/tests/test_worker.py b/tests/test_worker.py index 819fa52..8c031e3 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -31,10 +31,10 @@ def client(monkeypatch, tmp_path): worker._HANDSHAKE = None -def _wait_result(client, job_id: str, tries: int = 240): +def _wait_result(client, job_id: str, tries: int = 240, headers: dict | None = None): """Poll until the job is terminal; returns the 200 (frames) or 500 (error JSON) response.""" for _ in range(tries): - r = client.get(f"/v1/jobs/{job_id}/result") + r = client.get(f"/v1/jobs/{job_id}/result", headers=headers) if r.status_code in (200, 500): return r assert r.status_code == 409, f"unexpected: {r.status_code} {r.text[:300]}" @@ -335,3 +335,90 @@ def boom(_inputs, _progress): job_id = client.post("/v1/jobs", json={"task": "model_metrics", "inputs": {"model": "surrogate"}}).json()["job_id"] r = client.get(f"/v1/jobs/{job_id}/result") assert r.status_code == 500 and r.json()["error_type"] == "ZeroDivisionError" and "boom" in r.json()["detail"] + + +# ---------------------------------------------------------------- auth (opt-in shared token, #45) +AUTH_ROUTES = [ # (method, path, body): every /v1/* route — the 401 matrix walks all of them + ("get", "/v1/capabilities", None), + ("post", "/v1/jobs", {"task": "model_metrics", "inputs": {"model": "surrogate"}}), + ("get", "/v1/jobs/ghost", None), + ("get", "/v1/jobs/ghost/events", None), + ("get", "/v1/jobs/ghost/result", None), + ("post", "/v1/jobs/ghost/cancel", None), + ("post", f"/v1/artifacts/{'0' * 64}", b"x"), +] +_AUTH_IDS = [f"{m.upper()} {p}" for m, p, _ in AUTH_ROUTES] + + +def _call(client, method: str, path: str, body, headers: dict | None = None): + if isinstance(body, bytes): + return client.request(method, path, content=body, headers=headers) + if body is not None: + return client.request(method, path, json=body, headers=headers) + return client.request(method, path, headers=headers) + + +def test_token_unset_is_todays_behavior(client): + """Opt-in by absence: with no FUSIONLAB_WORKER_TOKEN, /v1 routes answer requests that carry no + Authorization header at all — the pre-token worker, pinned here (the other tests pin the rest).""" + assert worker.worker_token() == "" + assert client.get("/v1/capabilities").status_code == 200 + assert client.post("/v1/jobs", json={"task": "model_metrics", + "inputs": {"model": "surrogate"}}).status_code == 201 + + +def test_handshake_advertises_the_auth_mode(client, monkeypatch): + assert client.get("/v1/capabilities").json()["auth"] == "open" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "tk") + r = client.get("/v1/capabilities", headers={"Authorization": "Bearer tk"}) + assert r.json()["auth"] == "token" + + +@pytest.mark.parametrize(("method", "path", "body"), AUTH_ROUTES, ids=_AUTH_IDS) +def test_401_when_the_token_header_is_missing(client, monkeypatch, method, path, body): + """Auth runs before routing, validation, and queue state: a tokenless request to any /v1 route + is a 401 with WWW-Authenticate — never the 404/422/201 that route itself would have answered.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "tk") + r = _call(client, method, path, body) + assert r.status_code == 401 and r.headers["www-authenticate"] == "Bearer" + + +@pytest.mark.parametrize(("method", "path", "body"), AUTH_ROUTES, ids=_AUTH_IDS) +def test_401_when_the_token_is_wrong(client, monkeypatch, method, path, body): + """Same matrix, wrong Bearer value: still a 401 before the route ever sees the request.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "tk") + r = _call(client, method, path, body, headers={"Authorization": "Bearer not-the-token"}) + assert r.status_code == 401 and r.headers["www-authenticate"] == "Bearer" + + +def test_a_401_submission_never_touches_the_queue(client, monkeypatch): + """The gate sits before body parsing and the depth-1 book: a rejected submit leaves no job state + behind and the slot is free for the next (correctly authed) one.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "tk") + r = client.post("/v1/jobs", json={"task": "model_metrics", "inputs": {"model": "surrogate"}}, + headers={"Authorization": "Bearer wrong"}) + assert r.status_code == 401 + assert worker._JOBS._jobs == {} + ok = client.post("/v1/jobs", json={"task": "model_metrics", "inputs": {"model": "surrogate"}}, + headers={"Authorization": "Bearer tk"}) + assert ok.status_code == 201 + + +def test_health_stays_open_with_a_token_set(client, monkeypatch): + """The tunnel/autostart probe sends no token: /health must answer it even when the worker is + token-gated, or a running worker reads as absent and autostart starts a doomed second one.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "tk") + assert client.get("/health").json() == {"ok": True} + + +def test_the_correct_token_passes_end_to_end(client, monkeypatch): + """The happy path with the gate on: handshake and a full job roundtrip, every call bearing the + Bearer header.""" + monkeypatch.setenv("FUSIONLAB_WORKER_TOKEN", "tk") + headers = {"Authorization": "Bearer tk"} + assert client.get("/v1/capabilities", headers=headers).json()["auth"] == "token" + job_id = client.post("/v1/jobs", json={"task": "model_metrics", "inputs": {"model": "surrogate"}}, + headers=headers).json()["job_id"] + rr = _wait_result(client, job_id, headers=headers) + assert rr.status_code == 200 + assert "block_size" in frames.unpack_result(rr.content)