Protect sandbox control-plane routes (#142)

* Protect sandbox control-plane routes

The sandbox server exposed command and file APIs on public Spaces without validating the client. This adds bearer-token validation to every non-health endpoint and provisions a per-sandbox API secret so control-plane access is not coupled to URL knowledge.

Constraint: Existing sandboxes may only have HF_TOKEN configured, so the server accepts SANDBOX_API_TOKEN first and falls back to HF_TOKEN for compatibility.

Rejected: Make every sandbox private by default | that changes product behavior and still leaves the embedded API without a server-side auth boundary.

Confidence: high

Scope-risk: narrow

Directive: Do not add new sandbox API routes without the auth dependency unless the route is intentionally public health/status only.

Tested: UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/unit/test_sandbox_api_auth.py

* Preserve sandbox reconnect with API token

Review caught that Sandbox.connect would send only the HF token while newly-created sandboxes now expect SANDBOX_API_TOKEN. Add an api_token parameter, hide the field from repr, and cover the documented HF_TOKEN fallback for legacy sandboxes.

Constraint: Existing connect callers must keep working when they only pass an HF token to legacy sandboxes.

Confidence: high

Scope-risk: narrow

Tested: UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/unit/test_sandbox_api_auth.py

* Add live sandbox communication test

Add an opt-in integration test that provisions a real public cpu-basic sandbox, verifies unauthenticated API calls are rejected, exercises authenticated bash/write/exists/read communication, verifies reconnect with the sandbox API token, and deletes the Space afterward.

Constraint: Live sandbox creation requires HF_TOKEN and should not run in default CI because it creates a real Space.

Rejected: Use a private Space for this test | private hf.space requests return 404 before the sandbox FastAPI auth layer is reachable, so it does not validate the public-sandbox threat model fixed here.

Confidence: high

Scope-risk: narrow

Tested: ML_INTERN_LIVE_SANDBOX_TESTS=1 ML_INTERN_LIVE_ENV_FILE=/Users/akseljoonas/Documents/ml-intern/.env UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/integration/test_live_sandbox_auth.py -q -s

Tested: UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/unit/test_sandbox_api_auth.py tests/integration/test_live_sandbox_auth.py -q

agent/tools/sandbox_client.py

@@ -37,6 +37,7 @@ Tools: bash, read, write, edit, upload
 from __future__ import annotations
 
 import io
+import secrets as secrets_lib
 import sys
 import time
 import uuid
@@ -99,8 +100,8 @@ CMD ["python", "sandbox_server.py"]
 
 _SANDBOX_SERVER = '''\
 """Minimal FastAPI server for sandbox operations."""
-import os, subprocess, pathlib, signal, threading, re, tempfile
-from fastapi import FastAPI
+import hmac, os, subprocess, pathlib, signal, threading, re, tempfile
+from fastapi import Depends, FastAPI, HTTPException, Request
 from pydantic import BaseModel
 from typing import Optional
 import uvicorn
@@ -156,6 +157,22 @@ def _atomic_write(path: pathlib.Path, content: str):
 
 app = FastAPI()
 
+def _expected_api_token() -> str:
+    return os.environ.get("SANDBOX_API_TOKEN") or os.environ.get("HF_TOKEN") or ""
+
+def _require_auth(request: Request) -> None:
+    expected = _expected_api_token()
+    if not expected:
+        raise HTTPException(status_code=503, detail="Sandbox API token not configured")
+    auth_header = request.headers.get("authorization", "")
+    scheme, _, supplied = auth_header.partition(" ")
+    if scheme.lower() != "bearer" or not supplied:
+        raise HTTPException(status_code=401, detail="Missing bearer token")
+    if not hmac.compare_digest(supplied, expected):
+        raise HTTPException(status_code=401, detail="Invalid bearer token")
+
+_AUTH = [Depends(_require_auth)]
+
 # Track active bash processes so they can be killed on cancel
 _active_procs = {}  # pid -> subprocess.Popen
 _proc_lock = threading.Lock()
@@ -344,7 +361,7 @@ def _validate_python(content, path=""):
 def health():
     return {"status": "ok"}
 
-@app.post("/api/bash")
+@app.post("/api/bash", dependencies=_AUTH)
 def bash(req: BashReq):
     try:
         proc = subprocess.Popen(
@@ -371,7 +388,7 @@ def bash(req: BashReq):
     except Exception as e:
         return {"success": False, "output": "", "error": str(e)}
 
-@app.post("/api/kill")
+@app.post("/api/kill", dependencies=_AUTH)
 def kill_all():
     """Kill all active bash processes. Called when user cancels."""
     with _proc_lock:
@@ -389,7 +406,7 @@ def kill_all():
                 pass
     return {"success": True, "output": f"Killed {len(killed)} process(es): {killed}", "error": ""}
 
-@app.post("/api/read")
+@app.post("/api/read", dependencies=_AUTH)
 def read(req: ReadReq):
     try:
         p = pathlib.Path(req.path)
@@ -406,7 +423,7 @@ def read(req: ReadReq):
     except Exception as e:
         return {"success": False, "output": "", "error": str(e)}
 
-@app.post("/api/write")
+@app.post("/api/write", dependencies=_AUTH)
 def write(req: WriteReq):
     try:
         p = pathlib.Path(req.path)
@@ -420,7 +437,7 @@ def write(req: WriteReq):
     except Exception as e:
         return {"success": False, "output": "", "error": str(e)}
 
-@app.post("/api/edit")
+@app.post("/api/edit", dependencies=_AUTH)
 def edit(req: EditReq):
     try:
         p = pathlib.Path(req.path)
@@ -447,7 +464,7 @@ def edit(req: EditReq):
     except Exception as e:
         return {"success": False, "output": "", "error": str(e)}
 
-@app.post("/api/exists")
+@app.post("/api/exists", dependencies=_AUTH)
 def exists(req: ExistsReq):
     return {"success": True, "output": str(pathlib.Path(req.path).exists()).lower(), "error": ""}
 
@@ -482,6 +499,7 @@ class Sandbox:
 
     space_id: str
     token: str | None = None
+    api_token: str | None = field(default=None, repr=False)
     work_dir: str = "/app"
     timeout: int = DEFAULT_TIMEOUT
     _owns_space: bool = field(default=False, repr=False)
@@ -495,9 +513,10 @@ class Sandbox:
         # Trailing slash is critical: httpx resolves relative paths against base_url.
         # Without it, client.get("health") resolves to /health instead of /api/health.
         self._base_url = f"https://{slug}.hf.space/api/"
+        api_token = self.api_token or self.token
         self._client = httpx.Client(
             base_url=self._base_url,
-            headers={"Authorization": f"Bearer {self.token}"} if self.token else {},
+            headers={"Authorization": f"Bearer {api_token}"} if api_token else {},
             timeout=httpx.Timeout(MAX_TIMEOUT, connect=30),
             follow_redirects=True,
         )
@@ -563,6 +582,7 @@ class Sandbox:
         base = name or "sandbox"
         suffix = uuid.uuid4().hex[:8]
         space_id = f"{owner}/{base}-{suffix}"
+        sandbox_api_token = secrets_lib.token_urlsafe(32)
 
         _log(f"Creating sandbox: {space_id} (from {template})...")
 
@@ -583,8 +603,9 @@ class Sandbox:
         # Inject secrets BEFORE uploading server files (which triggers rebuild).
         # Secrets added after a Space is running aren't available until restart,
         # so they must be set before the build/start cycle.
-        if secrets:
-            for key, val in secrets.items():
+        sandbox_secrets = {**(secrets or {}), "SANDBOX_API_TOKEN": sandbox_api_token}
+        if sandbox_secrets:
+            for key, val in sandbox_secrets.items():
                 api.add_space_secret(space_id, key, val)
 
         # Upload sandbox server and Dockerfile (triggers rebuild)
@@ -617,7 +638,12 @@ class Sandbox:
         _check_cancel()
 
         # Wait for the API server to be responsive (non-fatal)
-        sb = cls(space_id=space_id, token=token, _owns_space=True)
+        sb = cls(
+            space_id=space_id,
+            token=token,
+            api_token=sandbox_api_token,
+            _owns_space=True,
+        )
         try:
             sb._wait_for_api(timeout=API_WAIT_TIMEOUT, log=_log)
         except TimeoutError as e:
@@ -648,13 +674,24 @@ class Sandbox:
         log("Server files uploaded, rebuild triggered.")
 
     @classmethod
-    def connect(cls, space_id: str, *, token: str | None = None) -> Sandbox:
+    def connect(
+        cls,
+        space_id: str,
+        *,
+        token: str | None = None,
+        api_token: str | None = None,
+    ) -> Sandbox:
         """
         Connect to an existing running Space.
 
         Does a health check to verify the Space is reachable.
         """
-        sb = cls(space_id=space_id, token=token, _owns_space=False)
+        sb = cls(
+            space_id=space_id,
+            token=token,
+            api_token=api_token,
+            _owns_space=False,
+        )
         sb._wait_for_api(timeout=60)
         return sb
 

tests/integration/test_live_sandbox_auth.py

@@ -0,0 +1,90 @@
+"""Opt-in live sandbox communication test.
+
+This test creates a real Hugging Face Space sandbox, verifies that unauthenticated
+requests are rejected, then exercises the authenticated agent client end-to-end.
+It is skipped unless ``ML_INTERN_LIVE_SANDBOX_TESTS=1`` and ``HF_TOKEN`` are set.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import httpx
+import pytest
+from dotenv import load_dotenv
+from huggingface_hub import HfApi
+
+from agent.tools.sandbox_client import Sandbox
+
+
+if env_file := os.environ.get("ML_INTERN_LIVE_ENV_FILE"):
+    load_dotenv(Path(env_file))
+
+
+def _skip_without_live_sandbox() -> None:
+    if os.environ.get("ML_INTERN_LIVE_SANDBOX_TESTS") != "1":
+        pytest.skip("set ML_INTERN_LIVE_SANDBOX_TESTS=1 to create a real sandbox")
+    if not os.environ.get("HF_TOKEN"):
+        pytest.skip("set HF_TOKEN to create a real sandbox")
+
+
+def test_live_sandbox_authenticated_agent_communication():
+    _skip_without_live_sandbox()
+
+    token = os.environ["HF_TOKEN"]
+    owner = HfApi(token=token).whoami()["name"]
+    sandbox = None
+
+    try:
+        sandbox = Sandbox.create(
+            owner=owner,
+            name="ml-intern-live-auth",
+            hardware="cpu-basic",
+            private=False,
+            token=token,
+            secrets={"HF_TOKEN": token},
+            wait_timeout=900,
+        )
+
+        unauthenticated = httpx.Client(
+            base_url=sandbox._base_url,
+            timeout=30,
+            follow_redirects=True,
+        )
+        try:
+            denied = unauthenticated.post("exists", json={"path": "/tmp"})
+            assert denied.status_code == 401
+        finally:
+            unauthenticated.close()
+
+        bash = sandbox.bash("printf sandbox-live-ok", timeout=30)
+        assert bash.success, bash.error
+        assert "sandbox-live-ok" in bash.output
+
+        write = sandbox.write("/tmp/ml_intern_live_auth.txt", "alpha\nbeta\n")
+        assert write.success, write.error
+
+        exists = sandbox._call("exists", {"path": "/tmp/ml_intern_live_auth.txt"})
+        assert exists.success, exists.error
+        assert exists.output == "true"
+
+        read = sandbox.read("/tmp/ml_intern_live_auth.txt")
+        assert read.success, read.error
+        assert "alpha" in read.output
+        assert "beta" in read.output
+
+        reattached = Sandbox.connect(
+            sandbox.space_id,
+            token=token,
+            api_token=sandbox.api_token,
+        )
+        try:
+            reread = reattached.read("/tmp/ml_intern_live_auth.txt")
+            assert reread.success, reread.error
+            assert "alpha" in reread.output
+        finally:
+            reattached._client.close()
+    finally:
+        if sandbox is not None:
+            sandbox.delete()

tests/unit/test_sandbox_api_auth.py

@@ -0,0 +1,87 @@
+from fastapi.testclient import TestClient
+
+from agent.tools.sandbox_client import _SANDBOX_SERVER, Sandbox
+
+
+def _sandbox_app(
+    monkeypatch,
+    token: str | None = "sandbox-secret",
+    *,
+    hf_token: str | None = None,
+):
+    monkeypatch.delenv("SANDBOX_API_TOKEN", raising=False)
+    monkeypatch.delenv("HF_TOKEN", raising=False)
+    if token is not None:
+        monkeypatch.setenv("SANDBOX_API_TOKEN", token)
+    if hf_token is not None:
+        monkeypatch.setenv("HF_TOKEN", hf_token)
+    namespace = {}
+    exec(_SANDBOX_SERVER, namespace)
+    return namespace["app"]
+
+
+def test_health_is_public(monkeypatch):
+    client = TestClient(_sandbox_app(monkeypatch))
+
+    response = client.get("/api/health")
+
+    assert response.status_code == 200
+    assert response.json() == {"status": "ok"}
+
+
+def test_file_and_command_routes_require_bearer_token(monkeypatch):
+    client = TestClient(_sandbox_app(monkeypatch, "sandbox-secret"))
+
+    response = client.post("/api/exists", json={"path": "/tmp"})
+
+    assert response.status_code == 401
+
+
+def test_file_and_command_routes_accept_valid_bearer_token(monkeypatch):
+    client = TestClient(_sandbox_app(monkeypatch, "sandbox-secret"))
+
+    response = client.post(
+        "/api/exists",
+        json={"path": "/tmp"},
+        headers={"Authorization": "Bearer sandbox-secret"},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["success"] is True
+
+
+def test_legacy_hf_token_fallback_is_accepted(monkeypatch):
+    client = TestClient(_sandbox_app(monkeypatch, token=None, hf_token="hf-secret"))
+
+    response = client.post(
+        "/api/exists",
+        json={"path": "/tmp"},
+        headers={"Authorization": "Bearer hf-secret"},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["success"] is True
+
+
+def test_protected_routes_fail_closed_without_configured_token(monkeypatch):
+    client = TestClient(_sandbox_app(monkeypatch, None))
+
+    response = client.post(
+        "/api/exists",
+        json={"path": "/tmp"},
+        headers={"Authorization": "Bearer anything"},
+    )
+
+    assert response.status_code == 503
+
+
+def test_sandbox_prefers_control_plane_token_for_api_headers():
+    sandbox = Sandbox("owner/name", token="hf-token", api_token="sandbox-secret")
+
+    assert sandbox._client.headers["authorization"] == "Bearer sandbox-secret"
+
+
+def test_sandbox_api_token_is_hidden_from_repr():
+    sandbox = Sandbox("owner/name", token="hf-token", api_token="sandbox-secret")
+
+    assert "sandbox-secret" not in repr(sandbox)