from __future__ import annotations import hashlib import io import os import shutil import tarfile import tempfile import unittest from dataclasses import replace from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, patch from uvm.config import Settings from uvm.errors import CommandError, UvmError from uvm.images import ImageStore from uvm.install import Installer, _ROOTFS_RECIPE from uvm.integrity import load_manifest, write_manifest from uvm.state import StateStore from uvm.system import CommandResult, CommandRunner class InstallerArchitectureTests(unittest.TestCase): def test_installs_matching_native_artifacts_with_a_verified_manifest(self) -> None: for machine, architecture in ( ("x86_64", "x86_64"), ("aarch64", "aarch64"), ("arm64", "aarch64"), ): with ( self.subTest(machine=machine), tempfile.TemporaryDirectory() as temporary_directory, patch.dict(os.environ, {"UVM_BASE": temporary_directory}, clear=True), patch("uvm.config.os.uname", return_value=SimpleNamespace(machine=machine)), ): settings = Settings.from_environment() version = settings.firecracker_version binaries = { "firecracker": f"firecracker for {architecture}".encode(), "jailer": f"jailer for {architecture}".encode(), } archive_buffer = io.BytesIO() with tarfile.open(fileobj=archive_buffer, mode="w:gz") as archive: for name, content in binaries.items(): member = tarfile.TarInfo( f"release-{version}-{architecture}/{name}-{version}-{architecture}" ) member.size = len(content) member.mode = 0o755 archive.addfile(member, io.BytesIO(content)) archive_bytes = archive_buffer.getvalue() kernel_bytes = f"kernel for {architecture}".encode() rootfs_bytes = f"rootfs for {architecture}".encode() settings = replace( settings, firecracker_sha256=hashlib.sha256(archive_bytes).hexdigest(), kernel_sha256=hashlib.sha256(kernel_bytes).hexdigest(), rootfs_sha256=hashlib.sha256(rootfs_bytes).hexdigest(), rootfs_url=( "https://s3.amazonaws.com/spec.ccfc.min/img/quickstart_guide/" f"{architecture}/rootfs/bionic.rootfs.ext4" ), ) guest_base = ( "https://s3.amazonaws.com/spec.ccfc.min/img/quickstart_guide/" f"{architecture}" ) payloads = { "https://github.com/firecracker-microvm/firecracker/releases/download/" f"{version}/firecracker-{version}-{architecture}.tgz": archive_bytes, f"{guest_base}/kernels/vmlinux.bin": kernel_bytes, f"{guest_base}/rootfs/bionic.rootfs.ext4": rootfs_bytes, } runner = Mock(spec=CommandRunner) installer = Installer(settings, runner, StateStore(settings), emit=None) with ( patch("uvm.install.require_root"), patch("uvm.install.check_kvm"), patch( "uvm.install.urllib.request.urlopen", side_effect=lambda url, timeout: io.BytesIO(payloads[url]), ) as download, ): firecracker, assets = installer.install() self.assertEqual(firecracker.read_bytes(), binaries["firecracker"]) self.assertEqual(settings.jailer_binary.read_bytes(), binaries["jailer"]) self.assertEqual(assets.kernel.read_bytes(), kernel_bytes) self.assertEqual(assets.rootfs.read_bytes(), rootfs_bytes) self.assertEqual(download.call_count, 3) self.assertEqual({call.args[0] for call in download.call_args_list}, set(payloads)) manifest = load_manifest(settings.integrity_manifest_path) self.assertTrue(manifest.verified) self.assertEqual(manifest.checksums, { name: hashlib.sha256(content).hexdigest() for name, content in { **binaries, "kernel": kernel_bytes, "rootfs": rootfs_bytes, }.items() }) def test_arm64_install_requires_checksums_before_host_commands_or_downloads(self) -> None: with tempfile.TemporaryDirectory() as temporary_directory: settings = Settings(base=Path(temporary_directory), architecture="aarch64") runner = Mock(spec=CommandRunner) installer = Installer(settings, runner, StateStore(settings), emit=None) with ( patch("uvm.install.require_root"), patch("uvm.install.check_kvm") as check_kvm, patch("uvm.install.urllib.request.urlopen") as download, self.assertRaisesRegex(UvmError, "Firecracker archive checksum is required"), ): installer.install() runner.run.assert_not_called() check_kvm.assert_not_called() download.assert_not_called() self.assertFalse(settings.integrity_manifest_path.exists()) def test_default_build_records_the_authenticated_local_image_digest(self) -> None: with tempfile.TemporaryDirectory() as temporary: settings = Settings( base=Path(temporary), firecracker_sha256="1" * 64, kernel_sha256=hashlib.sha256(b"kernel").hexdigest(), ) settings.bin_dir.mkdir() settings.firecracker_binary.write_bytes(b"firecracker") settings.jailer_binary.write_bytes(b"jailer") installer = Installer(settings, Mock(spec=CommandRunner), StateStore(settings), emit=None) with ( patch("uvm.install.require_root"), patch("uvm.install.check_kvm"), patch.object(installer, "install_apt_packages"), patch.object(installer, "install_firecracker_binary", return_value=settings.firecracker_binary), patch.object(installer, "download", side_effect=lambda _url, path, **_kwargs: path.write_bytes(b"kernel")), patch.object(installer, "_build_ubuntu_rootfs", side_effect=lambda path: path.write_bytes(b"full Ubuntu")) as build, ): installer.install() build.assert_called_once_with(settings.rootfs_image) manifest = load_manifest(settings.integrity_manifest_path) self.assertTrue(manifest.verified) self.assertEqual(manifest.checksums["rootfs"], hashlib.sha256(b"full Ubuntu").hexdigest()) self.assertEqual(ImageStore(settings).installed_assets().rootfs, settings.rootfs_image) def test_old_rootfs_checksum_requires_an_explicit_prebuilt_url_before_mutations(self) -> None: settings = Settings(firecracker_sha256="1" * 64, kernel_sha256="2" * 64, rootfs_sha256="3" * 64) runner = Mock(spec=CommandRunner) installer = Installer(settings, runner, Mock(spec=StateStore), emit=None) with self.assertRaisesRegex(UvmError, "unset UVM_ROOTFS_SHA256"): installer._validate_integrity_policy() runner.run.assert_not_called() def test_only_verified_current_recipe_images_are_reused(self) -> None: for verified, marker, force in ( (True, _ROOTFS_RECIPE, False), (False, _ROOTFS_RECIPE, False), (True, "", False), (True, _ROOTFS_RECIPE, True), ): with self.subTest(verified=verified, marker=marker, force=force), tempfile.TemporaryDirectory() as temporary: settings = Settings(base=Path(temporary)) settings.images_dir.mkdir() settings.rootfs_image.write_bytes(b"image") write_manifest(settings.integrity_manifest_path, {"rootfs": hashlib.sha256(b"image").hexdigest()}, verified=verified) runner = Mock(spec=CommandRunner) runner.run.return_value = CommandResult(("debugfs",), 0, stdout=marker) installer = Installer(settings, runner, StateStore(settings), emit=None) with patch.object(installer, "_build_ubuntu_rootfs") as build: installer._install_ubuntu_rootfs(settings.rootfs_image, force=force) self.assertEqual(build.called, not (verified and marker == _ROOTFS_RECIPE and not force)) class UbuntuBuildRunner: """Simulate package installation; exercise real ext4 formatting and readback.""" def __init__(self, *, failure: str | None = None) -> None: self.calls: list[tuple[str, ...]] = [] self.failure = failure self.runner = CommandRunner(emit=None) def run(self, command, **options) -> CommandResult: args = tuple(map(str, command)) self.calls.append(args) if "debootstrap" in args: root = Path(args[-2]) if self.failure == "bootstrap": raise CommandError("bootstrap interrupted") status = "\n\n".join( f"Package: {package}\nStatus: install ok installed\nVersion: 1.0" for package in ("apt", "sudo", "openssh-server", "systemd-sysv") ) + "\n" files = { "etc/passwd": "root:x:0:0::/root:/bin/bash\n", "etc/shadow": "root:*:1:0:99999:7:::\n", "etc/group": "root:x:0:\n", "etc/sudoers": "root ALL=(ALL:ALL) ALL\n", "etc/ssh/sshd_config": "UsePAM yes\n", "etc/ssh/ssh_host_rsa_key": "template host key", "usr/bin/sudo": "sudo binary", "usr/bin/apt-get": "apt binary", "usr/sbin/sshd": "sshd binary", "usr/lib/systemd/systemd": "init binary", "var/lib/dpkg/status": status, "var/lib/dpkg/info/apt.list": "/usr/bin/apt-get\n", "var/lib/dpkg/info/sudo.list": "/usr/bin/sudo\n", } if self.failure == "metadata": files["var/lib/dpkg/status"] = "" for relative, contents in files.items(): path = root / relative path.parent.mkdir(parents=True, exist_ok=True) path.write_text(contents, encoding="utf-8") (root / "usr/bin/sudo").chmod(0o4755) (root / "etc/resolv.conf").symlink_to("../run/systemd/resolve/stub-resolv.conf") return CommandResult(args, 0) if args[0] == "systemctl": return CommandResult(args, 0) return self.runner.run(command, **options) @unittest.skipUnless(all(shutil.which(tool) for tool in ("mkfs.ext4", "e2fsck", "debugfs")), "e2fsprogs is required for offline Ubuntu image checks") class UbuntuRootfsBuildTests(unittest.TestCase): def test_build_keeps_package_metadata_and_sudo_permissions_with_native_signed_bootstrap(self) -> None: for architecture, deb_arch, mirror in ( ("aarch64", "arm64", "https://ports.ubuntu.com/ubuntu-ports"), ("x86_64", "amd64", "https://archive.ubuntu.com/ubuntu"), ): with self.subTest(architecture=architecture), tempfile.TemporaryDirectory() as temporary: settings = Settings(base=Path(temporary), architecture=architecture) runner = UbuntuBuildRunner() installer = Installer(settings, runner, StateStore(settings), emit=None) # type: ignore[arg-type] installer._build_ubuntu_rootfs(settings.rootfs_image) bootstrap = next(call for call in runner.calls if "debootstrap" in call) self.assertIn(f"--arch={deb_arch}", bootstrap) self.assertEqual(bootstrap[-1], mirror) self.assertIn("--force-check-gpg", bootstrap) self.assertIn("--extra-suites=jammy-updates,jammy-security", bootstrap) self.assertIn("--mount-proc", bootstrap) self.assertIn("--kill-child", bootstrap) self.assertEqual(settings.rootfs_image.stat().st_mode & 0o777, 0o600) self.assertEqual(list(settings.images_dir.iterdir()), [settings.rootfs_image]) def guest(command): return runner.runner.run(("debugfs", "-R", command, settings.rootfs_image), capture=True) self.assertIn("Package: apt\nStatus: install ok installed", guest("cat /var/lib/dpkg/status").stdout) self.assertEqual(guest("cat /var/lib/dpkg/info/sudo.list").stdout, "/usr/bin/sudo\n") self.assertRegex(guest("stat /usr/bin/sudo").stdout, r"Mode:\s+0?4755") self.assertEqual(guest("cat /etc/uvm-image").stdout, _ROOTFS_RECIPE) self.assertIn("nameserver 1.1.1.1", guest("cat /etc/resolv.conf").stdout) self.assertIn("KeepConfiguration=static", guest("cat /etc/systemd/network/10-uvm.network").stdout) self.assertIn("jammy-security main universe", guest("cat /etc/apt/sources.list").stdout) for path in ("/usr/sbin/policy-rc.d", "/etc/ssh/ssh_host_rsa_key"): self.assertIn("File not found", guest(f"stat {path}").stderr) def test_failed_bootstrap_or_missing_package_metadata_preserves_the_old_template(self) -> None: for failure in ("bootstrap", "metadata"): with self.subTest(failure=failure), tempfile.TemporaryDirectory() as temporary: settings = Settings(base=Path(temporary)) settings.images_dir.mkdir() settings.rootfs_image.write_bytes(b"old template") runner = UbuntuBuildRunner(failure=failure) installer = Installer(settings, runner, StateStore(settings), emit=None) # type: ignore[arg-type] with self.assertRaises(UvmError): installer._build_ubuntu_rootfs(settings.rootfs_image) self.assertEqual(settings.rootfs_image.read_bytes(), b"old template") self.assertEqual(list(settings.images_dir.iterdir()), [settings.rootfs_image]) self.assertFalse(any(call[0] == "mkfs.ext4" for call in runner.calls))