diff --git a/docs/source/en/concepts/sandbox.md b/docs/source/en/concepts/sandbox.md index 0585850a38..a6553c5ada 100644 --- a/docs/source/en/concepts/sandbox.md +++ b/docs/source/en/concepts/sandbox.md @@ -76,7 +76,7 @@ One nonce is minted per **Job**, so the derived token is a *job* credential. In | **pool host token** | `HMAC(hf_token, nonce)` from the host job's label | pool management: create, list and delete sandboxes on that host, and recover their tokens | | **pooled sandbox token** | random 256 bits, minted by the host server per sandbox | that one sandbox — not a sibling, not the pool | -The client uses the narrow one automatically: `pool.create()` receives it in the create response, `Sandbox.connect(".")` recovers it with the host token, and [`proxy_headers`] hands out the sandbox's token rather than the host's — those headers usually end up in a browser or WebSocket client, so they should confer access to one sandbox and nothing more. +The client uses the narrow one automatically: `pool.create()` receives it in the create response, `Sandbox.connect(".")` recovers it with the host token, and [`proxy_headers`] hands out the sandbox's token rather than the host's (resolving the HF bearer at the moment you read it, so a long-lived handle does not hand out a stale one) — those headers usually end up in a browser or WebSocket client, so they should confer access to one sandbox and nothing more. What this means for a leak: a **pooled sandbox token** compromises that sandbox. A **host token** compromises the host — every sandbox on it, current and future, plus its management routes — so treat it as the pool's admin credential. Members of your namespace hold a different HF token and cannot derive yours, but see [Known limitations](#known-limitations) for how a token can be *delivered* to the wrong place. @@ -225,7 +225,7 @@ This section is deliberately exhaustive rather than reassuring: if you are decid | --- | --- | --- | | Anonymous internet user | Cannot reach the sandbox: the Jobs proxy requires an HF token with namespace read access. | Same. | | Namespace member with read access | Reaches the proxy, but not the API: they cannot derive your sandbox token. Can see the job exists, its labels, and `/health`. | Same, plus they can read the pool's labels and nonce. | -| Namespace member who can create Jobs | — | Can publish a Job carrying your pool's labels and nonce. Your client may adopt it as a host and send it the token derived for your real host (see "Host discovery" below). | +| Namespace member who can create Jobs | — | Can publish a Job carrying your pool's labels and nonce, but it is not adopted: `adopt_hosts` defaults to hosts this principal started. Relevant again if you opt into `adopt_hosts="namespace"`. | | Code running in the sandbox | Runs as root in the VM alongside `sbx-server`; owns the VM. | Confined by uid + Landlock as described above, but shares the kernel, VM and control plane with its neighbours. | | Your own client | Holds the HF token and the sandbox token. | Same; in a pool the sandbox token covers the whole host. | @@ -236,7 +236,7 @@ This section is deliberately exhaustive rather than reassuring: if you are decid - the **port proxy** connects to `$SBX_PROXY_DIR/.sock` without rejecting symlinks or checking the socket's owner, so a sandbox can point it at another sandbox's socket. Both require a legitimate caller to invoke the endpoint (the sandbox has no token of its own), so they are confused-deputy problems rather than direct escapes — but they do break confidentiality and integrity between pooled sandboxes. -- **Host discovery trusts Job labels.** Hosts are found by filtering Jobs on labels, which any Job creator in the namespace can set, and the nonce that derives the token is a public label. Nothing binds a Job to its creator, image, or pool. In a namespace whose members do not all trust each other, prefer `Sandbox.create`, or use a namespace you control for pools. +- **Host discovery starts from Job labels**, which any Job creator in the namespace can set — so a label match is a claim, not proof. By default (`adopt_hosts="own"`) the client only adopts hosts *this principal started*, per the Jobs API's `initiator`, and additionally checks the image, flavor, command and exposed URL. Setting `adopt_hosts="namespace"` restores cross-user host sharing and, with it, the ability of any Job creator in that namespace to publish a host your client will send a token to — only use it in a namespace whose members you trust. - **Landlock can degrade silently.** If Landlock is unavailable, or its ruleset cannot be built, the server currently falls back to uid-only isolation and creates the sandbox anyway — without telling the client. Under uid-only isolation, `/tmp`, `/dev/shm`, TCP bind and cross-home filesystem access are *not* denied. The server also accepts Landlock ABI 1, while the ✅ list above needs ABI 4 (TCP bind) and ABI 6 (abstract sockets); production kernels provide ABI 6, but a lower one would silently drop those two guarantees. - **Residual shared channels**, none of which Landlock or uid isolation closes: unrestricted outbound TCP; loopback access to the control server; **UDP bind is allowed** (Landlock has no UDP coverage); a sibling's `/proc//cmdline` and `status` are readable (`environ` is not — that is the part that would leak credentials, and it is denied); `/proc` and `/sys` are readable and `/dev` is broadly readable and writable; kernel IPC and all machine resources are shared. - **No CPU, disk, FD or total-memory quotas.** Only per-process `RLIMIT_NPROC` and `RLIMIT_AS` are set; cgroup delegation is not available on Jobs. One sandbox can starve its neighbours. The `max_procs`/`max_mem_mb` values are caller-supplied and not clamped server-side. diff --git a/src/huggingface_hub/_sandbox.py b/src/huggingface_hub/_sandbox.py index 760465b095..3997990ef1 100644 --- a/src/huggingface_hub/_sandbox.py +++ b/src/huggingface_hub/_sandbox.py @@ -23,6 +23,7 @@ from pathlib import Path from secrets import token_hex from typing import Any, BinaryIO, Callable, Iterator, List, Literal, overload +from urllib.parse import urlparse import httpx @@ -66,6 +67,13 @@ DEFAULT_SANDBOXES_PER_HOST = 50 +# Which pool hosts `create()` may adopt from job labels. Labels are set by whoever +# creates a Job, so "any job carrying our pool's labels" is not an identity claim -- +# see the `adopt_hosts` argument of [`SandboxPool`]. +ADOPT_OWN = "own" +ADOPT_NAMESPACE = "namespace" +ADOPT_NEVER = "never" + SHARED_ID_SEP = "." # Job stages in which a sandbox/host is finished and needs no teardown. @@ -354,6 +362,29 @@ def _raise_for_status(response: httpx.Response) -> None: raise SandboxError(f"Sandbox API error ({response.status_code}): {message}", status_code=response.status_code) +class _SandboxAuth(httpx.Auth): + """Attach the *current* HF bearer to every request. + + The bearer used to be captured once when the transport was built. A sandbox + host lives up to 24h, so an OIDC/OAuth credential can expire or rotate inside + that window, after which every request to the Jobs proxy failed even though + the handle was still perfectly good. Reading it per request fixes that. + + The `X-Sandbox-Token` is deliberately *not* re-derived: the server holds the + value derived from the bearer that created the job, so re-deriving from a + rotated bearer would produce a token the server has never seen. The two + credentials answer different gates -- the proxy wants a live bearer, the + server wants the original capability -- and only the first needs refreshing. + """ + + def __init__(self, api: HfApi) -> None: + self._api = api + + def auth_flow(self, request): + request.headers["Authorization"] = f"Bearer {_effective_token(self._api)}" + yield request + + class _SandboxServer: """HTTP transport to one `sbx-server` instance — a dedicated job or a shared host. @@ -394,10 +425,8 @@ def __init__( # httpx.Client is thread-safe, so a single client serves both sequential requests # and the concurrent workers used for parallel file transfers / many sandboxes. self._client = httpx.Client( - headers={ - "Authorization": f"Bearer {self._auth_token}", - "X-Sandbox-Token": sandbox_token, - }, + headers={"X-Sandbox-Token": sandbox_token}, + auth=_SandboxAuth(api), limits=httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections), follow_redirects=True, ) @@ -977,12 +1006,16 @@ def proxy_url_for(self, port: int | str, path: str = "/", *, scheme: str = "http def proxy_headers(self) -> dict[str, str]: """Auth headers to send with [`proxy_url_for`] requests (HF token + sandbox token). + Read these at the moment you use them: the HF token is resolved on access, + so a long-lived handle hands out a current bearer rather than one captured + when the sandbox was created. + For a pooled sandbox this is the sandbox's own capability token, not the pool host's: these headers typically end up in a browser or WebSocket client, so they should confer access to this sandbox and nothing else. """ return { - "Authorization": f"Bearer {self._server._auth_token}", + "Authorization": f"Bearer {_effective_token(self._server._api)}", "X-Sandbox-Token": self._sandbox_token or self._server._sandbox_token, } @@ -1059,6 +1092,7 @@ def __init__( idle_timeout: int | float | str | None = DEFAULT_IDLE_TIMEOUT, namespace: str | None = None, start_timeout: float = 120.0, + adopt_hosts: str = ADOPT_OWN, token: str | None = None, _connect_mode: bool = False, ) -> None: @@ -1101,9 +1135,26 @@ def __init__( User or org namespace to run hosts under. start_timeout (`float`, *optional*, defaults to `120.0`): Max seconds to wait for a host to become ready. + adopt_hosts (`str`, *optional*, defaults to `"own"`): + Which already-running hosts `create()` may pack onto. Hosts are found by + filtering Jobs on labels, and labels are set by whoever creates the Job — so a + label match is a claim, not proof. This decides how much more than the claim is + required: + + - `"own"` (default): only Jobs this principal started, per the Jobs API's + `initiator`. Safe in a shared namespace. + - `"namespace"`: any Job in the namespace carrying the pool's labels, provided + its image, flavor, command and exposed URL all match. Restores cross-user + host sharing; only use it in a namespace whose members you trust, since any + of them can publish a Job your client will then send a token to. + - `"never"`: no adoption at all — this handle only uses hosts it booted itself. token (`str`, *optional*): HF token override. """ + if adopt_hosts not in (ADOPT_OWN, ADOPT_NAMESPACE, ADOPT_NEVER): + raise ValueError( + f"adopt_hosts must be one of {ADOPT_OWN!r}, {ADOPT_NAMESPACE!r}, {ADOPT_NEVER!r}; got {adopt_hosts!r}." + ) if sandboxes_per_host < 1: raise ValueError("sandboxes_per_host must be >= 1.") if warm_up < 1: @@ -1118,6 +1169,11 @@ def __init__( self._idle_timeout = idle_timeout self._namespace = namespace self._start_timeout = start_timeout + self._adopt_hosts = adopt_hosts + # Resolved lazily and only when `adopt_hosts="own"` needs it: whoami() is + # heavily rate-limited, and a pool that never discovers never needs it. + self._principal: str | None = None + self._principal_resolved = False self._hosts: List[_SandboxServer] = [] self._lock = threading.Lock() # Held across the whole one-time warm-up so concurrent first create() calls block until @@ -1151,7 +1207,14 @@ def __init__( # ------------------------------------------------------------------ public API @classmethod - def connect(cls, pool_id: str, *, namespace: str | None = None, token: str | None = None) -> "SandboxPool": + def connect( + cls, + pool_id: str, + *, + namespace: str | None = None, + adopt_hosts: str = ADOPT_OWN, + token: str | None = None, + ) -> "SandboxPool": """Reattach to a running pool by id, from any machine — no local state needed. Finds a running host labelled with `pool_id` and rebuilds the pool's config @@ -1167,6 +1230,9 @@ def connect(cls, pool_id: str, *, namespace: str | None = None, token: str | Non The id returned when the pool was first created. namespace (`str`, *optional*): Namespace to search for the pool's hosts (defaults to yours). + adopt_hosts (`str`, *optional*, defaults to `"own"`): + Which hosts may be attached to. See [`SandboxPool`]. Reattaching to a pool whose + hosts another member of the namespace started needs `"namespace"`. token (`str`, *optional*): HF token override. """ @@ -1183,13 +1249,14 @@ def connect(cls, pool_id: str, *, namespace: str | None = None, token: str | Non name=pool_id, idle_timeout=cache.idle_timeout, namespace=cache.namespace if namespace is None else namespace, + adopt_hosts=adopt_hosts, token=token, _connect_mode=True, # attach to existing hosts; never boot during construction ) # Cold path: find a running host via labels and rebuild the config from its job spec. api = HfApi(token=token) - job = _find_pool_host_job(api, pool_id, namespace=namespace) + job = _find_pool_host_job(api, pool_id, namespace=namespace, policy=adopt_hosts) env = _host_env(api, job, namespace=namespace) idle_raw = env.get("SBX_IDLE_TIMEOUT") max_hosts_raw = env.get("SBX_MAX_HOSTS") @@ -1201,6 +1268,7 @@ def connect(cls, pool_id: str, *, namespace: str | None = None, token: str | Non name=pool_id, idle_timeout=int(idle_raw) if idle_raw is not None else None, namespace=namespace, + adopt_hosts=adopt_hosts, token=token, _connect_mode=True, # attach to existing hosts; never boot during construction ) @@ -1435,6 +1503,48 @@ def _ensure_warmed_up(self) -> None: if self._hosts: self._save_cache() + def _principal_id(self) -> str | None: + """Id of the principal this pool authenticates as, for the `"own"` policy.""" + if not self._principal_resolved: + self._principal_resolved = True + try: + self._principal = self._api.whoami(cache=True).get("id") + except Exception as e: + # Leave it None: `_adoptable` then refuses rather than guessing. + logger.warning(f"Could not resolve the current user to validate pool hosts: {e}") + return self._principal + + def _adoptable(self, job: JobInfo) -> bool: + """Whether `job` may be adopted as one of this pool's hosts.""" + reason = _host_rejection( + job, + policy=self._adopt_hosts, + principal_id=self._principal_id() if self._adopt_hosts == ADOPT_OWN else None, + namespace=self._namespace, + image=self.image, + flavor=self.flavor, + ) + if reason is not None: + logger.debug(f"Not adopting job {job.id} as a host for pool '{self.name}': {reason}.") + return False + return True + + def _adoptable_pending(self, job: JobInfo) -> bool: + """Like `_adoptable`, minus the checks a not-yet-running job cannot pass.""" + reason = _host_rejection( + job, + policy=self._adopt_hosts, + principal_id=self._principal_id() if self._adopt_hosts == ADOPT_OWN else None, + namespace=self._namespace, + image=self.image, + flavor=self.flavor, + check_url=False, + ) + if reason is not None: + logger.debug(f"Not waiting for job {job.id} as a host for pool '{self.name}': {reason}.") + return False + return True + def _reserve_one(self) -> "_SandboxServer | None": """Reserve one slot on the first host with free capacity (under lock), else None.""" with self._lock: @@ -1489,7 +1599,10 @@ def _adopt_pending_host(self) -> bool: labels={MODE_LABEL: MODE_POOL, POOL_LABEL: self.name}, namespace=self._namespace, ) - if job.id not in known + # A SCHEDULING job has no exposed URL yet, so only the ownership + # and spec checks apply here; discovery re-validates it in full + # once it is RUNNING. + if job.id not in known and self._adoptable_pending(job) ), None, ) @@ -1531,7 +1644,7 @@ def _discover_hosts(self) -> None: labels={MODE_LABEL: MODE_POOL, POOL_LABEL: self.name}, namespace=self._namespace, ) - if job.id not in known + if job.id not in known and self._adoptable(job) ] for job in matches: @@ -1769,6 +1882,91 @@ def _save_cache(self) -> None: ) +def _normalize_image(image: str | None) -> str | None: + """Compare Docker images leniently: registries and tags get normalized + server-side, and a mis-parse here must not stop a pool from working.""" + if image is None: + return None + image = image.strip().lower() + for prefix in ("docker.io/", "index.docker.io/", "library/"): + if image.startswith(prefix): + image = image[len(prefix) :] + return image + + +def _server_url_rejection(job: JobInfo) -> str | None: + """Why this Job's exposed URL is not one we should send credentials to. + + The hostname must be derived from *this* job's id, so a Job cannot name an + arbitrary destination for the credentials that are about to follow. The + domain is not hard-coded, so staging endpoints keep working. + """ + urls = list(job.status.expose_urls or []) if job.status is not None else [] + if len(urls) != 1: + return f"exposes {len(urls)} URL(s), expected exactly one for port {SANDBOX_SERVER_PORT}" + parsed = urlparse(urls[0]) + if parsed.scheme != "https": + return f"exposes the server over {parsed.scheme!r}, not https" + # `urlparse` lowercases the hostname, so compare case-insensitively. + if not (parsed.hostname or "").startswith(f"{job.id}--{SANDBOX_SERVER_PORT}.".lower()): + return f"exposes {urls[0]!r}, which is not this job's port-{SANDBOX_SERVER_PORT} URL" + return None + + +def _host_rejection( + job: JobInfo, + *, + policy: str, + principal_id: str | None, + namespace: str | None, + image: str | None, + flavor: str | None, + check_url: bool = True, +) -> str | None: + """Why this Job must not be adopted as one of our pool hosts, or `None`. + + Hosts are found by filtering Jobs on labels, and labels are set by whoever + creates the Job -- so a label match says "this Job claims to be one of our + hosts", not "this Job is one of our hosts". Anyone who can create a Job in + the namespace can make that claim, copy the public nonce from a real host's + labels, and receive the token derived for that host. + + Everything checked here is asserted by the backend rather than by the Job's + creator: the initiator, the owner, the image, the flavor, the command, and + the exposed URL. `initiator` is the load-bearing one -- there is no + client-side way to set it -- and it is what makes `policy="own"` meaningful. + """ + if policy == ADOPT_NEVER: + return "host adoption is disabled (adopt_hosts='never')" + + if policy == ADOPT_OWN: + initiator = job.initiator + if initiator is None or initiator.id is None: + return "the Jobs API did not say who started it, so it cannot be confirmed as ours" + if principal_id is None: + return "could not resolve the current principal to compare against its initiator" + if initiator.id != principal_id: + return f"started by a different principal ({initiator.name or initiator.id})" + + if namespace is not None and job.owner is not None and job.owner.name != namespace: + return f"runs under namespace {job.owner.name!r}, not {namespace!r}" + + job_image = job.docker_image or job.space_id + if image is not None and _normalize_image(job_image) != _normalize_image(image): + return f"runs image {job_image!r}, not {image!r}" + + if flavor is not None and job.flavor is not None and str(job.flavor) != flavor: + return f"runs flavor {str(job.flavor)!r}, not {flavor!r}" + + # A sandbox host runs the bootstrap script and nothing else. A Job running + # anything else is not a host, whatever its labels say. + if job.command is not None and list(job.command) != _bootstrap_command(): + return "does not run the sandbox bootstrap command" + + # A SCHEDULING job has no exposed URL yet; the caller re-checks once it runs. + return _server_url_rejection(job) if check_url else None + + def _effective_token(api: HfApi) -> str: token = api.token if isinstance(api.token, str) else get_token() if not token: @@ -1810,8 +2008,12 @@ def _bootstrap_job_spec( job_volumes.append( Volume(type="bucket", source=constants.SANDBOX_SERVER_BUCKET, mount_path=_SERVER_MOUNT_PATH, read_only=True) ) - command = ["/bin/sh", "-c", _BOOTSTRAP_DOWNLOAD] - return command, job_env, job_secrets, job_volumes + return _bootstrap_command(), job_env, job_secrets, job_volumes + + +def _bootstrap_command() -> list[str]: + """The Job command every sandbox host and dedicated sandbox runs.""" + return ["/bin/sh", "-c", _BOOTSTRAP_DOWNLOAD] def _host_env(api: HfApi, job: JobInfo, *, namespace: str | None) -> dict[str, Any]: @@ -1823,12 +2025,37 @@ def _host_env(api: HfApi, job: JobInfo, *, namespace: str | None) -> dict[str, A return env -def _find_pool_host_job(api: HfApi, pool_id: str, *, namespace: str | None = None) -> JobInfo: - """Return any running host job belonging to `pool_id` (found via the pool label).""" +def _find_pool_host_job(api: HfApi, pool_id: str, *, namespace: str | None = None, policy: str = ADOPT_OWN) -> JobInfo: + """Return a running host job belonging to `pool_id`, found via the pool label. + + This one feeds `SandboxPool.connect`'s cold path, which rebuilds the pool's + whole configuration (image, flavor, density, caps) from the job it returns -- + so an unvalidated match would let a Job in the namespace choose the image + that the pool's *next* hosts boot. Image and flavor are unknown here (they + are what we are about to learn), so this checks ownership and shape. + """ + principal = None + if policy == ADOPT_OWN: + try: + principal = api.whoami(cache=True).get("id") + except Exception as e: + logger.warning(f"Could not resolve the current user to validate pool hosts: {e}") + rejected = [] for job in api.list_jobs( status="RUNNING", labels={MODE_LABEL: MODE_POOL, POOL_LABEL: pool_id}, namespace=namespace ): - return job + reason = _host_rejection( + job, policy=policy, principal_id=principal, namespace=namespace, image=None, flavor=None + ) + if reason is None: + return job + rejected.append(f"{job.id} ({reason})") + if rejected: + raise SandboxError( + f"Found host job(s) labelled for pool '{pool_id}' but none usable: {'; '.join(rejected)}. " + "Pass adopt_hosts='namespace' to SandboxPool if you intend to share hosts with other " + "members of this namespace." + ) raise SandboxError( f"No running host found for pool '{pool_id}'. The pool has stopped " "(all its hosts were killed or idle-timed-out); create a new one." @@ -1844,6 +2071,14 @@ def _connect_host(api: HfApi, host_job_id: str, *, namespace: str | None = None) raise SandboxError(f"Job {host_job_id} is not a sandbox host.") if job.status.stage != "RUNNING": raise SandboxError(f"Sandbox host {host_job_id} is not running (status: {job.status.stage}).") + # The caller named this job explicitly, so who started it is their call -- + # but the credentials below must still only reach this job's own URL, over + # HTTPS, on a job that actually runs the sandbox bootstrap. + reason = _host_rejection( + job, policy=ADOPT_NAMESPACE, principal_id=None, namespace=namespace, image=None, flavor=None + ) + if reason is not None: + raise SandboxError(f"Job {host_job_id} does not look like a usable sandbox host: {reason}.") return _SandboxServer.from_job( job=job, nonce=nonce, diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py index 5f9213df3b..6e104d3aca 100644 --- a/tests/test_sandbox.py +++ b/tests/test_sandbox.py @@ -386,9 +386,7 @@ def test_falls_back_to_the_host_token_on_an_older_server(self, fake_server: str) # in the create response; those sandboxes keep working with the host one. server = _make_server(fake_server, capacity=10) _FakeServer.sandboxes.add("local1") - sandbox = Sandbox( - id="job123.local1", server=server, local_id="local1", owns_sandbox=True, owns_server=False - ) + sandbox = Sandbox(id="job123.local1", server=server, local_id="local1", owns_sandbox=True, owns_server=False) assert sandbox._sandbox_token is None assert sandbox.proxy_headers["X-Sandbox-Token"] == "secret" @@ -474,21 +472,57 @@ def test_full_host_triggers_duplicate(self, monkeypatch) -> None: assert {b.id.split(".")[0] for b in boxes} == {"h0", "h1"} +# Id `whoami()` returns in these tests; a host must be started by this principal +# to be adopted under the default policy. +PRINCIPAL_ID = "principal-self" + + +@pytest.fixture(autouse=True) +def _principal(monkeypatch): + """Resolve the calling principal without a network call. + + Host adoption compares a Job's backend-asserted initiator against this, so + every test that discovers a host needs it. + """ + monkeypatch.setattr(sandbox_mod.HfApi, "whoami", lambda self, **kwargs: {"id": PRINCIPAL_ID}) + + +def _pool_host_job( + job_id: str = "host9", + *, + capacity: int = 4, + pool_name: str = "p1", + image: str = "python:3.12", + env: dict | None = None, +) -> MagicMock: + """A Job as the Jobs API would really describe one of our pool hosts. + + Everything beyond the labels is asserted by the backend and checked by the + client before a credential is sent, so a fixture that set only labels would + let an adoption bug through unnoticed. + """ + job = MagicMock() + job.id = job_id + job.owner.name = "user" + job.initiator.id = PRINCIPAL_ID + job.initiator.name = "user" + job.docker_image = image + job.space_id = None + job.flavor = "cpu-basic" + job.command = ["/bin/sh", "-c", sandbox_mod._BOOTSTRAP_DOWNLOAD] + job.status.stage = "RUNNING" + job.status.expose_urls = [f"https://{job_id}--{sandbox_mod.SANDBOX_SERVER_PORT}.hf.jobs"] + job.labels = {SANDBOX_LABEL: "1", MODE_LABEL: MODE_POOL, POOL_LABEL: pool_name, NONCE_LABEL: "nonce"} + job.environment = {"SBX_CAPACITY": str(capacity)} if env is None else env + return job + + class TestHostDiscovery: """`create()` should attach to a warm host found via job labels (e.g. left by another process) before booting a new one.""" def _host_job(self, job_id: str = "host9", capacity: int = 4, pool_name: str = "p1") -> MagicMock: - job = MagicMock() - job.id = job_id - job.owner.name = "user" - job.docker_image = "python:3.12" - job.space_id = None - job.flavor = "cpu-basic" - job.status.stage = "RUNNING" - job.labels = {SANDBOX_LABEL: "1", MODE_LABEL: MODE_POOL, POOL_LABEL: pool_name, NONCE_LABEL: "nonce"} - job.environment = {"SBX_CAPACITY": str(capacity)} # config lives in env vars, not labels - return job + return _pool_host_job(job_id, capacity=capacity, pool_name=pool_name) def _pool(self, fake_server, monkeypatch, jobs, name: str = "p1", **kwargs) -> SandboxPool: # Patch discovery + boot before construction: the constructor warms up, adopting any @@ -549,21 +583,146 @@ def fake_inspect(self, **kwargs) -> MagicMock: assert [host.job_id for host in pool._hosts] == ["host-sched"] # adopted, not booted +class TestHostAdmission: + """A Job carrying our pool's labels is a *claim* to be one of our hosts, not + proof: labels are set by whoever creates the Job, and the nonce that derives + the host token is a public label. These assert that everything else about the + Job -- all of it backend-asserted -- has to line up before a credential is + sent to it.""" + + def _reject(self, job, **kwargs) -> str | None: + defaults = dict( + policy=sandbox_mod.ADOPT_OWN, + principal_id=PRINCIPAL_ID, + namespace=None, + image="python:3.12", + flavor="cpu-basic", + ) + defaults.update(kwargs) + return sandbox_mod._host_rejection(job, **defaults) + + def test_a_genuine_host_is_accepted(self) -> None: + assert self._reject(_pool_host_job()) is None + + def test_a_job_started_by_someone_else_is_refused(self) -> None: + # The attack: a namespace member creates a Job with our pool's labels and + # the real host's nonce, and our client sends it the real host's token. + impostor = _pool_host_job("impostor") + impostor.initiator.id = "principal-attacker" + impostor.initiator.name = "attacker" + assert "different principal" in (self._reject(impostor) or "") + # `initiator` is the only field here the Jobs API does not let a client + # set, which is what makes this check worth anything. + + def test_an_unknown_initiator_is_refused_rather_than_assumed(self) -> None: + job = _pool_host_job() + job.initiator = None + assert self._reject(job) is not None + assert self._reject(_pool_host_job(), principal_id=None) is not None + + def test_namespace_policy_accepts_a_sibling_member_but_still_checks_the_spec(self) -> None: + other = _pool_host_job("other") + other.initiator.id = "principal-colleague" + assert self._reject(other, policy=sandbox_mod.ADOPT_NAMESPACE) is None + # Opting into sharing does not opt out of the rest. + other.docker_image = "evil:latest" + assert self._reject(other, policy=sandbox_mod.ADOPT_NAMESPACE) is not None + + def test_never_policy_refuses_everything(self) -> None: + assert self._reject(_pool_host_job(), policy=sandbox_mod.ADOPT_NEVER) is not None + + def test_a_mismatched_spec_is_refused(self) -> None: + wrong_image = _pool_host_job() + wrong_image.docker_image = "attacker/image:latest" + assert "image" in (self._reject(wrong_image) or "") + + wrong_flavor = _pool_host_job() + wrong_flavor.flavor = "a10g-large" + assert "flavor" in (self._reject(wrong_flavor) or "") + + # A Job labelled as a host but running something else is not a host. + wrong_command = _pool_host_job() + wrong_command.command = ["/bin/sh", "-c", "nc -l -p 49983"] + assert "bootstrap" in (self._reject(wrong_command) or "") + + wrong_namespace = _pool_host_job() + wrong_namespace.owner.name = "someone-else" + assert "namespace" in (self._reject(wrong_namespace, namespace="mine") or "") + + def test_registry_prefixes_do_not_cause_a_spurious_mismatch(self) -> None: + # Image names get normalized server-side; a mis-parse here must not stop a + # legitimate pool from working. + job = _pool_host_job() + job.docker_image = "docker.io/library/Python:3.12" + assert self._reject(job) is None + + def test_the_exposed_url_must_belong_to_this_job(self) -> None: + # Otherwise a Job could simply name where the credentials should go. + elsewhere = _pool_host_job("victim") + elsewhere.status.expose_urls = ["https://attacker--49983.hf.jobs"] + assert self._reject(elsewhere) is not None + + insecure = _pool_host_job("plain") + insecure.status.expose_urls = ["http://plain--49983.hf.jobs"] + assert "https" in (self._reject(insecure) or "") + + extra = _pool_host_job("extra") + extra.status.expose_urls = [ + "https://extra--49983.hf.jobs", + "https://extra--8080.hf.jobs", + ] + assert self._reject(extra) is not None + + none_exposed = _pool_host_job("bare") + none_exposed.status.expose_urls = None + assert self._reject(none_exposed) is not None + + def test_discovery_does_not_adopt_an_impostor(self, fake_server: str, monkeypatch) -> None: + # End to end through `create()`: the impostor is listed and matches on + # labels, and must not be adopted -- the pool boots its own host instead. + impostor = _pool_host_job("impostor") + impostor.initiator.id = "principal-attacker" + monkeypatch.setattr(sandbox_mod.HfApi, "list_jobs", _fake_list_jobs([impostor])) + monkeypatch.setattr( + sandbox_mod.SandboxPool, "_boot_host", lambda self: _make_server(fake_server, job_id="mine", capacity=4) + ) + monkeypatch.setattr( + sandbox_mod, "_connect_host", lambda api, jid, namespace=None: _make_server(fake_server, job_id=jid) + ) + pool = SandboxPool(image="python:3.12", flavor="cpu-basic", name="p1", token="hf_test") + + assert [host.job_id for host in pool._hosts] == ["mine"] + assert pool.create().host_id == "mine" + + def test_connect_refuses_a_pool_whose_only_host_is_an_impostor(self, monkeypatch) -> None: + impostor = _pool_host_job("impostor", pool_name="pool-x") + impostor.initiator.id = "principal-attacker" + monkeypatch.setattr(sandbox_mod.HfApi, "list_jobs", _fake_list_jobs([impostor])) + with pytest.raises(SandboxError, match="none usable"): + SandboxPool.connect("pool-x", token="hf_test") + # The error names the opt-in, so a legitimate shared-host user is not stuck. + with pytest.raises(SandboxError, match="none usable") as exc_info: + SandboxPool.connect("pool-x", token="hf_test") + assert "adopt_hosts='namespace'" in str(exc_info.value) + + def test_an_invalid_policy_is_rejected_at_construction(self) -> None: + with pytest.raises(ValueError, match="adopt_hosts"): + SandboxPool(name="p", adopt_hosts="anything", token="hf_test") + + class TestPoolConnect: """`SandboxPool.connect(pool_id)` rebuilds a pool from a running host's job spec + env vars — no local state, no config endpoint — then packs onto that host.""" def test_connect_reads_config_from_host_env(self, monkeypatch) -> None: url, _ = _spawn_fake(capacity=7) - job = MagicMock() - job.id = "hostA" - job.owner.name = "user" - job.status.stage = "RUNNING" - job.docker_image = "alpine:3.20" - job.space_id = None - job.flavor = "cpu-basic" + job = _pool_host_job( + "hostA", + pool_name="pool-x", + image="alpine:3.20", + env={"SBX_CAPACITY": "7", "SBX_IDLE_TIMEOUT": "600", "SBX_MAX_HOSTS": "3"}, + ) job.labels = {SANDBOX_LABEL: "1", MODE_LABEL: MODE_POOL, POOL_LABEL: "pool-x", NONCE_LABEL: "n"} - job.environment = {"SBX_CAPACITY": "7", "SBX_IDLE_TIMEOUT": "600", "SBX_MAX_HOSTS": "3"} monkeypatch.setattr(sandbox_mod.HfApi, "list_jobs", _fake_list_jobs([job])) monkeypatch.setattr( sandbox_mod, "_connect_host", lambda api, jid, namespace=None: _make_server(url, job_id=jid, capacity=7) @@ -688,12 +847,8 @@ def test_stale_host_falls_back_to_discovery_and_prunes(self, fake_server: str, m monkeypatch.setattr(sandbox_mod, "_derive_sandbox_token", lambda *a: "secret") pool = SandboxPool.connect("pool-stale", token="hf_test") - live_job = MagicMock() - live_job.id, live_job.flavor = "live", "cpu-basic" - live_job.owner.name, live_job.docker_image, live_job.space_id = "user", "python:3.12", None - live_job.status.stage = "RUNNING" + live_job = _pool_host_job("live", pool_name="pool-stale") live_job.labels = {SANDBOX_LABEL: "1", MODE_LABEL: MODE_POOL, POOL_LABEL: "pool-stale", NONCE_LABEL: "n"} - live_job.environment = {"SBX_CAPACITY": "4"} pool._api.list_jobs = MagicMock(return_value=[live_job]) monkeypatch.setattr( sandbox_mod, "_connect_host", lambda api, jid, namespace=None: _make_server(fake_server, job_id=jid)