feat: Add storage management for VM creation and resizing

- Introduced storage parameter in VmCreateRequest and VmResizeRequest models.
- Implemented storage validation to ensure minimum and maximum limits.
- Enhanced CLI to support storage options during VM creation and resizing.
- Updated lifecycle management to handle pending resize operations for VM resources.
- Added methods for resizing VM disks and validating storage capacity.
- Updated API endpoints to accommodate new storage functionality.
- Refactored related validation and configuration logic to support storage management.

uvm/images.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
+import os
 import re
-import shutil
+import stat
 import tempfile
 from dataclasses import dataclass
 from pathlib import Path
 
-from .config import Settings
-from .errors import UvmError
+from .config import MAX_STORAGE_MIB, MIN_STORAGE_MIB, Settings
+from .errors import UvmError, ValidationError
 from .integrity import load_manifest, verify_file
 from .system import CommandRunner
 
@@ -88,23 +89,92 @@ class ImageStore:
                     allow_unverified=self._settings.allow_unverified_downloads)
         return kernel
 
-    def create_vm_disk(self, source: Path, destination: Path) -> Path:
+    def create_vm_disk(
+        self, source: Path, destination: Path, *, storage_mib: int = MIN_STORAGE_MIB,
+    ) -> Path:
         """Copy the template before Firecracker opens it read-write for a VM."""
 
+        self._validate_storage(storage_mib)
         if destination.exists():
             raise UvmError(f"VM disk already exists: {destination}")
         destination.parent.mkdir(parents=True, exist_ok=True)
         try:
-            shutil.copy2(source, destination)
+            self._copy_sparse(source, destination)
             destination.chmod(0o600)
-        except OSError as error:
+            self._resize_filesystem(destination, storage_mib)
+        except BaseException as error:
             try:
                 destination.unlink()
             except FileNotFoundError:
                 pass
-            raise UvmError(f"could not create VM disk {destination}: {error}") from error
+            if isinstance(error, OSError):
+                raise UvmError(f"could not create VM disk {destination}: {error}") from error
+            raise
         return destination
 
+    def resize_vm_disk(self, disk: Path, storage_mib: int) -> None:
+        """Resize a proven-offline private disk on a copy, then publish atomically."""
+        self._validate_storage(storage_mib)
+        metadata = disk.lstat()
+        if (not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1
+                or disk.resolve() == self._settings.rootfs_image.resolve()):
+            raise UvmError(f"private VM disk is unsafe to resize: {disk}")
+        target_bytes = storage_mib * 1024 * 1024
+        if metadata.st_size == target_bytes and self._filesystem_size(disk) == target_bytes:
+            return  # A previous resize may have published its disk before saving metadata.
+
+        descriptor, name = tempfile.mkstemp(prefix=".rootfs-resize-", suffix=".ext4", dir=disk.parent)
+        os.close(descriptor)
+        temporary = Path(name)
+        try:
+            self._copy_sparse(disk, temporary)
+            temporary.chmod(0o600)
+            self._resize_filesystem(temporary, storage_mib)
+            os.replace(temporary, disk)
+            directory = os.open(disk.parent, os.O_RDONLY)
+            try:
+                os.fsync(directory)
+            finally:
+                os.close(directory)
+        except OSError as error:
+            raise UvmError(f"could not resize VM disk {disk}: {error}") from error
+        finally:
+            temporary.unlink(missing_ok=True)
+
+    @staticmethod
+    def _validate_storage(storage_mib: int) -> None:
+        if type(storage_mib) is not int or not MIN_STORAGE_MIB <= storage_mib <= MAX_STORAGE_MIB:
+            raise ValidationError("storage must be a supported whole-MiB capacity of at least 30 GiB")
+
+    def _copy_sparse(self, source: Path, destination: Path) -> None:
+        self._runner.run(("cp", "--reflink=auto", "--sparse=always", "--", source, destination), capture=True)
+
+    def _filesystem_size(self, disk: Path) -> int:
+        result = self._runner.run(("env", "LC_ALL=C", "dumpe2fs", "-h", disk), capture=True)
+        blocks = re.search(r"^Block count:\s+(\d+)\s*$", result.stdout, re.MULTILINE)
+        block_size = re.search(r"^Block size:\s+(\d+)\s*$", result.stdout, re.MULTILINE)
+        if blocks is None or block_size is None:
+            raise UvmError(f"could not determine ext4 filesystem capacity: {disk}")
+        return int(blocks.group(1)) * int(block_size.group(1))
+
+    def _resize_filesystem(self, disk: Path, storage_mib: int) -> None:
+        """Grow the file before ext4; shrink ext4 before truncating the file."""
+        checked = self._runner.run(("e2fsck", "-f", "-p", disk), check=False, capture=True)
+        if checked.returncode not in (0, 1):
+            detail = checked.stderr.strip() or checked.stdout.strip()
+            raise UvmError(f"filesystem check failed before resize ({checked.returncode}): {detail}")
+        target_bytes = storage_mib * 1024 * 1024
+        if disk.stat().st_size < target_bytes:
+            with disk.open("r+b") as image:
+                image.truncate(target_bytes)
+        self._runner.run(("resize2fs", disk, f"{storage_mib}M"), capture=True)
+        if self._filesystem_size(disk) != target_bytes:
+            raise UvmError("resized ext4 filesystem does not match the requested capacity")
+        with disk.open("r+b") as image:
+            image.truncate(target_bytes)
+            image.flush()
+            os.fsync(image.fileno())
+
     def provision_credentials(self, disk: Path, username: str, password: str) -> None:
         """Set an existing guest account password in an offline ext4 disk."""