[sandbox audit] Pin the sandbox server binary by digest and verify it before running it

src/huggingface_hub/_sandbox.py

@@ -65,6 +65,20 @@
 
 DEFAULT_IMAGE = "python:3.12"
 
+# The exact sbx-server build every sandbox job downloads and runs as root, pinned by digest.
+# The bucket object is fetched by its sha256 name and the download is verified before it is
+# made executable (see `_BOOTSTRAP_DOWNLOAD`), so the bytes that run are the bytes that were
+# reviewed for this release -- not whatever the mutable `sbx-server` alias points at today.
+# Both constants move together on every server release; the digest is printed by the server
+# repo's publish workflow.
+SANDBOX_SERVER_VERSION = "0.6.0"
+SANDBOX_SERVER_SHA256 = "501290eacb36a3bd8746b2f2ee20e85190ec3f697b4756c3e51a574384e61db4"
+
+# The sbx-server wire contract this client drives, checked against `/health`'s `protocol` on
+# startup. A pool host keeps the binary it downloaded at boot for up to 24h, so pinning a
+# digest does not stop this client from meeting a server it did not pin -- the check does.
+SANDBOX_SERVER_PROTOCOL = 2
+
 DEFAULT_IDLE_TIMEOUT = 10 * 60  # 10 minutes
 SANDBOX_MAX_LIFETIME = "24h"
 
@@ -101,14 +115,42 @@
 
 # Job startup script (needs only /bin/sh). The server bucket is public, so the download is
 # unauthenticated: no HF credential is ever placed in the job environment (see `_derive_sandbox_token`).
+#
+# It fetches a public object and runs it as root, as PID 1, holding the sandbox token -- so it
+# checks the bytes against `SANDBOX_SERVER_SHA256` and refuses to run them if they don't match.
+# The check happens *before* `chmod +x`: an unverified file that is already executable is one
+# slip away from being executed. Neither `sha256sum` nor `openssl` is guaranteed to exist in an
+# arbitrary image, so both are tried; with neither, the default is to refuse rather than to run
+# unverified code, and `SBX_ALLOW_UNVERIFIED_SERVER=1` (passed via `env=`) is the way out for an
+# image that has to. It never bypasses a *failed* check, only a missing tool.
 _BOOTSTRAP_DOWNLOAD = """\
 set -e
 d=/tmp/.sbx-server
 if command -v wget >/dev/null 2>&1; then wget -q -O "$d" "$SBX_SERVER_URL"
 elif command -v curl >/dev/null 2>&1; then curl -fsSL -o "$d" "$SBX_SERVER_URL"
-else cp "$SBX_SERVER_MOUNT/sbx-server" "$d"; fi
+else cp "$SBX_SERVER_MOUNT/sbx-server-$SBX_SERVER_SHA256" "$d"; fi
+if command -v sha256sum >/dev/null 2>&1; then
+  actual=$(sha256sum "$d" | cut -d' ' -f1)
+elif command -v openssl >/dev/null 2>&1; then
+  actual=$(openssl dgst -sha256 -r "$d" | cut -d' ' -f1)
+elif [ "${SBX_ALLOW_UNVERIFIED_SERVER:-}" = 1 ]; then
+  echo "sbx: SBX_ALLOW_UNVERIFIED_SERVER=1: running the sandbox server unverified" >&2
+  actual=$SBX_SERVER_SHA256  # opted out, so there is nothing left to compare against
+else
+  echo "sbx: cannot verify the sandbox server digest: this image has neither sha256sum nor" >&2
+  echo "sbx: openssl, so refusing to run it. To accept an unverified server, pass" >&2
+  echo "sbx: env={'SBX_ALLOW_UNVERIFIED_SERVER': '1'} when creating the sandbox." >&2
+  exit 1
+fi
+if [ "$actual" != "$SBX_SERVER_SHA256" ]; then
+  echo "sbx: sandbox server digest mismatch: got $actual," >&2
+  echo "sbx: expected $SBX_SERVER_SHA256." >&2
+  echo "sbx: refusing to run it. Upgrade huggingface_hub if this client is pinned to a" >&2
+  echo "sbx: digest that is no longer published." >&2
+  exit 1
+fi
 chmod +x "$d"
-unset SBX_SERVER_URL SBX_SERVER_MOUNT
+unset SBX_SERVER_URL SBX_SERVER_MOUNT SBX_SERVER_SHA256 SBX_ALLOW_UNVERIFIED_SERVER
 exec "$d"
 """
 
@@ -497,6 +539,32 @@ def _raise_for_status(response: httpx.Response) -> None:
     raise SandboxError(f"Sandbox API error ({response.status_code}): {message}", status_code=response.status_code)
 
 
+def _check_server_protocol(response: httpx.Response, job_id: str) -> None:
+    """Refuse a server whose wire contract this client cannot drive.
+
+    Permissive in one direction only. A *newer* server is accepted: it declares its own
+    backwards compatibility by keeping the protocol integer it still supports. An *older* one
+    is refused, because the mismatches are things like the per-sandbox capability token, which
+    surfaces as a 403 on an unrelated route several calls later instead of here.
+
+    `version` is no use for this: it moves for a doc fix as readily as for a protocol break.
+    """
+    try:
+        protocol = response.json().get("protocol")
+    except Exception:
+        protocol = None
+    if isinstance(protocol, int) and not isinstance(protocol, bool) and protocol >= SANDBOX_SERVER_PROTOCOL:
+        return
+    # No field at all means a server published before protocols were declared, i.e. 1.
+    reported = protocol if protocol is not None else "1 (declares no protocol)"
+    raise SandboxError(
+        f"Sandbox job {job_id} runs sbx-server protocol {reported}, but this version of "
+        f"huggingface_hub needs protocol {SANDBOX_SERVER_PROTOCOL} or newer. A pool host keeps the "
+        "server binary it downloaded at boot for up to 24h, so recycle this pool's hosts (kill them, "
+        "or let them idle-time-out) to pick up the current server."
+    )
+
+
 class _SandboxAuth(httpx.Auth):
     """Attach the *current* HF bearer to every request.
 
@@ -665,6 +733,7 @@ def wait_ready(self, start_timeout: float) -> None:
             try:
                 response = self._client.get(self.base_url + "/health", timeout=httpx.Timeout(5.0))
                 if response.status_code == 200:
+                    _check_server_protocol(response, self.job_id)
                     return
             except httpx.RequestError:
                 pass
@@ -2352,7 +2421,8 @@ def _bootstrap_job_spec(
 
     Shared by dedicated sandboxes and shared hosts: both fetch and exec the same unified
     `sbx-server` binary at startup (via `/bin/sh`), downloading it with wget/curl, or
-    reading it off the always-mounted server bucket when the image ships neither.
+    reading it off the always-mounted server bucket when the image ships neither. Either
+    way the download is pinned to `SANDBOX_SERVER_SHA256` and verified before it runs.
     """
     # Reserved SBX_* keys go last so user-provided env/secrets can't override them
     # (e.g. clobbering SBX_PORT would break the proxy, SBX_TOKEN would break auth).
@@ -2364,7 +2434,12 @@ def _bootstrap_job_spec(
     if forward_hf_token:
         job_secrets["HF_TOKEN"] = hf_token
 
-    job_env["SBX_SERVER_URL"] = f"{api.endpoint}/buckets/{constants.SANDBOX_SERVER_BUCKET}/resolve/sbx-server"
+    # The digest-named object, not the mutable `sbx-server` alias: a fetch that names the
+    # content it wants cannot be answered with different content, and the bootstrap re-checks
+    # the digest anyway so a rewritten object is caught rather than executed.
+    server_object = f"sbx-server-{SANDBOX_SERVER_SHA256}"
+    job_env["SBX_SERVER_URL"] = f"{api.endpoint}/buckets/{constants.SANDBOX_SERVER_BUCKET}/resolve/{server_object}"
+    job_env["SBX_SERVER_SHA256"] = SANDBOX_SERVER_SHA256
     # Always mount the server bucket as a transparent fallback for images without wget/curl.
     # It's only read (paying the ~2-3s FUSE cost) when the bootstrap script can't download.
     job_env["SBX_SERVER_MOUNT"] = _SERVER_MOUNT_PATH

tests/test_sandbox.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 import stat
 import threading
 import time
@@ -112,6 +113,8 @@ class _FakeServer(BaseHTTPRequestHandler):
     requests: list = []
     writes: list = []  # (path, body) of every /files/write received
     redirect_to = ""  # where the /v1/redirect route points (set per-subclass)
+    # Extra /health fields, so a test can stand in for a server of another vintage.
+    health: dict = {"protocol": sandbox_mod.SANDBOX_SERVER_PROTOCOL}
 
     def log_message(self, *args) -> None:
         pass
@@ -224,8 +227,11 @@ def do_DELETE(self) -> None:
 
     def do_GET(self) -> None:
         self._record()
-        self._expect_token()
         cls = type(self)
+        if self.path == "/health":  # unauthenticated on the real server, so no token check
+            self._json({"status": "ok", "version": "0.6.0", **cls.health})
+            return
+        self._expect_token()
         if self.path == "/v1/redirect":  # a redirect the client must not follow
             self.send_response(302)
             self.send_header("Location", cls.redirect_to)
@@ -270,6 +276,7 @@ def fake_server():
     _FakeServer.requests = []
     _FakeServer.writes = []
     _FakeServer.redirect_to = ""
+    _FakeServer.health = {"protocol": sandbox_mod.SANDBOX_SERVER_PROTOCOL}
     server = ThreadingHTTPServer(("127.0.0.1", 0), _FakeServer)
     thread = threading.Thread(target=server.serve_forever, daemon=True)
     thread.start()
@@ -506,6 +513,92 @@ def test_a_process_without_an_id_refuses_to_be_killed(self, fake_server: str) ->
             process.kill()
 
 
+class TestServerBinaryPinning:
+    """The server binary is fetched from a public bucket and run as root, as PID 1, holding
+    the sandbox token. These assert the client says which bytes that may be, and that the job
+    refuses anything else rather than executing it and finding out."""
+
+    def _job_env(self) -> dict:
+        api = MagicMock()
+        api.endpoint = "https://huggingface.co"
+        _, job_env, _, _ = sandbox_mod._bootstrap_job_spec(
+            api,
+            "hf_test",
+            env=None,
+            secrets=None,
+            volumes=None,
+            idle_timeout=None,
+            forward_hf_token=False,
+            sandbox_token="secret",
+        )
+        return job_env
+
+    def test_the_pinned_digest_is_a_sha256(self) -> None:
+        assert re.fullmatch(r"[0-9a-f]{64}", sandbox_mod.SANDBOX_SERVER_SHA256)
+
+    def test_the_download_names_the_digest_not_the_mutable_alias(self) -> None:
+        # `/resolve/sbx-server` is overwritten by every publish, so a fetch of it cannot be
+        # verified in advance -- a fetch that names its content can.
+        job_env = self._job_env()
+        assert job_env["SBX_SERVER_URL"].endswith(f"/resolve/sbx-server-{sandbox_mod.SANDBOX_SERVER_SHA256}")
+        assert job_env["SBX_SERVER_SHA256"] == sandbox_mod.SANDBOX_SERVER_SHA256
+
+    def test_the_digest_is_checked_before_the_binary_becomes_executable(self) -> None:
+        script = sandbox_mod._BOOTSTRAP_DOWNLOAD
+        assert script.index("$SBX_SERVER_SHA256") < script.index("chmod +x")
+        assert script.index("exit 1") < script.index("chmod +x")
+        # `sha256sum` is not guaranteed in an arbitrary image, so openssl is the second try.
+        assert "sha256sum" in script
+        assert "openssl dgst -sha256" in script
+        # The bucket-mount fallback is content-addressed too; it used to read the alias.
+        assert 'cp "$SBX_SERVER_MOUNT/sbx-server-$SBX_SERVER_SHA256"' in script
+
+    def test_an_image_with_no_hash_tool_has_to_opt_in_by_name(self) -> None:
+        # Failing closed is the whole point of the pin, so the only way past it is an
+        # explicit env var -- and it can only skip a *missing tool*, never a failed compare.
+        script = sandbox_mod._BOOTSTRAP_DOWNLOAD
+        assert script.count("SBX_ALLOW_UNVERIFIED_SERVER") > 0
+        assert script.index("SBX_ALLOW_UNVERIFIED_SERVER") > script.index("openssl dgst -sha256")
+
+
+class TestProtocolNegotiation:
+    """A pool host keeps the binary it downloaded at boot for up to 24h, so pinning a digest
+    does not stop this client from meeting a server it did not pin. `/health` reports the wire
+    contract; the client refuses the ones it cannot drive, instead of failing later on an
+    unrelated route."""
+
+    def test_a_matching_protocol_is_accepted(self, fake_server: str) -> None:
+        _make_server(fake_server).wait_ready(5.0)
+
+    def test_a_newer_server_may_serve_this_client(self, fake_server: str) -> None:
+        # Permissive in this direction only: a server that still reports a protocol this
+        # client knows is declaring that it kept serving it.
+        _FakeServer.health = {"protocol": sandbox_mod.SANDBOX_SERVER_PROTOCOL + 1}
+        _make_server(fake_server).wait_ready(5.0)
+
+    def test_an_older_server_is_refused_and_the_error_names_the_fix(self, fake_server: str) -> None:
+        _FakeServer.health = {"protocol": sandbox_mod.SANDBOX_SERVER_PROTOCOL - 1}
+        with pytest.raises(SandboxError) as exc_info:
+            _make_server(fake_server).wait_ready(5.0)
+        message = str(exc_info.value)
+        assert f"protocol {sandbox_mod.SANDBOX_SERVER_PROTOCOL - 1}" in message
+        assert "recycle" in message  # the actionable part: the host has to be replaced
+
+    def test_a_server_that_declares_nothing_is_refused(self, fake_server: str) -> None:
+        # Every server published before the field existed. Silence is not compatibility.
+        _FakeServer.health = {}
+        with pytest.raises(SandboxError, match="declares no protocol"):
+            _make_server(fake_server).wait_ready(5.0)
+
+    def test_a_non_integer_protocol_is_refused(self, fake_server: str) -> None:
+        # Anything that is not a number this client can compare is a server it cannot
+        # reason about, so it is refused rather than assumed compatible.
+        for value in ("2", 2.5, True, None, {"major": 2}):
+            _FakeServer.health = {"protocol": value}
+            with pytest.raises(SandboxError, match="protocol"):
+                _make_server(fake_server).wait_ready(5.0)
+
+
 class TestSharedSandbox:
     """A shared sandbox routes operations under /v1/sandboxes/<local_id>/ and is
     terminated with a DELETE on the host (the host job keeps running)."""