from __future__ import annotations import hashlib import shutil import tempfile import unittest from pathlib import Path from unittest.mock import Mock, patch from uvm.config import MIN_STORAGE_MIB, Settings from uvm.errors import UvmError, ValidationError from uvm.images import ImageStore from uvm.system import CommandResult, CommandRunner MIB = 1024 * 1024 class DiskResizeTests(unittest.TestCase): def test_grow_extends_the_file_first_and_shrink_truncates_only_after_verified_resize(self) -> None: for original_mib, target_mib in ((8, 16), (32, 16)): with self.subTest(original=original_mib), tempfile.TemporaryDirectory() as temporary: disk = Path(temporary) / "rootfs.ext4" with disk.open("wb") as image: image.truncate(original_mib * MIB) sizes = [] def run(command, **_options): sizes.append((command[0], disk.stat().st_size)) return CommandResult(args=tuple(map(str, command)), returncode=1 if command[0] == "e2fsck" else 0) runner = Mock(spec=CommandRunner) runner.run.side_effect = run store = ImageStore(Settings(), runner) with patch.object(store, "_filesystem_size", return_value=target_mib * MIB): store._resize_filesystem(disk, target_mib) self.assertEqual(sizes, [("e2fsck", original_mib * MIB), ("resize2fs", max(original_mib, target_mib) * MIB)]) self.assertEqual(disk.stat().st_size, target_mib * MIB) def test_failed_or_interrupted_staging_never_changes_the_original_disk(self) -> None: for failure in (UvmError("filesystem cannot fit"), KeyboardInterrupt()): with self.subTest(failure=type(failure)), tempfile.TemporaryDirectory() as temporary: disk = Path(temporary) / "rootfs.ext4" disk.write_bytes(b"data, password hash, and SSH host keys") disk.chmod(0o600) before = (disk.stat().st_ino, disk.read_bytes()) store = ImageStore(Settings(base=Path(temporary)), Mock(spec=CommandRunner)) def fail(candidate, _capacity): self.assertNotEqual(candidate, disk) candidate.write_bytes(b"partly resized staging image") raise failure with ( patch.object(store, "_copy_sparse", side_effect=shutil.copyfile), patch.object(store, "_resize_filesystem", side_effect=fail), self.assertRaises(type(failure)), ): store.resize_vm_disk(disk, MIN_STORAGE_MIB) self.assertEqual((disk.stat().st_ino, disk.read_bytes()), before) self.assertEqual(list(disk.parent.iterdir()), [disk]) def test_filesystem_check_errors_prevent_resize_and_truncation(self) -> None: with tempfile.TemporaryDirectory() as temporary: disk = Path(temporary) / "rootfs.ext4" disk.write_bytes(b"original") runner = Mock(spec=CommandRunner) runner.run.return_value = CommandResult(("e2fsck",), 4, stderr="unrepaired errors") store = ImageStore(Settings(), runner) with self.assertRaisesRegex(UvmError, "filesystem check failed"): store._resize_filesystem(disk, 16) self.assertEqual(disk.read_bytes(), b"original") self.assertEqual(runner.run.call_count, 1) def test_geometry_mismatch_prevents_truncating_a_larger_filesystem(self) -> None: with tempfile.TemporaryDirectory() as temporary: disk = Path(temporary) / "rootfs.ext4" with disk.open("wb") as image: image.truncate(32 * MIB) runner = Mock(spec=CommandRunner) runner.run.return_value = CommandResult(("tool",), 0) store = ImageStore(Settings(), runner) with ( patch.object(store, "_filesystem_size", return_value=32 * MIB), self.assertRaisesRegex(UvmError, "does not match"), ): store._resize_filesystem(disk, 16) self.assertEqual(disk.stat().st_size, 32 * MIB) def test_retry_after_disk_publication_recognizes_the_completed_capacity(self) -> None: with tempfile.TemporaryDirectory() as temporary: disk = Path(temporary) / "rootfs.ext4" with disk.open("wb") as image: image.truncate(MIN_STORAGE_MIB * MIB) runner = Mock(spec=CommandRunner) store = ImageStore(Settings(), runner) with patch.object(store, "_filesystem_size", return_value=MIN_STORAGE_MIB * MIB): store.resize_vm_disk(disk, MIN_STORAGE_MIB) runner.run.assert_not_called() def test_resize_rejects_subminimum_capacity_and_shared_or_linked_disks(self) -> None: with tempfile.TemporaryDirectory() as temporary: settings = Settings(base=Path(temporary)) settings.images_dir.mkdir() settings.rootfs_image.write_bytes(b"template") disk = Path(temporary) / "rootfs.ext4" disk.write_bytes(b"private disk") linked = Path(temporary) / "linked.ext4" linked.hardlink_to(disk) symlink = Path(temporary) / "symlink.ext4" symlink.symlink_to(disk) runner = Mock(spec=CommandRunner) store = ImageStore(settings, runner) with self.assertRaises(ValidationError): store.resize_vm_disk(disk, MIN_STORAGE_MIB - 1) for target in (settings.rootfs_image, linked, symlink): with self.subTest(target=target), self.assertRaisesRegex(UvmError, "unsafe"): store.resize_vm_disk(target, MIN_STORAGE_MIB) runner.run.assert_not_called() @unittest.skipUnless(all(shutil.which(tool) for tool in ("mkfs.ext4", "e2fsck", "resize2fs", "dumpe2fs", "debugfs", "cp")), "e2fsprogs and coreutils are required for offline disk checks") class OfflineExt4ResizeTests(unittest.TestCase): def test_real_sparse_disk_grows_and_shrinks_with_guest_data_and_keys_intact(self) -> None: with tempfile.TemporaryDirectory() as temporary: base = Path(temporary) source = base / "template.ext4" root = base / "guest" root.mkdir() contents = { "data": "retained application data\n", "shadow": "root:$6$unchanged-password-hash:1:0:99999:7:::\n", "host-key": "unchanged-private-host-key\n", } for name, value in contents.items(): (root / name).write_text(value, encoding="utf-8") with source.open("wb") as image: image.truncate(32 * MIB) runner = CommandRunner(emit=None) runner.run(("mkfs.ext4", "-q", "-F", "-b", "4096", "-N", "128", "-d", root, source), capture=True) template_digest = hashlib.sha256(source.read_bytes()).hexdigest() store = ImageStore(Settings(base=base), runner) disk = store.create_vm_disk(source, base / "rootfs.ext4") for capacity in (MIN_STORAGE_MIB, 31 * 1024, MIN_STORAGE_MIB): store.resize_vm_disk(disk, capacity) self.assertEqual(disk.stat().st_size, capacity * MIB) self.assertEqual(store._filesystem_size(disk), capacity * MIB) self.assertEqual(disk.stat().st_mode & 0o777, 0o600) self.assertLess(disk.stat().st_blocks * 512, disk.stat().st_size // 2) runner.run(("e2fsck", "-f", "-n", disk), capture=True) for name, value in contents.items(): result = runner.run(("debugfs", "-R", f"cat /{name}", disk), capture=True) self.assertEqual(result.stdout, value) self.assertEqual(hashlib.sha256(source.read_bytes()).hexdigest(), template_digest)