resize safety constraints

Author: evakhoni

packages/jumpstarter-driver-qemu/jumpstarter_driver_qemu/driver.py

@@ -4,6 +4,7 @@
 import logging
 import os
 import platform
+import shutil
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from functools import cached_property
@@ -20,7 +21,7 @@
 from jumpstarter_driver_opendal.driver import FlasherInterface
 from jumpstarter_driver_power.driver import PowerInterface, PowerReading
 from jumpstarter_driver_pyserial.driver import PySerial
-from pydantic import BaseModel, Field, validate_call
+from pydantic import BaseModel, ByteSize, Field, TypeAdapter, validate_call
 from qemu.qmp import QMPClient
 from qemu.qmp.protocol import ConnectError, Runstate
 
@@ -153,19 +154,6 @@ async def on(self) -> None:  # noqa: C901
             ]
 
         if root.exists():
-            # Resize disk if configured
-            if self.parent.disk_size:
-                self.logger.info(f"Resizing disk to {self.parent.disk_size}")
-                resize_proc = await run_process(
-                    ["qemu-img", "resize", str(root), self.parent.disk_size],
-                    stdout=PIPE,
-                    stderr=PIPE,
-                )
-                try:
-                    resize_proc.check_returncode()
-                except CalledProcessError as e:
-                    raise RuntimeError(f"Failed to resize disk: {resize_proc.stderr.decode()}") from e
-
             proc = await run_process(
                 ["qemu-img", "info", "--output=json", str(root)],
                 stdout=PIPE,
@@ -175,6 +163,7 @@ async def on(self) -> None:  # noqa: C901
                 proc.check_returncode()
                 info = json.loads(proc.stdout)
                 image_format = info.get("format", "raw")
+                current_virtual_size = info.get("virtual-size") or root.stat().st_size
                 match image_format:
                     case "raw" | "qcow2" | "qcow" | "vmdk":
                         image_driver = image_format
@@ -183,6 +172,36 @@ async def on(self) -> None:  # noqa: C901
             except CalledProcessError:
                 self.logger.warning("unable to detect image format, assuming raw")
                 image_driver = "raw"
+                current_virtual_size = root.stat().st_size
+
+            # Resize disk if configured
+            if self.parent.disk_size:
+                # Convert QEMU binary format (20G) to Pydantic format (20GiB)
+                s = self.parent.disk_size
+                requested = int(TypeAdapter(ByteSize).validate_python(s + "iB" if s[-1] in "kmgtKMGT" else s))
+
+                if requested < current_virtual_size:
+                    raise RuntimeError(
+                        f"Shrinking disk is not supported: current {ByteSize(current_virtual_size).human_readable()}, "
+                        f"requested {self.parent.disk_size}"
+                    )
+
+                available = shutil.disk_usage(root.parent).free
+                if requested > available:
+                    raise RuntimeError(
+                        f"Not enough disk space: need {ByteSize(requested).human_readable()}, "
+                        f"only {ByteSize(available).human_readable()} available"
+                    )
+
+                if requested > current_virtual_size:
+                    self.logger.info(f"Resizing disk to {ByteSize(requested).human_readable()}")
+                    proc = await run_process(
+                        ["qemu-img", "resize", str(root), str(requested)],
+                        stdout=PIPE,
+                        stderr=PIPE,
+                    )
+                    if proc.returncode != 0:
+                        raise RuntimeError(f"Failed to resize disk: {proc.stderr.decode()}")
 
             cmdline += [
                 "-blockdev",