fix(test): align snapshot disk gate with runtime contracts

.github/workflows/ci.yml

@@ -107,6 +107,12 @@ jobs:
         with:
           workspaces: ${{ matrix.workspace }}
 
+      - name: Test snapshot-disk harness contract
+        if: matrix.workspace == 'orch'
+        run: |
+          python3 -m unittest orch/tests/test_e2e_snapshot_disk.py -v
+          bash -n orch/tests/e2e_snapshot_disk.sh
+
       - name: Format
         working-directory: ${{ matrix.workspace }}
         run: cargo fmt --all -- --check

orch/tests/e2e_snapshot_disk.sh

@@ -54,6 +54,40 @@ json_field() {
   python3 -c 'import json,sys; print(json.load(sys.stdin)[sys.argv[1]])' "$1"
 }
 
+snapshot_artifact_path() {
+  # Resolve the public handle through this test's local store for file checks.
+  python3 - "$DIR/fleet.db" "$1" <<'PY'
+from pathlib import Path
+import sqlite3
+import sys
+
+with sqlite3.connect(Path(sys.argv[1]).resolve().as_uri() + "?mode=ro", uri=True) as db:
+    row = db.execute(
+        "SELECT path FROM snapshots WHERE snapshot_id = ?", (sys.argv[2],)
+    ).fetchone()
+if row is None:
+    raise SystemExit("FAIL: snapshot is missing from the local test database")
+print(row[0])
+PY
+}
+
+vm_overlay_path() {
+  # Restore uppers use snapshot-specific names; the runtime record is canonical.
+  python3 - "$DIR/fleet.db" "$1" <<'PY'
+from pathlib import Path
+import sqlite3
+import sys
+
+with sqlite3.connect(Path(sys.argv[1]).resolve().as_uri() + "?mode=ro", uri=True) as db:
+    row = db.execute(
+        "SELECT runtime_overlay_path FROM vms WHERE id = ?", (sys.argv[2],)
+    ).fetchone()
+if row is None or not row[0]:
+    raise SystemExit("FAIL: VM has no recorded runtime overlay")
+print(row[0])
+PY
+}
+
 exec_json() {
   local vm_id=$1 command=$2
   api -H 'Content-Type: application/json' \
@@ -109,6 +143,8 @@ echo "== create source and checkpoint disk =="
 SOURCE_JSON=$(api -H 'Content-Type: application/json' \
   -d '{"vcpus":1,"memory_mib":256}' "$BASE_URL/v1/vms")
 SOURCE_ID=$(printf '%s' "$SOURCE_JSON" | json_field id)
+SOURCE_OVERLAY=$(vm_overlay_path "$SOURCE_ID")
+[ -f "$SOURCE_OVERLAY" ]
 exec_json "$SOURCE_ID" "sh -c 'echo snapshot-checkpoint > /root/tarit-snapshot-state; sync'" |
   grep -q '"exit_code":0'
 
@@ -118,40 +154,52 @@ SNAPSHOT_JSON=$(api -H 'Content-Type: application/json' -d '{"diff":false}' \
   "$BASE_URL/v1/vms/$SOURCE_ID/snapshot")
 SNAPSHOT_END_MS=$(python3 -c 'import time; print(time.monotonic_ns() // 1000000)')
 SNAPSHOT_CAPTURE_MS=$((SNAPSHOT_END_MS - SNAPSHOT_START_MS))
-SNAPSHOT_PATH=$(printf '%s' "$SNAPSHOT_JSON" | json_field path)
+SNAPSHOT_ID=$(printf '%s' "$SNAPSHOT_JSON" | json_field snapshot_id)
+SNAPSHOT_PATH=$(snapshot_artifact_path "$SNAPSHOT_ID")
 [ -f "$SNAPSHOT_PATH" ]
 exec_json "$SOURCE_ID" "sh -c 'echo post-snapshot-mutation > /root/tarit-snapshot-state; sync'" |
   grep -q '"exit_code":0'
 api -X DELETE "$BASE_URL/v1/vms/$SOURCE_ID" >/dev/null
-[ ! -e "$DIR/sockets/overlays/$SOURCE_ID.cow" ] || {
+[ -f "$SNAPSHOT_PATH" ]
+[ ! -e "$SOURCE_OVERLAY" ] || {
   echo "FAIL: source VM overlay survived deletion"
   exit 1
 }
 
 echo "== restore twice from the snapshot-owned disk artifact =="
 RESTORE_A=$(api -H 'Content-Type: application/json' \
-  -d "$(python3 -c 'import json,sys; print(json.dumps({"snapshot_path":sys.argv[1]}))' "$SNAPSHOT_PATH")" \
+  -d "$(python3 -c 'import json,sys; print(json.dumps({"snapshot_id":sys.argv[1]}))' "$SNAPSHOT_ID")" \
   "$BASE_URL/v1/restore")
 RESTORE_B=$(api -H 'Content-Type: application/json' \
-  -d "$(python3 -c 'import json,sys; print(json.dumps({"snapshot_path":sys.argv[1]}))' "$SNAPSHOT_PATH")" \
+  -d "$(python3 -c 'import json,sys; print(json.dumps({"snapshot_id":sys.argv[1]}))' "$SNAPSHOT_ID")" \
   "$BASE_URL/v1/restore")
 A_ID=$(printf '%s' "$RESTORE_A" | json_field id)
 B_ID=$(printf '%s' "$RESTORE_B" | json_field id)
 [ "$A_ID" != "$B_ID" ]
-[ -f "$DIR/sockets/overlays/$A_ID.cow" ]
-[ -f "$DIR/sockets/overlays/$B_ID.cow" ]
+A_OVERLAY=$(vm_overlay_path "$A_ID")
+B_OVERLAY=$(vm_overlay_path "$B_ID")
+for overlay in "$A_OVERLAY" "$B_OVERLAY"; do
+  [ -f "$overlay" ] || {
+    echo "FAIL: restored VM overlay is missing"
+    exit 1
+  }
+done
+[ "$A_OVERLAY" != "$B_OVERLAY" ] || {
+  echo "FAIL: restored VMs shared a writable overlay"
+  exit 1
+}
 
 exec_json "$A_ID" 'cat /root/tarit-snapshot-state' | grep -q 'snapshot-checkpoint'
 exec_json "$B_ID" 'cat /root/tarit-snapshot-state' | grep -q 'snapshot-checkpoint'
 exec_json "$A_ID" "sh -c 'echo restore-a-private > /root/tarit-snapshot-state; sync'" |
   grep -q '"exit_code":0'
 exec_json "$A_ID" 'cat /root/tarit-snapshot-state' | grep -q 'restore-a-private'
 B_STATE=$(exec_json "$B_ID" 'cat /root/tarit-snapshot-state')
-printf '%s' "$B_STATE" | grep -q 'snapshot-checkpoint'
 if printf '%s' "$B_STATE" | grep -q 'restore-a-private'; then
   echo "FAIL: restored VMs shared writable disk state"
   exit 1
 fi
+printf '%s' "$B_STATE" | grep -q 'snapshot-checkpoint'
 
 api -X DELETE "$BASE_URL/v1/vms/$A_ID" >/dev/null
 api -X DELETE "$BASE_URL/v1/vms/$B_ID" >/dev/null

orch/tests/test_e2e_snapshot_disk.py

@@ -0,0 +1,251 @@
+"""Exercise the real snapshot-disk shell harness without a VMM or KVM.
+
+The fake API follows SnapshotResponse/RestoreRequest's opaque-handle contract.
+It models disk markers only to check the harness's requests and assertions;
+actual snapshot/restore isolation still requires the live KVM gate.
+"""
+
+import json
+import os
+import re
+import shlex
+import shutil
+import signal
+import sqlite3
+import subprocess
+import sys
+import tempfile
+import unittest
+import uuid
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+
+HARNESS = Path(__file__).with_name("e2e_snapshot_disk.sh")
+SNAPSHOT_ID = "dd3efb1c-883a-48c6-9065-995580a6b925"
+
+
+class FakeApi(BaseHTTPRequestHandler):
+    def log_message(self, *_args):
+        pass
+
+    def reply(self, data, status=200):
+        body = json.dumps(data, separators=(",", ":")).encode()
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self):
+        if self.path == "/health":
+            self.reply({"status": "ok"})
+        else:
+            self.reply({"error": "unexpected endpoint"}, 404)
+
+    def record(self, body=None):
+        with self.server.requests.open("a") as log:
+            log.write(
+                json.dumps({"method": self.command, "path": self.path, "body": body})
+                + "\n"
+            )
+
+    def create_vm(self, state, *, restored=False):
+        vm_id = str(uuid.UUID(int=self.server.next_vm))
+        self.server.next_vm += 1
+        self.server.vms[vm_id] = state
+        # Runtime metadata owns the filename; restored uppers need not use the
+        # cold-start VM filename or any naming convention known to the harness.
+        filename = f"restored-private-{vm_id}.cow" if restored else f"{vm_id}.cow"
+        overlay = self.server.overlays / filename
+        self.server.overlay_paths[vm_id] = overlay
+        if not (restored and self.server.mode == "missing-restore-overlay"):
+            overlay.touch()
+        if not (restored and self.server.mode == "missing-runtime-record"):
+            self.server.db.execute(
+                "INSERT INTO vms(id, runtime_overlay_path) VALUES (?, ?)",
+                (vm_id, str(overlay)),
+            )
+            self.server.db.commit()
+        self.reply({"id": vm_id}, 201)
+
+    def do_POST(self):
+        body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+        self.record(body)
+        if self.path == "/v1/vms":
+            self.create_vm({"marker": ""})
+        elif self.path.endswith("/snapshot"):
+            source_id = self.path.split("/")[-2]
+            self.server.snapshot = self.server.vms[source_id].copy()
+            artifact = self.server.overlays.parent / "snapshot.ram"
+            if self.server.mode != "missing-artifact":
+                artifact.write_bytes(b"fake snapshot")
+            if self.server.mode != "missing-record":
+                self.server.db.execute(
+                    "INSERT INTO snapshots(path, snapshot_id) VALUES (?, ?)",
+                    (str(artifact), SNAPSHOT_ID),
+                )
+                self.server.db.commit()
+            # Public responses must not expose a host path.
+            self.reply({"snapshot_id": SNAPSHOT_ID})
+        elif self.path == "/v1/restore":
+            if body != {"snapshot_id": SNAPSHOT_ID}:
+                self.reply(
+                    {"error": "restore requires only the opaque snapshot_id"}, 400
+                )
+                return
+            state = self.server.snapshot
+            if self.server.mode != "shared-restores":
+                state = state.copy()
+            self.create_vm(state, restored=True)
+        elif self.path == "/v1/execute":
+            state = self.server.vms[body["vm_id"]]
+            marker = re.search(r"echo ([a-z-]+) >", body["command"])
+            if marker:
+                state["marker"] = marker.group(1)
+                stdout = ""
+            elif body["command"] == "cat /root/tarit-snapshot-state":
+                stdout = state["marker"] + "\n"
+            else:
+                self.reply({"error": "unexpected guest command"}, 400)
+                return
+            self.reply({"exit_code": 0, "stdout": stdout, "stderr": ""})
+        else:
+            self.reply({"error": "unexpected endpoint"}, 404)
+
+    def do_DELETE(self):
+        self.record()
+        vm_id = self.path.split("/")[-1]
+        del self.server.vms[vm_id]
+        self.server.overlay_paths.pop(vm_id).unlink(missing_ok=True)
+        self.server.db.execute("DELETE FROM vms WHERE id = ?", (vm_id,))
+        self.server.db.commit()
+        self.reply({})
+
+
+def fake_taritd():
+    host, port = os.environ["TARIT_LISTEN"].rsplit(":", 1)
+    Path(os.environ["FAKE_TARITD_PID"]).write_text(str(os.getpid()))
+    with HTTPServer((host, int(port)), FakeApi) as server:
+        server.mode = os.environ["FAKE_SNAPSHOT_MODE"]
+        server.requests = Path(os.environ["FAKE_SNAPSHOT_REQUESTS"])
+        server.overlays = Path(os.environ["TARIT_SOCKET_DIR"]) / "overlays"
+        server.overlays.mkdir()
+        server.db = sqlite3.connect(os.environ["TARIT_DB"])
+        server.db.execute(
+            "CREATE TABLE snapshots(path TEXT PRIMARY KEY, snapshot_id TEXT UNIQUE)"
+        )
+        server.db.execute(
+            "CREATE TABLE vms(id TEXT PRIMARY KEY, runtime_overlay_path TEXT)"
+        )
+        server.db.commit()
+        server.vms = {}
+        server.overlay_paths = {}
+        server.next_vm = 1
+        server.serve_forever()
+
+
+class SnapshotDiskContractTests(unittest.TestCase):
+    def run_harness(self, mode="normal"):
+        with tempfile.TemporaryDirectory(
+            prefix="tarit snapshot contract "
+        ) as temporary:
+            root = Path(temporary)
+            commands = root / "bin"
+            commands.mkdir()
+            fake = commands / "taritd"
+            fake.write_text(
+                "#!/bin/sh\nexec "
+                + shlex.join(
+                    [sys.executable, str(Path(__file__).resolve()), "--fake-taritd"]
+                )
+                + ' "$@"\n'
+            )
+            fake.chmod(0o755)
+            if shutil.which("setsid") is None:
+                # macOS lacks the launcher; retain its process-group behavior.
+                launcher = commands / "setsid"
+                launcher.write_text(
+                    "#!/usr/bin/env python3\nimport os, sys\n"
+                    "os.setsid()\nos.execvp(sys.argv[1], sys.argv[1:])\n"
+                )
+                launcher.chmod(0o755)
+            requests = root / "requests.jsonl"
+            pid_file = root / "taritd.pid"
+            env = {
+                **os.environ,
+                "PATH": str(commands) + os.pathsep + os.environ["PATH"],
+                "TMPDIR": str(root),
+                "TARITD_BIN": str(fake),
+                "SNAPSHOT_DISK_E2E_PORT": "",
+                "FAKE_TARITD_PID": str(pid_file),
+                "FAKE_SNAPSHOT_MODE": mode,
+                "FAKE_SNAPSHOT_REQUESTS": str(requests),
+            }
+            try:
+                result = subprocess.run(
+                    ["bash", str(HARNESS)],
+                    env=env,
+                    text=True,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    timeout=20,
+                )
+            finally:
+                # Also stop the fixture if a broken harness times out.
+                if pid_file.exists():
+                    try:
+                        os.killpg(int(pid_file.read_text()), signal.SIGTERM)
+                    except ProcessLookupError:
+                        pass
+            self.assertFalse(
+                list(root.glob("tarit-snapshot-disk-e2e.*")), "harness cleanup failed"
+            )
+            self.assertTrue(requests.exists(), result.stdout)
+            calls = [json.loads(line) for line in requests.read_text().splitlines()]
+            return result, calls
+
+    def test_opaque_handle_response_and_restore_requests(self):
+        result, calls = self.run_harness()
+        self.assertEqual(result.returncode, 0, result.stdout)
+        self.assertIn("RESULT: SNAPSHOT_DISK_PASS", result.stdout)
+        restores = [call["body"] for call in calls if call["path"] == "/v1/restore"]
+        self.assertEqual(restores, [{"snapshot_id": SNAPSHOT_ID}] * 2)
+        self.assertEqual(sum(call["method"] == "DELETE" for call in calls), 3)
+
+    def test_missing_snapshot_record_fails_before_restore(self):
+        result, calls = self.run_harness("missing-record")
+        self.assertNotEqual(result.returncode, 0)
+        self.assertIn("snapshot is missing from the local test database", result.stdout)
+        self.assertFalse(any(call["path"] == "/v1/restore" for call in calls))
+
+    def test_missing_snapshot_artifact_fails_before_restore(self):
+        result, calls = self.run_harness("missing-artifact")
+        self.assertNotEqual(result.returncode, 0)
+        self.assertFalse(any(call["path"] == "/v1/restore" for call in calls))
+        self.assertNotIn("RESULT: SNAPSHOT_DISK_PASS", result.stdout)
+
+    def test_shared_restore_state_is_rejected(self):
+        result, calls = self.run_harness("shared-restores")
+        self.assertNotEqual(result.returncode, 0)
+        self.assertIn("restored VMs shared writable disk state", result.stdout)
+        self.assertEqual(sum(call["path"] == "/v1/restore" for call in calls), 2)
+        self.assertNotIn("RESULT: SNAPSHOT_DISK_PASS", result.stdout)
+
+    def test_missing_runtime_record_is_rejected(self):
+        result, _ = self.run_harness("missing-runtime-record")
+        self.assertNotEqual(result.returncode, 0)
+        self.assertIn("VM has no recorded runtime overlay", result.stdout)
+        self.assertNotIn("RESULT: SNAPSHOT_DISK_PASS", result.stdout)
+
+    def test_missing_restore_overlay_is_rejected(self):
+        result, _ = self.run_harness("missing-restore-overlay")
+        self.assertNotEqual(result.returncode, 0)
+        self.assertIn("restored VM overlay is missing", result.stdout)
+        self.assertNotIn("RESULT: SNAPSHOT_DISK_PASS", result.stdout)
+
+
+if __name__ == "__main__":
+    if sys.argv[1:2] == ["--fake-taritd"]:
+        fake_taritd()
+    else:
+        unittest.main()