From b92000b47f3c36284a01f56015705dfdcac4e3fd Mon Sep 17 00:00:00 2001
From: Jake Correnti <jakecorrenti+github@proton.me>
Date: Mon, 16 Dec 2024 17:17:09 -0500
Subject: [PATCH] Support additional image types

Link with libkrun's new `krun_add_disk2` API to allow the user to use
either a Raw or Qcow2 disk image.

Adds the `format` argument to the `virtio-blk` device.

Signed-off-by: Jake Correnti <jakecorrenti+github@proton.me>
---
 src/cmdline.rs | 11 ++++++++---
 src/virtio.rs  | 37 ++++++++++++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/cmdline.rs b/src/cmdline.rs
index d3a4adb..cb94142 100644
--- a/src/cmdline.rs
+++ b/src/cmdline.rs
@@ -191,13 +191,13 @@ mod tests {
             "--bootloader",
             "efi,variable-store=/Users/user/bootloader,create",
             "--device",
-            "virtio-blk,path=/Users/user/root.raw",
+            "virtio-blk,path=/Users/user/root.qcow2,format=qcow2",
             "--device",
             "virtio-rng",
             "--device",
             "virtio-serial,logFilePath=/Users/user/serial.log",
             "--device",
-            "virtio-blk,path=/Users/user/data.raw",
+            "virtio-blk,path=/Users/user/data.raw,format=raw",
             "--device",
             "virtio-vsock,port=1024,socketURL=/Users/user/vsock1.sock,listen",
             "--device",
@@ -299,6 +299,7 @@ mod tests {
             .expect("expected 4th virtio device config");
         if let VirtioDeviceConfig::Blk(blk) = blk {
             assert_eq!(blk.path, PathBuf::from_str("/Users/user/data.raw").unwrap());
+            assert_eq!(blk.format, KRUN_DISK_FORMAT_RAW);
         } else {
             panic!("expected virtio-blk device as 4th device config argument");
         }
@@ -330,7 +331,11 @@ mod tests {
             .pop()
             .expect("expected 1st virtio device config");
         if let VirtioDeviceConfig::Blk(blk) = blk {
-            assert_eq!(blk.path, PathBuf::from_str("/Users/user/root.raw").unwrap());
+            assert_eq!(
+                blk.path,
+                PathBuf::from_str("/Users/user/root.qcow2").unwrap()
+            );
+            assert_eq!(blk.format, KRUN_DISK_FORMAT_QCOW2);
         } else {
             panic!("expected virtio-blk device as 1st device config argument");
         }
diff --git a/src/virtio.rs b/src/virtio.rs
index ec18f69..bfb6775 100644
--- a/src/virtio.rs
+++ b/src/virtio.rs
@@ -14,10 +14,11 @@ use mac_address::MacAddress;
 
 #[link(name = "krun-efi")]
 extern "C" {
-    fn krun_add_disk(
+    fn krun_add_disk2(
         ctx_id: u32,
         c_block_id: *const c_char,
         c_disk_path: *const c_char,
+        disk_format: u32,
         read_only: bool,
     ) -> i32;
     fn krun_add_vsock_port(ctx_id: u32, port: u32, c_filepath: *const c_char) -> i32;
@@ -27,6 +28,25 @@ extern "C" {
     fn krun_set_console_output(ctx_id: u32, c_filepath: *const c_char) -> i32;
 }
 
+#[repr(u32)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum DiskImageFormat {
+    Raw = 0,
+    Qcow2 = 1,
+}
+
+impl FromStr for DiskImageFormat {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "raw" => Ok(DiskImageFormat::Raw),
+            "qcow2" => Ok(DiskImageFormat::Qcow2),
+            _ => Err(anyhow!("unsupported disk image format")),
+        }
+    }
+}
+
 /// Each virito device configures itself with krun differently. This is used by each virtio device
 /// to set their respective configurations with libkrun.
 pub trait KrunContextSet {
@@ -100,17 +120,21 @@ impl KrunContextSet for VirtioDeviceConfig {
 pub struct BlkConfig {
     /// Path of the file to store as the root disk.
     pub path: PathBuf,
+
+    /// Format of the disk image.
+    pub format: DiskImageFormat,
 }
 
 impl FromStr for BlkConfig {
     type Err = anyhow::Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        let args = args_parse(s.to_string(), "virtio-blk", Some(1))?;
+        let args = args_parse(s.to_string(), "virtio-blk", Some(2))?;
 
         Ok(Self {
             path: PathBuf::from_str(&val_parse(&args[0], "path")?)
                 .context("path argument not a valid path")?,
+            format: DiskImageFormat::from_str(val_parse(&args[1], "format")?.as_str())?,
         })
     }
 }
@@ -125,7 +149,14 @@ impl KrunContextSet for BlkConfig {
         let block_id_cstr = CString::new(basename).context("can't convert basename to cstring")?;
         let path_cstr = path_to_cstring(&self.path)?;
 
-        if krun_add_disk(id, block_id_cstr.as_ptr(), path_cstr.as_ptr(), false) < 0 {
+        if krun_add_disk2(
+            id,
+            block_id_cstr.as_ptr(),
+            path_cstr.as_ptr(),
+            self.format as u32,
+            false,
+        ) < 0
+        {
             return Err(anyhow!(format!(
                 "unable to set virtio-blk disk for {}",
                 self.path.display()