diff --git a/tools/prepare_hosted_shells.py b/tools/prepare_hosted_shells.py index 65a29e16..c2f497d2 100644 --- a/tools/prepare_hosted_shells.py +++ b/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 diff --git a/tools/test_prepare_hosted_shells.py b/tools/test_prepare_hosted_shells.py index b7110612..9036d951 100644 --- a/tools/test_prepare_hosted_shells.py +++ b/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()