block: AlignedOperation owns its bounce buffer via Drop

The bounce buffer for an unaligned descriptor was allocated in
execute_async and leaked on error paths, even though, for the sync case
the kernel already had a pointer to the buffer.

Clean this up by moving ownership of the buffer to the AlignedOperation
type. To make it actually safe, stop stashing a guest memory pointer for
the duration of the op. Instead, save the guest address and pass guest
memory back to the complete function.

Signed-off-by: Dylan Reid <dgreid@fb.com>

block/src/aligned_operation.rs

@@ -0,0 +1,90 @@
+// Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
+
+use std::alloc::{Layout, alloc_zeroed, dealloc};
+use std::io;
+
+use vm_memory::GuestAddress;
+
+/// Owns an aligned bounce buffer used when a guest descriptor's host VA
+/// does not meet the disk backend's alignment requirement.
+#[derive(Debug)]
+pub struct AlignedOperation {
+    data_addr: GuestAddress,
+    aligned_ptr: *mut u8,
+    size: usize,
+    layout: Layout,
+}
+
+impl AlignedOperation {
+    /// Allocate a zero-initialized buffer of `size` bytes aligned to
+    /// `alignment`. Returns `InvalidInput` if `size` is zero;
+    /// `alignment` must be a power of two and not exceed `isize::MAX`
+    /// after rounding up.
+    pub fn new(data_addr: GuestAddress, size: usize, alignment: usize) -> io::Result<Self> {
+        if size == 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "AlignedOperation requires a non-zero size",
+            ));
+        }
+        let layout = Layout::from_size_align(size, alignment)
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
+        // SAFETY: size is non-zero (checked above) and Layout::from_size_align
+        // rejects alignments that are not a power of two or that overflow.
+        let aligned_ptr = unsafe { alloc_zeroed(layout) };
+        if aligned_ptr.is_null() {
+            return Err(io::Error::last_os_error());
+        }
+        Ok(Self {
+            data_addr,
+            aligned_ptr,
+            size,
+            layout,
+        })
+    }
+
+    /// Gets the raw pointer to the aligned buffer.
+    pub fn as_mut_ptr(&mut self) -> *mut u8 {
+        self.aligned_ptr
+    }
+
+    /// Returns the aligned buffer as a slice.
+    pub fn as_bytes(&self) -> &[u8] {
+        // SAFETY: `new` allocates `size` bytes via alloc_zeroed (so they
+        // are initialized) and AlignedOperation owns the buffer
+        // exclusively.
+        unsafe { std::slice::from_raw_parts(self.aligned_ptr, self.size) }
+    }
+
+    /// Returns the aligned buffer as a mutable slice.
+    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
+        // SAFETY: same invariant as as_bytes; &mut self rules out other
+        // simultaneous borrows.
+        unsafe { std::slice::from_raw_parts_mut(self.aligned_ptr, self.size) }
+    }
+
+    /// Returns the guest address for this op.
+    pub fn data_addr(&self) -> GuestAddress {
+        self.data_addr
+    }
+}
+
+impl Drop for AlignedOperation {
+    fn drop(&mut self) {
+        // SAFETY: `new` is the only constructor, and it stores a pointer
+        // returned by `alloc_zeroed` paired with the exact `layout` used
+        // for that allocation. Ownership has not escaped (the type is
+        // neither `Clone` nor `Copy`).
+        unsafe {
+            dealloc(self.aligned_ptr, self.layout);
+        }
+    }
+}
+
+// SAFETY: AlignedOperation owns its heap allocation exclusively (no Clone/
+// Copy, no shared aliases) and the allocation's lifetime is tied to the
+// value's. Moving an AlignedOperation between threads transfers that
+// ownership — the same rationale Box<T> uses for its Send impl.
+unsafe impl Send for AlignedOperation {}

block/src/lib.rs

@@ -8,6 +8,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
 
+mod aligned_operation;
 pub mod async_io;
 pub mod disk_file;
 pub mod error;
@@ -49,16 +50,14 @@ use std::path::Path;
 use std::str::FromStr;
 use std::{cmp, mem, result};
 
+pub use aligned_operation::AlignedOperation;
 #[cfg(feature = "io_uring")]
 use io_uring::{IoUring, Probe, opcode};
 use libc::{
     FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE, FALLOC_FL_ZERO_RANGE, S_IFBLK, S_IFMT, ioctl,
 };
 use log::{debug, info, warn};
-pub use request::{
-    AlignedOperation, BatchRequest, ExecuteAsync, MAX_DISCARD_WRITE_ZEROES_SEG, Request,
-    RequestType,
-};
+pub use request::{BatchRequest, ExecuteAsync, MAX_DISCARD_WRITE_ZEROES_SEG, Request, RequestType};
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
 use thiserror::Error;

block/src/request.rs

@@ -8,7 +8,6 @@
 //
 // SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
 
-use std::alloc::{Layout, alloc_zeroed, dealloc};
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::mem;
 use std::time::Instant;
@@ -27,6 +26,7 @@ use vm_memory::{
 };
 use vm_virtio::{AccessPlatform, Translatable as _};
 
+use crate::aligned_operation::AlignedOperation;
 use crate::async_io::AsyncIo;
 use crate::{Error, ExecuteError, request_type, sector};
 
@@ -43,13 +43,6 @@ const DISCARD_WZ_SECTOR_OFFSET: u64 =
 const DISCARD_WZ_NUM_SECTORS_OFFSET: u64 =
     mem::offset_of!(virtio_blk_discard_write_zeroes, num_sectors) as u64;
 const DISCARD_WZ_FLAGS_OFFSET: u64 = mem::offset_of!(virtio_blk_discard_write_zeroes, flags) as u64;
-#[derive(Debug)]
-pub struct AlignedOperation {
-    origin_ptr: u64,
-    aligned_ptr: u64,
-    size: usize,
-    layout: Layout,
-}
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum RequestType {
@@ -284,31 +277,18 @@ impl Request {
             let iov_base = if (origin_ptr.as_ptr() as u64).is_multiple_of(alignment) {
                 origin_ptr.as_ptr().cast()
             } else {
-                let layout = Layout::from_size_align(data_len, alignment as usize).unwrap();
-                // SAFETY: layout has non-zero size
-                let aligned_ptr = unsafe { alloc_zeroed(layout) };
-                if aligned_ptr.is_null() {
-                    return Err(ExecuteError::TemporaryBufferAllocation(
-                        std::io::Error::last_os_error(),
-                    ));
-                }
+                let mut aligned_op = AlignedOperation::new(data_addr, data_len, alignment as usize)
+                    .map_err(ExecuteError::TemporaryBufferAllocation)?;
 
                 // We need to perform the copy beforehand in case we're writing
                 // data out.
                 if request_type == RequestType::Out {
-                    // SAFETY: destination buffer has been allocated with
-                    // the proper size.
-                    unsafe { std::ptr::copy(origin_ptr.as_ptr(), aligned_ptr, data_len) };
+                    mem.read_slice(aligned_op.as_bytes_mut(), data_addr)
+                        .map_err(ExecuteError::Read)?;
                 }
 
-                // Store both origin and aligned pointers for complete_async()
-                // to process them.
-                self.aligned_operations.push(AlignedOperation {
-                    origin_ptr: origin_ptr.as_ptr() as u64,
-                    aligned_ptr: aligned_ptr as u64,
-                    size: data_len,
-                    layout,
-                });
+                let aligned_ptr = aligned_op.as_mut_ptr();
+                self.aligned_operations.push(aligned_op);
 
                 aligned_ptr.cast()
             };
@@ -522,31 +502,17 @@ impl Request {
         Ok(ret)
     }
 
-    pub fn complete_async(&mut self) -> Result<(), Error> {
-        for aligned_operation in self.aligned_operations.drain(..) {
+    pub fn complete_async<B: Bitmap + 'static>(
+        &mut self,
+        mem: &vm_memory::GuestMemoryMmap<B>,
+    ) -> Result<(), Error> {
+        for aligned_op in self.aligned_operations.drain(..) {
             // We need to perform the copy after the data has been read inside
             // the aligned buffer in case we're reading data in.
             if self.request_type == RequestType::In {
-                // SAFETY: origin buffer has been allocated with the
-                // proper size.
-                unsafe {
-                    std::ptr::copy(
-                        aligned_operation.aligned_ptr as *const u8,
-                        aligned_operation.origin_ptr as *mut u8,
-                        aligned_operation.size,
-                    );
-                };
+                mem.write_slice(aligned_op.as_bytes(), aligned_op.data_addr())
+                    .map_err(Error::GuestMemory)?;
             }
-
-            // Free the temporary aligned buffer.
-            // SAFETY: aligned_ptr was allocated by alloc_zeroed with the same
-            // layout
-            unsafe {
-                dealloc(
-                    aligned_operation.aligned_ptr as *mut u8,
-                    aligned_operation.layout,
-                );
-            };
         }
 
         Ok(())

virtio-devices/src/block.rs

@@ -456,7 +456,9 @@ impl BlockEpollHandler {
 
             let mut request = self.find_inflight_request(desc_index)?;
 
-            request.complete_async().map_err(Error::RequestCompleting)?;
+            request
+                .complete_async(&mem)
+                .map_err(Error::RequestCompleting)?;
 
             let latency = request.start().elapsed().as_micros() as u64;
             let read_ops_last = self.counters.read_ops.load(Ordering::Relaxed);