fix(block): punch holes for freed qcow2 clusters

src/devices/Cargo.toml

@@ -41,7 +41,7 @@ arch = { package = "krun-arch", version = "=0.1.0-1.18.0", path = "../arch" }
 utils = { package = "krun-utils", version = "=0.1.0-1.18.0", path = "../utils" }
 polly = { package = "krun-polly", version = "=0.1.0-1.18.0", path = "../polly" }
 rutabaga_gfx = { package = "krun-rutabaga-gfx", version = "=0.1.0-1.18.0", path = "../rutabaga_gfx", features = ["virgl_renderer", "virgl_renderer_next"], optional = true }
-imago = { version = "0.2.2", features = ["sync-wrappers", "vm-memory"] }
+imago = { git = "https://github.com/ltstriker/imago", rev = "0ff6430e7ee05c9920dd84809388935a1a340396", features = ["sync-wrappers", "vm-memory"] }
 
 [target.'cfg(target_os = "macos")'.dependencies]
 hvf = { package = "krun-hvf", version = "=0.1.0-1.18.0", path = "../hvf" }

src/devices/src/virtio/block/device.rs

@@ -253,7 +253,7 @@ impl Block {
         #[cfg(target_os = "macos")]
         let file_opts = file_opts.relaxed_sync(sync_mode == SyncMode::Relaxed);
         let file = ImagoFile::open_sync(file_opts)?;
-        let discard_alignment = file.discard_align();
+        let mut discard_alignment = file.discard_align();
 
         let disk_image = match disk_image_format {
             ImageType::Qcow2 => {
@@ -263,6 +263,8 @@ impl Block {
                         !is_disk_read_only,
                     )?;
                 qcow2.open_implicit_dependencies_sync()?;
+                // qcow2 only frees whole clusters; smaller discards are dropped.
+                discard_alignment = discard_alignment.max(qcow2.cluster_size());
                 SyncFormatAccess::new(qcow2)?
             }
             ImageType::Raw => {
@@ -288,8 +290,6 @@ impl Block {
 
         let mut avail_features = (1u64 << VIRTIO_F_VERSION_1)
             | (1u64 << VIRTIO_BLK_F_SEG_MAX)
-            | (1u64 << VIRTIO_BLK_F_DISCARD)
-            | (1u64 << VIRTIO_BLK_F_WRITE_ZEROES)
             | (1u64 << VIRTIO_RING_F_EVENT_IDX);
 
         if sync_mode != SyncMode::None {
@@ -298,7 +298,9 @@ impl Block {
 
         if is_disk_read_only {
             avail_features |= 1u64 << VIRTIO_BLK_F_RO;
-        };
+        } else {
+            avail_features |= (1u64 << VIRTIO_BLK_F_DISCARD) | (1u64 << VIRTIO_BLK_F_WRITE_ZEROES);
+        }
 
         let config = VirtioBlkConfig {
             capacity: disk_properties.nsectors(),

src/devices/src/virtio/block/worker.rs

@@ -20,12 +20,30 @@ pub enum RequestError {
     DiscardingToZero(io::Error),
     FlushingToDisk(io::Error),
     InvalidDataLength,
+    InvalidRange,
     ReadingFromDescriptor(io::Error),
     WritingToDescriptor(io::Error),
     WritingZeroes(io::Error),
     UnknownRequest,
 }
 
+impl RequestError {
+    /// virtio-blk status reported to the guest for this error.
+    fn status(&self) -> u32 {
+        match self {
+            RequestError::UnknownRequest => VIRTIO_BLK_S_UNSUPP,
+            RequestError::Discarding(e)
+            | RequestError::DiscardingToZero(e)
+            | RequestError::WritingZeroes(e)
+                if e.kind() == io::ErrorKind::Unsupported =>
+            {
+                VIRTIO_BLK_S_UNSUPP
+            }
+            _ => VIRTIO_BLK_S_IOERR,
+        }
+    }
+}
+
 /// The request header represents the mandatory fields of each block device request.
 ///
 /// A request header contains the following fields:
@@ -187,7 +205,7 @@ impl BlockWorker {
                     Ok(l) => (VIRTIO_BLK_S_OK.try_into().unwrap(), l),
                     Err(e) => {
                         error!("error processing request: {e:?}");
-                        (VIRTIO_BLK_S_IOERR.try_into().unwrap(), 0)
+                        (e.status().try_into().unwrap(), 0)
                     }
                 };
 
@@ -259,50 +277,63 @@ impl BlockWorker {
                     Ok(disk_id.len())
                 }
             }
+            // Discards may punch holes into the image file, which is only safe because this
+            // worker handles all requests of the disk one at a time (no I/O can be in flight
+            // on the discarded range).  Revisit this before adding more queues or workers.
             VIRTIO_BLK_T_DISCARD => {
-                let discard_write_data: DiscardWriteData = reader
-                    .read_obj()
-                    .map_err(RequestError::ReadingFromDescriptor)?;
+                let (offset, length) = self.read_discard_range(reader)?.0;
                 self.disk
                     .file
                     .lock()
                     .unwrap()
-                    .discard_to_any(
-                        discard_write_data.sector * 512,
-                        discard_write_data.num_sectors as u64 * 512,
-                    )
+                    .discard_to_any(offset, length)
                     .map_err(RequestError::Discarding)?;
                 Ok(0)
             }
             VIRTIO_BLK_T_WRITE_ZEROES => {
-                let discard_write_data: DiscardWriteData = reader
-                    .read_obj()
-                    .map_err(RequestError::ReadingFromDescriptor)?;
-                let unmap = (discard_write_data.flags & VIRTIO_BLK_WRITE_ZEROES_FLAG_UNMAP) != 0;
+                let ((offset, length), flags) = self.read_discard_range(reader)?;
+                let unmap = (flags & VIRTIO_BLK_WRITE_ZEROES_FLAG_UNMAP) != 0;
                 if unmap {
                     self.disk
                         .file
                         .lock()
                         .unwrap()
-                        .discard_to_zero(
-                            discard_write_data.sector * 512,
-                            discard_write_data.num_sectors as u64 * 512,
-                        )
+                        .discard_to_zero(offset, length)
                         .map_err(RequestError::DiscardingToZero)?;
                 } else {
                     self.disk
                         .file
                         .lock()
                         .unwrap()
-                        .write_zeroes(
-                            discard_write_data.sector * 512,
-                            discard_write_data.num_sectors as u64 * 512,
-                        )
+                        .write_zeroes(offset, length)
                         .map_err(RequestError::WritingZeroes)?;
                 }
                 Ok(0)
             }
             _ => Err(RequestError::UnknownRequest),
         }
     }
+
+    /// Read a DISCARD / WRITE_ZEROES segment, returning its byte range and flags.
+    ///
+    /// Fails if the range does not lie entirely within the disk.
+    fn read_discard_range(
+        &self,
+        reader: &mut Reader,
+    ) -> result::Result<((u64, u64), u32), RequestError> {
+        let data: DiscardWriteData = reader
+            .read_obj()
+            .map_err(RequestError::ReadingFromDescriptor)?;
+        let end_sector = data
+            .sector
+            .checked_add(data.num_sectors as u64)
+            .ok_or(RequestError::InvalidRange)?;
+        if end_sector > self.disk.nsectors() {
+            return Err(RequestError::InvalidRange);
+        }
+        Ok((
+            (data.sector * 512, data.num_sectors as u64 * 512),
+            data.flags,
+        ))
+    }
 }