from __future__ import annotations import hashlib import json import shutil import sys import tempfile import unittest from pathlib import Path from unittest.mock import patch from uvm.config import Settings from uvm.errors import UvmError from uvm.images import GUEST_AGENT_PREPARE, ImageStore, guest_agent_files from uvm.system import CommandRunner class GuestRuntimeBootstrapTests(unittest.TestCase): def test_runtime_install_is_guest_only_skipped_when_ready_and_stops_on_apt_failure(self) -> None: with tempfile.TemporaryDirectory() as temporary: directory = Path(temporary) python = directory / "python3" apt = directory / "apt-get" ready = directory / "ready" python.write_text( f"#!{sys.executable}\nimport os, sys\nfrom pathlib import Path\n" "sys.exit(0 if Path(os.environ['READY']).exists() else 1)\n", encoding="utf-8", ) apt.write_text( f"#!{sys.executable}\nimport json, os, sys\nfrom pathlib import Path\n" "print(json.dumps(sys.argv[1:]))\n" "if sys.argv[-1] == 'update' and os.environ['FAIL_UPDATE'] == '1': sys.exit(100)\n" "if sys.argv[-1] == 'python3': Path(os.environ['READY']).touch()\n", encoding="utf-8", ) python.chmod(0o755) apt.chmod(0o755) script = GUEST_AGENT_PREPARE.replace("/usr/bin/python3", str(python)).replace("/usr/bin/apt-get", str(apt)) for installed, fail in ((True, False), (False, False), (False, True)): with self.subTest(installed=installed, fail=fail): ready.unlink(missing_ok=True) if installed: ready.touch() result = CommandRunner(emit=None).run(( "env", f"READY={ready}", f"FAIL_UPDATE={int(fail)}", "/bin/sh", "-c", script, ), check=False, capture=True) calls = [json.loads(line) for line in result.stdout.splitlines()] self.assertEqual(result.returncode, 100 if fail else 0) self.assertEqual(len(calls), 0 if installed else 1 if fail else 2) if calls: self.assertEqual(calls[0][-1], "update") self.assertIn("APT::Update::Error-Mode=any", calls[0]) if not installed and not fail: self.assertEqual(calls[-1][-2:], ["install", "python3"]) @unittest.skipUnless(all(shutil.which(tool) for tool in ("mkfs.ext4", "debugfs", "e2fsck", "cp")), "e2fsprogs and coreutils are required for offline guest agent checks") class GuestAgentImageTests(unittest.TestCase): def setUp(self) -> None: temporary = tempfile.TemporaryDirectory() self.addCleanup(temporary.cleanup) self.base = Path(temporary.name) self.settings = Settings(base=self.base / "uvm") self.settings.images_dir.mkdir(parents=True) self.disk = self.base / "VM with spaces" / "rootfs.ext4" self.disk.parent.mkdir() self.runner = CommandRunner(emit=None) self.images = ImageStore(self.settings, self.runner) self.retained = { "/etc/shadow": "worker:$6$retained-hash:1:0:99999:7:::\n", "/etc/ssh/ssh_host_ed25519_key": "retained private SSH key\n", "/home/worker/data": "retained application data\n", } def make_disk(self, marker="uvm-ubuntu-jammy-v3\n") -> None: root = self.base / "tree" if root.exists(): shutil.rmtree(root) payloads = { **self.retained, "/usr/lib/systemd/systemd": "systemd fixture", "/usr/bin/apt-get": "apt fixture", } if marker is not None: payloads["/etc/uvm-image"] = marker for path, content in payloads.items(): target = root / path.lstrip("/") target.parent.mkdir(parents=True, exist_ok=True) target.write_text(content, encoding="utf-8") target.chmod(0o644) for directory in (root, *(path for path in root.rglob("*") if path.is_dir())): directory.chmod(0o755) with self.settings.rootfs_image.open("wb") as image: image.truncate(32 * 1024 * 1024) self.runner.run(("mkfs.ext4", "-q", "-F", "-b", "4096", "-d", root, self.settings.rootfs_image), capture=True) # mkfs -d retains the creating user's ownership. Model the root-built # guest while allowing this test itself to run without root. paths = {"/"} for path in payloads: paths.add(path) paths.update(str(parent) for parent in Path(path).parents) commands = "\n".join(f"set_inode_field {path} {field} 0" for path in paths for field in ("uid", "gid")) self.runner.run(("debugfs", "-w", self.settings.rootfs_image), input_text=commands + "\n", capture=True) self.runner.run(("cp", "--sparse=always", self.settings.rootfs_image, self.disk), capture=True) self.disk.chmod(0o600) def guest(self, command: str, *, disk: Path | None = None) -> str: return self.runner.run(("debugfs", "-R", command, disk or self.disk), capture=True).stdout def snapshot(self) -> tuple[int, bytes]: return self.disk.stat().st_ino, hashlib.sha256(self.disk.read_bytes()).digest() def replace_guest_file(self, path: str, contents: str, *, disk: Path | None = None) -> None: source = self.base / "edited" source.write_text(contents, encoding="utf-8") source.chmod(0o644) self.runner.run(("debugfs", "-w", disk or self.disk), input_text=( f"rm {path}\nwrite {source} {path}\nset_inode_field {path} uid 0\nset_inode_field {path} gid 0\n" ), capture=True) def test_existing_guest_is_upgraded_atomically_with_private_identity_and_retained_data(self) -> None: self.make_disk() template = hashlib.sha256(self.settings.rootfs_image.read_bytes()).digest() self.assertTrue(self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True)) self.assertEqual(json.loads(self.guest("cat /etc/uvm-agent.json")), {"vm_id": "vm-test", "username": "worker"}) self.assertRegex(self.guest("stat /etc/uvm-agent.json"), r"Mode:\s+0?600") self.assertRegex(self.guest("stat /etc/uvm-agent.json"), r"User:\s+0\s+Group:\s+0") for path, content in guest_agent_files().items(): self.assertEqual(self.guest(f"cat {path}"), content.decode("utf-8")) self.assertIn('Fast link dest: "/etc/systemd/system/uvm-agent.service"', self.guest("stat /etc/systemd/system/multi-user.target.wants/uvm-agent.service")) self.assertEqual(self.guest("cat /etc/uvm-image"), "uvm-ubuntu-jammy-v3\n") for path, content in self.retained.items(): self.assertEqual(self.guest(f"cat {path}"), content) self.assertEqual(hashlib.sha256(self.settings.rootfs_image.read_bytes()).digest(), template) self.assertEqual(self.disk.stat().st_mode & 0o777, 0o600) self.runner.run(("e2fsck", "-f", "-n", self.disk), capture=True) self.assertEqual(list(self.disk.parent.iterdir()), [self.disk]) self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) self.assertEqual(self.guest("cat /home/worker/data"), self.retained["/home/worker/data"]) def test_all_compatible_images_get_vsock_agent_setup_without_opt_in(self) -> None: for marker in ("uvm-ubuntu-jammy-v1\n", "uvm-ubuntu-jammy-v3\n", "uvm-ubuntu-jammy-v4\n", "uvm-ubuntu-jammy-v5\n", "custom-image\n", None): with self.subTest(marker=marker): self.make_disk(marker) template = hashlib.sha256(self.settings.rootfs_image.read_bytes()).digest() self.assertTrue(self.images.provision_guest_agent(self.disk, "vm-test", "worker")) self.assertEqual(json.loads(self.guest("cat /etc/uvm-agent.json")), {"vm_id": "vm-test", "username": "worker"}) self.assertIn('Fast link dest: "/etc/systemd/system/uvm-agent.service"', self.guest("stat /etc/systemd/system/multi-user.target.wants/uvm-agent.service")) self.assertEqual(self.guest("cat /usr/lib/uvm/agent.py"), guest_agent_files()["/usr/lib/uvm/agent.py"].decode()) for path, content in self.retained.items(): self.assertEqual(self.guest(f"cat {path}"), content) self.assertEqual(hashlib.sha256(self.settings.rootfs_image.read_bytes()).digest(), template) def test_current_agent_setup_is_read_only_on_later_boots_and_tracks_vm_identity(self) -> None: self.make_disk() self.images.provision_guest_agent(self.disk, "vm-test", "worker") before = self.snapshot() with patch.object(self.images, "_copy_sparse") as copy: self.assertFalse(self.images.provision_guest_agent(self.disk, "vm-test", "worker")) copy.assert_not_called() self.assertEqual(self.snapshot(), before) self.assertTrue(self.images.provision_guest_agent(self.disk, "vm-recreated", "worker")) self.assertEqual(json.loads(self.guest("cat /etc/uvm-agent.json")), {"vm_id": "vm-recreated", "username": "worker"}) def test_default_setup_reenables_an_existing_agent_service(self) -> None: self.make_disk() self.images.provision_guest_agent(self.disk, "vm-test", "worker") self.runner.run(("debugfs", "-w", "-R", "rm /etc/systemd/system/multi-user.target.wants/uvm-agent.service", self.disk), capture=True) self.assertTrue(self.images.provision_guest_agent(self.disk, "vm-test", "worker")) self.assertIn('Fast link dest: "/etc/systemd/system/uvm-agent.service"', self.guest("stat /etc/systemd/system/multi-user.target.wants/uvm-agent.service")) def test_incompatible_images_fail_default_setup_without_changing_the_disk(self) -> None: for binary in ("/usr/lib/systemd/systemd", "/usr/bin/apt-get"): with self.subTest(binary=binary): self.make_disk(marker=None) self.runner.run(("debugfs", "-w", "-R", f"rm {binary}", self.disk), capture=True) before = self.snapshot() with self.assertRaisesRegex(UvmError, "automatic vsock guest-agent setup requires a systemd/apt image"): self.images.provision_guest_agent(self.disk, "vm-test", "worker") self.assertEqual(self.snapshot(), before) def test_cached_template_without_a_manifest_only_adopts_the_known_v4_payload(self) -> None: self.make_disk() self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) self.runner.run(("debugfs", "-w", "-R", "rm /etc/uvm-agent-files.json", self.disk), capture=True) old_payload = "previous packaged agent\n" self.replace_guest_file("/usr/lib/uvm/agent.py", old_payload) before = self.snapshot() with self.assertRaisesRegex(UvmError, "customized guest agent"): self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) self.assertEqual(self.snapshot(), before) with patch("uvm.images._V4_TEMPLATE_AGENT_SHA256", hashlib.sha256(old_payload.encode()).hexdigest()): self.assertTrue(self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True)) self.assertEqual(self.guest("cat /usr/lib/uvm/agent.py"), guest_agent_files()["/usr/lib/uvm/agent.py"].decode()) def test_agent_updates_use_recorded_hashes_but_preserve_guest_customizations(self) -> None: self.make_disk() self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) updated = guest_agent_files() updated["/usr/lib/uvm/agent.py"] += b"\n# next packaged revision\n" with patch("uvm.images.guest_agent_files", return_value=updated): self.images.provision_guest_agent(self.disk, "vm-test", "worker") self.assertTrue(self.guest("cat /usr/lib/uvm/agent.py").endswith("# next packaged revision\n")) self.replace_guest_file("/usr/lib/uvm/agent.py", "guest customization\n") before = self.snapshot() with self.assertRaisesRegex(UvmError, "customized guest agent file"): self.images.provision_guest_agent(self.disk, "vm-test", "worker") self.assertEqual(self.snapshot(), before) self.assertEqual(list(self.disk.parent.iterdir()), [self.disk]) def test_failure_during_provisioning_keeps_the_original_disk(self) -> None: self.make_disk() before = self.snapshot() write = self.images._write_guest_file def fail(candidate, path, source): self.assertNotEqual(candidate, self.disk) write(candidate, path, source) raise UvmError("simulated write failure") with patch.object(self.images, "_write_guest_file", side_effect=fail), self.assertRaisesRegex(UvmError, "write failure"): self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) self.assertEqual(self.snapshot(), before) self.assertEqual(list(self.disk.parent.iterdir()), [self.disk]) def test_guest_edits_revealed_by_journal_replay_are_checked_before_replacement(self) -> None: self.make_disk() self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) before = self.snapshot() run = self.runner.run def replay(command, **options): result = run(command, **options) if command[:3] == ("e2fsck", "-f", "-p"): self.replace_guest_file("/usr/lib/uvm/agent.py", "journal-committed guest edit\n", disk=command[-1]) return result with patch.object(self.runner, "run", side_effect=replay), self.assertRaisesRegex(UvmError, "customized"): self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) self.assertEqual(self.snapshot(), before) def test_unsafe_guest_directory_and_shared_disk_are_rejected(self) -> None: self.make_disk() with self.assertRaisesRegex(UvmError, "unsafe"): self.images.provision_guest_agent(self.settings.rootfs_image, "vm-test", "worker", install=True) self.runner.run(("debugfs", "-w", "-R", "set_inode_field /etc mode 040777", self.disk), capture=True) before = self.snapshot() with self.assertRaisesRegex(UvmError, "unsafe guest agent directory"): self.images.provision_guest_agent(self.disk, "vm-test", "worker", install=True) self.assertEqual(self.snapshot(), before)