From 3bc618c2641801c8900abc875af5e948cf5c7405 Mon Sep 17 00:00:00 2001 From: codex Date: Mon, 27 Jul 2026 03:01:26 +0000 Subject: [PATCH] fix(cleanup): revalidate destructive backend plans --- bot_bottle/backend/firecracker/cleanup.py | 48 +++++++++++++++---- bot_bottle/backend/firecracker/launch.py | 42 ++++++++-------- .../backend/firecracker/lifecycle_lock.py | 30 ++++++++++++ bot_bottle/backend/macos_container/cleanup.py | 10 ++-- bot_bottle/cli/commands/cleanup.py | 6 ++- tests/unit/test_cli_cleanup_cross_backend.py | 25 ++++++++-- tests/unit/test_firecracker_cleanup.py | 42 +++++++++++++--- tests/unit/test_macos_container_cleanup.py | 16 +++++++ 8 files changed, 176 insertions(+), 43 deletions(-) create mode 100644 bot_bottle/backend/firecracker/lifecycle_lock.py diff --git a/bot_bottle/backend/firecracker/cleanup.py b/bot_bottle/backend/firecracker/cleanup.py index ac3e855f..21ec618c 100644 --- a/bot_bottle/backend/firecracker/cleanup.py +++ b/bot_bottle/backend/firecracker/cleanup.py @@ -30,7 +30,7 @@ from pathlib import Path from ...log import info from .. import EnumerationError -from . import util +from . import lifecycle_lock, util from .bottle_cleanup_plan import FirecrackerBottleCleanupPlan @@ -126,12 +126,42 @@ def prepare_cleanup() -> FirecrackerBottleCleanupPlan: def cleanup(plan: FirecrackerBottleCleanupPlan) -> None: - for pid in plan.vm_pids: - info(f"kill firecracker VM pid {pid}") + """Revalidate the preview under the launch lock, then remove its survivors.""" + with lifecycle_lock.hold(): + fresh = prepare_cleanup() + approved_pids = set(plan.vm_pids).intersection(fresh.vm_pids) + approved_dirs = set(plan.run_dirs).intersection(fresh.run_dirs) + for pid in sorted(approved_pids): + _terminate_orphan(pid, _run_root()) + for path in sorted(approved_dirs): + info(f"rm -rf {path}") + shutil.rmtree(path, ignore_errors=True) + + +def _terminate_orphan(pid: int, run_root: Path) -> None: + """Signal exactly the process identity that still owns an orphan config.""" + try: + pidfd = os.pidfd_open(pid) + except ProcessLookupError: + return + except OSError as exc: + raise EnumerationError( + f"could not pin Firecracker pid {pid} for cleanup: {exc}" + ) from exc + try: try: - os.kill(pid, signal.SIGTERM) - except ProcessLookupError: - pass - for path in plan.run_dirs: - info(f"rm -rf {path}") - shutil.rmtree(path, ignore_errors=True) + raw = Path(f"/proc/{pid}/cmdline").read_bytes() + except FileNotFoundError: + return + except OSError as exc: + raise EnumerationError( + f"could not revalidate Firecracker pid {pid}: {exc}" + ) from exc + command = raw.replace(b"\0", b" ").decode(errors="replace") + run_dir = _run_dir_of(command, run_root) + if run_dir is None or run_dir.is_dir(): + return + info(f"kill firecracker VM pid {pid}") + signal.pidfd_send_signal(pidfd, signal.SIGTERM) + finally: + os.close(pidfd) diff --git a/bot_bottle/backend/firecracker/launch.py b/bot_bottle/backend/firecracker/launch.py index 866e1dac..6fa0d0b8 100644 --- a/bot_bottle/backend/firecracker/launch.py +++ b/bot_bottle/backend/firecracker/launch.py @@ -46,7 +46,7 @@ from ...log import die, info, warn from ...supervisor.types import SUPERVISE_PORT from ..docker.egress import EGRESS_PORT from ..util import AGENT_CA_BUNDLE, AGENT_CA_PATH -from . import firecracker_vm, image_builder, isolation_probe, netpool, util +from . import firecracker_vm, image_builder, isolation_probe, lifecycle_lock, netpool, util from .bottle import FirecrackerBottle from .bottle_plan import FirecrackerBottlePlan from ...orchestrator.store.config_store import resolve_teardown_timeout @@ -164,25 +164,29 @@ def launch( ) # Step 6: build the per-bottle rootfs + SSH key, then boot. - run_dir = util.cache_dir() / "run" / plan.slug - run_dir.mkdir(parents=True, exist_ok=True) - # Remove the run dir on teardown so the per-bottle rootfs.ext4 (~1G) - # doesn't leak. Registered before vm.terminate below so it runs *after* - # it (ExitStack is LIFO): the VM is gone before we rm its rootfs. - stack.callback(lambda: shutil.rmtree(run_dir, ignore_errors=True)) - rootfs = run_dir / "rootfs.ext4" - util.build_rootfs_ext4(agent_base, rootfs) - private_key, pubkey = util.generate_keypair(run_dir) + # Cleanup takes the same lock while refreshing its process snapshot. + # Hold it until the VMM exists so a newly-created run dir can never be + # mistaken for an orphan in the build-before-boot window. + with lifecycle_lock.hold(): + run_dir = util.cache_dir() / "run" / plan.slug + run_dir.mkdir(parents=True, exist_ok=True) + # Remove the run dir on teardown so the per-bottle rootfs.ext4 (~1G) + # doesn't leak. Registered before vm.terminate below so it runs + # *after* it (ExitStack is LIFO). + stack.callback(lambda: shutil.rmtree(run_dir, ignore_errors=True)) + rootfs = run_dir / "rootfs.ext4" + util.build_rootfs_ext4(agent_base, rootfs) + private_key, pubkey = util.generate_keypair(run_dir) - vm = firecracker_vm.boot( - name=plan.container_name, - rootfs=rootfs, - tap=slot.iface, - guest_ip=slot.guest_ip, - host_ip=slot.host_ip, - pubkey=pubkey, - run_dir=run_dir, - ) + vm = firecracker_vm.boot( + name=plan.container_name, + rootfs=rootfs, + tap=slot.iface, + guest_ip=slot.guest_ip, + host_ip=slot.host_ip, + pubkey=pubkey, + run_dir=run_dir, + ) stack.callback(vm.terminate) firecracker_vm.wait_for_ssh(vm, private_key) persist_env_var_secret(private_key, slot.guest_ip, ctx.env_var_secret) diff --git a/bot_bottle/backend/firecracker/lifecycle_lock.py b/bot_bottle/backend/firecracker/lifecycle_lock.py new file mode 100644 index 00000000..3dead046 --- /dev/null +++ b/bot_bottle/backend/firecracker/lifecycle_lock.py @@ -0,0 +1,30 @@ +"""Serialize Firecracker run-directory creation with orphan cleanup.""" + +from __future__ import annotations + +import fcntl +from contextlib import contextmanager +from pathlib import Path +from typing import Generator + +from . import util + + +def _lock_path() -> Path: + return util.cache_dir() / "run.lifecycle.lock" + + +@contextmanager +def hold() -> Generator[None]: + """Exclude cleanup while a launch directory lacks a visible VMM.""" + path = _lock_path() + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + fcntl.flock(handle, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(handle, fcntl.LOCK_UN) + + +__all__ = ["hold"] diff --git a/bot_bottle/backend/macos_container/cleanup.py b/bot_bottle/backend/macos_container/cleanup.py index bd7d5ee5..f3beb2de 100644 --- a/bot_bottle/backend/macos_container/cleanup.py +++ b/bot_bottle/backend/macos_container/cleanup.py @@ -4,7 +4,8 @@ from __future__ import annotations import subprocess -from ...log import info, warn +from .. import EnumerationError +from ...log import info from . import util as container_mod from .bottle_cleanup_plan import MacosContainerBottleCleanupPlan @@ -19,8 +20,8 @@ def _list_prefixed_containers() -> list[str]: check=False, ) if result.returncode != 0: - warn(f"container list failed: {result.stderr.strip()}") - return [] + detail = result.stderr.strip() or f"exit {result.returncode}" + raise EnumerationError(f"container list failed: {detail}") return sorted( name for name in (line.strip() for line in result.stdout.splitlines()) if name.startswith(_PREFIX) @@ -35,7 +36,8 @@ def _list_prefixed_networks() -> list[str]: check=False, ) if result.returncode != 0: - return [] + detail = result.stderr.strip() or f"exit {result.returncode}" + raise EnumerationError(f"container network list failed: {detail}") return sorted( name for name in (line.strip() for line in result.stdout.splitlines()) if name.startswith(_PREFIX) diff --git a/bot_bottle/cli/commands/cleanup.py b/bot_bottle/cli/commands/cleanup.py index d83ef14c..d64013d6 100644 --- a/bot_bottle/cli/commands/cleanup.py +++ b/bot_bottle/cli/commands/cleanup.py @@ -52,7 +52,11 @@ def cmd_cleanup(_argv: list[str]) -> int: info("cleanup: skipped") return 0 - for name, backend, plan in prepared: + # Confirmation authorizes a fresh authoritative snapshot, not blind use of + # identities that may have changed while the operator reviewed the preview. + refreshed = [(name, backend, backend.prepare_cleanup()) + for name, backend, _plan in prepared] + for name, backend, plan in refreshed: if plan.empty: continue backend.cleanup(plan) diff --git a/tests/unit/test_cli_cleanup_cross_backend.py b/tests/unit/test_cli_cleanup_cross_backend.py index 47ab1334..d66fa673 100644 --- a/tests/unit/test_cli_cleanup_cross_backend.py +++ b/tests/unit/test_cli_cleanup_cross_backend.py @@ -41,8 +41,8 @@ class TestCmdCleanup(unittest.TestCase): ): self.assertEqual(0, cmd.cmd_cleanup([])) - docker.prepare_cleanup.assert_called_once() - fc.prepare_cleanup.assert_called_once() + self.assertEqual(2, docker.prepare_cleanup.call_count) + self.assertEqual(2, fc.prepare_cleanup.call_count) docker.cleanup.assert_called_once_with(docker_plan) fc.cleanup.assert_called_once_with(fc_plan) @@ -68,7 +68,7 @@ class TestCmdCleanup(unittest.TestCase): ): self.assertEqual(0, cmd.cmd_cleanup([])) - docker.prepare_cleanup.assert_called_once() + self.assertEqual(2, docker.prepare_cleanup.call_count) docker.cleanup.assert_called_once_with(docker_plan) macos.prepare_cleanup.assert_not_called() @@ -135,6 +135,25 @@ class TestCmdCleanup(unittest.TestCase): docker.cleanup.assert_called_once_with(docker_plan) fc.cleanup.assert_not_called() + def test_executes_refreshed_plan_after_confirmation(self): + backend = MagicMock() + preview = MagicMock(empty=False) + refreshed = MagicMock(empty=False) + backend.prepare_cleanup.side_effect = [preview, refreshed] + + with patch.object( + cmd, "known_backend_names", return_value=("firecracker",), + ), patch.object( + cmd, "get_bottle_backend", return_value=backend, + ), patch.object( + cmd, "has_backend", return_value=True, + ), patch.object( + cmd, "_prompt_yes", return_value=True, + ): + self.assertEqual(0, cmd.cmd_cleanup([])) + + backend.cleanup.assert_called_once_with(refreshed) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_firecracker_cleanup.py b/tests/unit/test_firecracker_cleanup.py index c6cd55c9..069cfe38 100644 --- a/tests/unit/test_firecracker_cleanup.py +++ b/tests/unit/test_firecracker_cleanup.py @@ -132,18 +132,46 @@ class TestCleanupRemoval(unittest.TestCase): vm_pids=(101,), run_dirs=("/run/dev-x",), ) - with patch.object(fc_cleanup.os, "kill") as kill, \ + with patch.object(fc_cleanup, "prepare_cleanup", return_value=plan), \ + patch.object(fc_cleanup, "_run_root", return_value=Path("/run")), \ + patch.object(fc_cleanup, "_terminate_orphan") as terminate, \ patch.object(fc_cleanup.shutil, "rmtree") as rmtree, \ patch.object(fc_cleanup, "info"): fc_cleanup.cleanup(plan) - kill.assert_called_once() + terminate.assert_called_once_with(101, Path("/run")) rmtree.assert_called_once_with("/run/dev-x", ignore_errors=True) - def test_cleanup_tolerates_dead_pid(self): - plan = FirecrackerBottleCleanupPlan(vm_pids=(999,)) - with patch.object(fc_cleanup.os, "kill", side_effect=ProcessLookupError), \ - patch.object(fc_cleanup, "info"): - fc_cleanup.cleanup(plan) # must not raise + def test_cleanup_skips_resources_no_longer_in_refreshed_plan(self): + preview = FirecrackerBottleCleanupPlan( + vm_pids=(999,), run_dirs=("/run/reused",), + ) + with patch.object( + fc_cleanup, "prepare_cleanup", + return_value=FirecrackerBottleCleanupPlan(), + ), patch.object(fc_cleanup, "_terminate_orphan") as terminate, \ + patch.object(fc_cleanup.shutil, "rmtree") as rmtree: + fc_cleanup.cleanup(preview) + terminate.assert_not_called() + rmtree.assert_not_called() + + def test_pidfd_prevents_pid_reuse_from_signalling_unrelated_process(self): + with patch.object(fc_cleanup.os, "pidfd_open", return_value=7), \ + patch.object( + fc_cleanup.Path, "read_bytes", + return_value=b"/usr/bin/python\0worker.py\0", + ), patch.object(fc_cleanup.signal, "pidfd_send_signal") as send, \ + patch.object(fc_cleanup.os, "close"): + fc_cleanup._terminate_orphan(101, Path("/run")) + send.assert_not_called() + + def test_pidfd_signals_revalidated_orphan(self): + command = b"firecracker\0--config-file\0/run/gone/config.json\0" + with patch.object(fc_cleanup.os, "pidfd_open", return_value=7), \ + patch.object(fc_cleanup.Path, "read_bytes", return_value=command), \ + patch.object(fc_cleanup.signal, "pidfd_send_signal") as send, \ + patch.object(fc_cleanup.os, "close"), patch.object(fc_cleanup, "info"): + fc_cleanup._terminate_orphan(101, Path("/run")) + send.assert_called_once_with(7, fc_cleanup.signal.SIGTERM) class TestCleanupPlan(unittest.TestCase): diff --git a/tests/unit/test_macos_container_cleanup.py b/tests/unit/test_macos_container_cleanup.py index fcc9886d..fbe13a4e 100644 --- a/tests/unit/test_macos_container_cleanup.py +++ b/tests/unit/test_macos_container_cleanup.py @@ -42,6 +42,22 @@ class TestMacosContainerCleanup(unittest.TestCase): run.call_args_list[1].args[0], ) + def test_container_enumeration_failure_aborts(self): + completed = cleanup.subprocess.CompletedProcess( + args=[], returncode=1, stdout="", stderr="service unavailable", + ) + with patch.object(cleanup.subprocess, "run", return_value=completed), \ + self.assertRaisesRegex(EnumerationError, "service unavailable"): + cleanup._list_prefixed_containers() + + def test_network_enumeration_failure_aborts(self): + completed = cleanup.subprocess.CompletedProcess( + args=[], returncode=1, stdout="", stderr="service unavailable", + ) + with patch.object(cleanup.subprocess, "run", return_value=completed), \ + self.assertRaisesRegex(EnumerationError, "service unavailable"): + cleanup._list_prefixed_networks() + class TestMacosContainerEnumerate(unittest.TestCase): """The backend launches bottles again (PRD 0070), so enumeration is real