ci: observe hosted preparation on the experimental stack (#81)

tools/prepare_hosted_shells.py

@@ -29,6 +29,11 @@
 ROOTFS_SHA256 = "8251e27ffff381a4af5f41dcb94d867de3e0d9774a9241908ab34555d99315ea"
 MAX_DOWNLOAD = 400 * 1024 * 1024
 STEPS = {"preflight": 45, "download": 120, "import": 90, "check-before": 25, "check-after": 25}
+COMMANDS_BY_STEP = {"preflight": ("cmd", "ps51", "ps7", "list"), "download": (),
+                    "import": ("list-before-import", "import"),
+                    "check-before": ("list", "startup"), "check-after": ("list", "startup")}
+COMMAND_DEADLINES = {"cmd": 10, "ps51": 10, "ps7": 10, "list": 10,
+                     "list-before-import": 10, "import": 80, "startup": 20}
 
 
 class PreparationError(Exception):
@@ -165,16 +170,95 @@ def copy_verified_rootfs(response, destination, *, expected_hash=ROOTFS_SHA256,
         raise PreparationError("download_hash")
 
 
+def validate_command_events(value, step):
+    """Whitelist an ordered prefix; an invoking tail is not a completed command."""
+    if (step not in COMMANDS_BY_STEP or not isinstance(value, dict) or set(value) != {"events"} or
+            not isinstance(value["events"], list) or len(value["events"]) > 8):
+        raise ValueError("command event schema")
+    expected = COMMANDS_BY_STEP[step]
+    index, waiting, timed_out = 0, False, False
+    for event in value["events"]:
+        if (not isinstance(event, dict) or
+                set(event) != {"command", "phase", "deadline_ms", "elapsed_ms", "returncode"} or
+                index >= len(expected) or timed_out or event["command"] != expected[index]):
+            raise ValueError("command event identity")
+        phase = event["phase"]
+        if (phase not in (("returned", "timed_out") if waiting else ("invoking",)) or
+                type(event["deadline_ms"]) is not int or
+                event["deadline_ms"] != 1000 * COMMAND_DEADLINES[expected[index]] or
+                type(event["elapsed_ms"]) is not int or not 0 <= event["elapsed_ms"] <= STEPS[step] * 1000 + 10000):
+            raise ValueError("command event boundary")
+        if phase == "returned":
+            if type(event["returncode"]) is not int or not -(2**31) <= event["returncode"] < 2**32:
+                raise ValueError("command exit code")
+            waiting = False
+            index += 1
+        else:
+            if event["returncode"] is not None or (phase == "invoking" and event["elapsed_ms"] != 0):
+                raise ValueError("unobserved command exit")
+            waiting = phase == "invoking"
+            timed_out = phase == "timed_out"
+    return value["events"]
+
+
+def read_command_events(directory, step):
+    with (directory / "command-events.json").open("rb") as stream:
+        raw = stream.read(8193)
+    if len(raw) > 8192:
+        raise ValueError("command event size")
+    return validate_command_events(json.loads(raw), step)
+
+
+def write_command_event(directory, name, phase, elapsed_ms, returncode=None):
+    try:
+        path = directory / "command-events.json"
+        events = read_command_events(directory, directory.name) if path.exists() else []
+        event = {"command": name, "phase": phase, "deadline_ms": 1000 * COMMAND_DEADLINES[name],
+                 "elapsed_ms": elapsed_ms, "returncode": returncode}
+        value = {"events": events + [event]}
+        validate_command_events(value, directory.name)
+        write_json(path, value)
+    except (OSError, ValueError, KeyError, RecursionError):
+        raise PreparationError("command_observation") from None
+
+
+def public_command_observation(directory, step):
+    # Read only after the worker is reaped. Partial traces survive its deadline;
+    # raw arguments/output and malformed data never enter the public report.
+    if step not in COMMANDS_BY_STEP:
+        return {"status": "invalid", "events": []}
+    if not COMMANDS_BY_STEP[step]:
+        return {"status": "not_applicable", "events": []}
+    try:
+        return {"status": "observed", "events": read_command_events(directory, step)}
+    except FileNotFoundError:
+        return {"status": "not_observed", "events": []}
+    except (OSError, ValueError, RecursionError):
+        return {"status": "invalid", "events": []}
+
+
 def run_command(command, directory, name, *, cwd, timeout=10, data=None):
     # This function only runs inside a GO-gated worker already in the parent's
     # exact Job. No shell string construction or global-process cleanup.
+    if name not in COMMANDS_BY_STEP.get(directory.name, ()) or timeout != COMMAND_DEADLINES[name]:
+        raise PreparationError("command_observation")
+    def observe(phase, elapsed_ms, returncode=None):
+        try:
+            write_command_event(directory, name, phase, elapsed_ms, returncode)
+        except PreparationError:
+            # A missing/partial trace cannot replace the actual command result.
+            pass
     stdout_path, stderr_path = directory / (name + ".stdout"), directory / (name + ".stderr")
     with stdout_path.open("xb") as stdout, stderr_path.open("xb") as stderr:
+        observe("invoking", 0)
+        started = time.monotonic_ns()
         try:
             result = subprocess.run(command, input=data, stdout=stdout, stderr=stderr,
                                     cwd=cwd, timeout=timeout, creationflags=subprocess.CREATE_NO_WINDOW)
         except subprocess.TimeoutExpired:
+            observe("timed_out", (time.monotonic_ns() - started) // 1000000)
             raise PreparationError("command_deadline") from None
+        observe("returned", (time.monotonic_ns() - started) // 1000000, result.returncode)
     if stdout_path.stat().st_size > 65536:
         raise PreparationError("command_report_size")
     return result.returncode, stdout_path.read_bytes()
@@ -268,7 +352,8 @@ def worker(step, out):
 FAILURES = {"hosted_guard", "workspace_guard", "windows_admin_guard", "worker_job_guard", "cwd_identity", "wsl_report_size",
             "wsl_list_encoding", "wsl_list_identity", "marker_encoding", "shell_startup", "download_deadline",
             "download_size", "download_hash", "command_deadline", "command_report_size", "worker_directory_guard",
-            "wsl_list", "download_response", "ubuntu_already_present", "wsl_import", "ubuntu_missing", "worker_step"}
+            "wsl_list", "download_response", "ubuntu_already_present", "wsl_import", "ubuntu_missing", "worker_step",
+            "command_observation"}
 
 
 def run_step(step, out):
@@ -375,6 +460,8 @@ def terminate_job():
                 job.close()
         (private / "worker.stdout").write_bytes(stdout)
         (private / "worker.stderr").write_bytes(stderr)
+    record["command_observation"] = (public_command_observation(private, step) if cleanup["root_reaped"] else
+                                     {"status": "worker_not_reaped", "events": []})
     return record
 
 

tools/test_prepare_hosted_shells.py

@@ -166,5 +166,155 @@ def test_local_runtime_is_rejected_before_any_subprocess_or_download(self):
         urlopen.assert_not_called()
 
 
+class CommandBoundaryTests(unittest.TestCase):
+    @staticmethod
+    def event(command="cmd", phase="invoking", elapsed=0, code=None):
+        return {"command": command, "phase": phase, "deadline_ms": 10000,
+                "elapsed_ms": elapsed, "returncode": code}
+
+    def test_timeout_keeps_actual_command_boundary_without_private_arguments(self):
+        secret = "synthetic-private-command-path"
+        with tempfile.TemporaryDirectory() as temporary:
+            directory = Path(temporary) / "preflight"
+            directory.mkdir()
+            failure = prep.subprocess.TimeoutExpired([secret], 10, output=b"synthetic-private-output")
+            with mock.patch.object(prep.subprocess, "run", side_effect=failure) as command, \
+                 mock.patch.object(prep.time, "monotonic_ns", side_effect=[0, 10000000000]), \
+                 mock.patch.object(prep.subprocess, "CREATE_NO_WINDOW", 0, create=True):
+                with self.assertRaisesRegex(prep.PreparationError, "^command_deadline$"):
+                    prep.run_command([secret], directory, "cmd", cwd=directory)
+            command.assert_called_once()
+            self.assertEqual(command.call_args.kwargs["timeout"], 10)
+            data = json.loads((directory / "command-events.json").read_text(encoding="utf-8"))
+            self.assertEqual(data, {"events": [
+                {"command": "cmd", "phase": "invoking", "deadline_ms": 10000, "elapsed_ms": 0, "returncode": None},
+                {"command": "cmd", "phase": "timed_out", "deadline_ms": 10000, "elapsed_ms": 10000, "returncode": None}]})
+            self.assertNotIn(secret, json.dumps(data))
+            self.assertNotIn("synthetic-private-output", json.dumps(data))
+
+    def test_returned_command_and_partial_invocation_are_reported_independently(self):
+        with tempfile.TemporaryDirectory() as temporary:
+            directory = Path(temporary) / "preflight"
+            directory.mkdir()
+            def returned(command, **kwargs):
+                kwargs["stdout"].write(b"private-native-output")
+                return prep.subprocess.CompletedProcess(command, 7)
+            with mock.patch.object(prep.subprocess, "run", side_effect=returned), \
+                 mock.patch.object(prep.time, "monotonic_ns", side_effect=[0, 123000000]), \
+                 mock.patch.object(prep.subprocess, "CREATE_NO_WINDOW", 0, create=True):
+                self.assertEqual(prep.run_command(["private-command"], directory, "cmd", cwd=directory),
+                                 (7, b"private-native-output"))
+            self.assertEqual(prep.public_command_observation(directory, "preflight"),
+                             {"status": "observed", "events": [self.event(), self.event(phase="returned", elapsed=123, code=7)]})
+            prep.write_json(directory / "command-events.json", {"events": [self.event()]})
+            self.assertEqual(prep.public_command_observation(directory, "preflight"),
+                             {"status": "observed", "events": [self.event()]})
+
+    def test_unknown_or_unordered_data_is_rejected_without_echoing_it(self):
+        good = self.event()
+        values = [
+            {"events": [{**good, "command": "private-path"}]},
+            {"events": [{**good, "argv": ["private-token"]}]},
+            {"events": [self.event(phase="returned", code=0)]},
+            {"events": [good, self.event(phase="returned", code=True)]},
+            {"events": [good, self.event(phase="timed_out", code=91)]},
+            {"events": [good, self.event(phase="returned", code=0), self.event(command="ps7")]},
+            {"events": [{**good, "deadline_ms": 11000}]},
+            {"events": [good, self.event(phase="returned", elapsed=-1, code=0)]},
+            {"events": [good] * 9},
+        ]
+        with tempfile.TemporaryDirectory() as temporary:
+            directory = Path(temporary) / "preflight"
+            directory.mkdir()
+            for value in values:
+                with self.subTest(value=value):
+                    prep.write_json(directory / "command-events.json", value)
+                    self.assertEqual(prep.public_command_observation(directory, "preflight"),
+                                     {"status": "invalid", "events": []})
+            (directory / "command-events.json").write_bytes(b"private-token" * 1000)
+            self.assertEqual(prep.public_command_observation(directory, "preflight"),
+                             {"status": "invalid", "events": []})
+            (directory / "command-events.json").write_bytes(b"[" * 2000 + b"]" * 2000)
+            self.assertEqual(prep.public_command_observation(directory, "preflight"),
+                             {"status": "invalid", "events": []})
+
+    def test_trace_write_failure_cannot_replace_observed_timeout(self):
+        with tempfile.TemporaryDirectory() as temporary:
+            directory = Path(temporary) / "preflight"
+            directory.mkdir()
+            with mock.patch.object(prep.subprocess, "run", side_effect=prep.subprocess.TimeoutExpired(["synthetic"], 10)), \
+                 mock.patch.object(prep, "write_command_event", side_effect=[None, prep.PreparationError("command_observation")]), \
+                 mock.patch.object(prep.subprocess, "CREATE_NO_WINDOW", 0, create=True):
+                with self.assertRaisesRegex(prep.PreparationError, "^command_deadline$"):
+                    prep.run_command(["synthetic"], directory, "cmd", cwd=directory)
+
+    def test_trace_write_failure_preserves_command_and_marker_result(self):
+        marker = "TERMFORGE_HOSTED_CMD_READY"
+        output = (marker + "\r\n").encode("ascii")
+        write_json = prep.write_json
+        for phase in ("invoking", "returned"):
+            for code in (0, 7):
+                with self.subTest(phase=phase, code=code), tempfile.TemporaryDirectory() as temporary:
+                    directory = Path(temporary) / "preflight"
+                    directory.mkdir()
+                    def failed_write(path, value):
+                        if value["events"][-1]["phase"] == phase:
+                            raise OSError("synthetic-private-storage-error")
+                        write_json(path, value)
+                    def returned(command, **kwargs):
+                        kwargs["stdout"].write(output)
+                        return prep.subprocess.CompletedProcess(command, code)
+                    with mock.patch.object(prep, "write_json", side_effect=failed_write), \
+                         mock.patch.object(prep.subprocess, "run", side_effect=returned) as command, \
+                         mock.patch.object(prep.subprocess, "CREATE_NO_WINDOW", 0, create=True):
+                        result_code, result_output = prep.run_command(["synthetic"], directory, "cmd", cwd=directory)
+                    command.assert_called_once()
+                    self.assertEqual((result_code, result_output), (code, output))
+                    if code:
+                        with self.assertRaisesRegex(prep.PreparationError, "^shell_startup$"):
+                            prep.validate_marker(result_output, result_code, marker)
+                    else:
+                        prep.validate_marker(result_output, result_code, marker)
+                    observation = prep.public_command_observation(directory, "preflight")
+                    self.assertEqual(observation, {"status": "not_observed", "events": []} if phase == "invoking" else
+                                     {"status": "observed", "events": [self.event()]})
+
+    def test_parent_retains_command_or_worker_deadline_after_cleanup(self):
+        for worker_timeout in (False, True):
+            with self.subTest(worker_timeout=worker_timeout), tempfile.TemporaryDirectory() as temporary:
+                out = Path(temporary)
+                events = [self.event(), self.event(phase="returned", elapsed=2, code=0), self.event(command="ps51")]
+                if not worker_timeout:
+                    events.append(self.event(command="ps51", phase="timed_out", elapsed=10000))
+                process = mock.Mock(pid=123, returncode=None)
+                process.poll.side_effect = lambda: process.returncode
+                calls = []
+                def communicate(*args, **kwargs):
+                    calls.append(kwargs["timeout"])
+                    prep.write_json(out / "_private/preflight/command-events.json", {"events": events})
+                    if worker_timeout and len(calls) == 1:
+                        raise prep.subprocess.TimeoutExpired(["synthetic-worker"], 45)
+                    process.returncode = 99 if worker_timeout else 1
+                    return (b"" if worker_timeout else b'{"failure_code":"command_deadline"}', b"private-worker-error")
+                process.communicate.side_effect = communicate
+                job = mock.Mock()
+                job.drain.return_value = {"active": 0, "total": 2, "terminated": 0}
+                audit = mock.Mock(stable=True, members=[], private=[])
+                audit.process_ids.return_value = []
+                with mock.patch.object(prep, "OwnedJob", return_value=job), \
+                     mock.patch.object(prep, "JobMembers", return_value=audit), \
+                     mock.patch.object(prep.subprocess, "Popen", return_value=process), \
+                     mock.patch.object(prep.subprocess, "CREATE_NO_WINDOW", 0, create=True):
+                    result = prep.run_step("preflight", out)
+                self.assertEqual(result["status"], "failed")
+                self.assertEqual(result["failure_code"], "worker_deadline" if worker_timeout else "command_deadline")
+                self.assertEqual(result["command_observation"], {"status": "observed", "events": events})
+                self.assertTrue(result["cleanup"]["root_reaped"])
+                self.assertEqual(result["cleanup"]["root_natural_exit"], not worker_timeout)
+                self.assertEqual(result["cleanup"]["job_termination_requests"], int(worker_timeout))
+                self.assertEqual(calls, [45, 5] if worker_timeout else [45])
+                self.assertNotIn("private-worker-error", json.dumps(result))
+
+
 if __name__ == "__main__":
     unittest.main()