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>
(cherry picked from commit 1b8c92dd5e3c0c58316826486ce5ee30eeb71407))
[backport: adapted to stable/v51.x; v51.x has no block/src/request.rs
 split, so the new aligned_operation module is added next to
 block/src/lib.rs and the in tree struct, alloc, free path is
 replaced in place.]
Signed-off-by: Anatol Belski <anbelski@linux.microsoft.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 fcntl;
 pub mod fixed_vhd;
@@ -28,7 +29,7 @@ pub mod vhd;
 pub mod vhdx;
 pub mod vhdx_sync;
 
-use std::alloc::{Layout, alloc_zeroed, dealloc};
+use std::alloc::{Layout, alloc_zeroed};
 use std::collections::VecDeque;
 use std::fmt::{self, Debug};
 use std::fs::File;
@@ -57,6 +58,8 @@ use vm_virtio::{AccessPlatform, Translatable};
 use vmm_sys_util::eventfd::EventFd;
 use vmm_sys_util::{aio, ioctl_io_nr};
 
+pub use aligned_operation::AlignedOperation;
+
 use crate::async_io::{AsyncIo, AsyncIoError, AsyncIoResult};
 use crate::vhdx::VhdxError;
 
@@ -232,14 +235,6 @@ fn sector<B: Bitmap + 'static>(
 
 const DEFAULT_DESCRIPTOR_VEC_SIZE: usize = 32;
 
-#[derive(Debug)]
-pub struct AlignedOperation {
-    origin_ptr: u64,
-    aligned_ptr: u64,
-    size: usize,
-    layout: Layout,
-}
-
 pub struct BatchRequest {
     pub offset: libc::off_t,
     pub iovecs: SmallVec<[libc::iovec; DEFAULT_DESCRIPTOR_VEC_SIZE]>,
@@ -473,31 +468,19 @@ impl Request {
             let iov_base = if (origin_ptr.as_ptr() as u64).is_multiple_of(SECTOR_SIZE) {
                 origin_ptr.as_ptr() as *mut libc::c_void
             } else {
-                let layout = Layout::from_size_align(data_len, SECTOR_SIZE 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(
-                        io::Error::last_os_error(),
-                    ));
-                }
+                let mut aligned_op =
+                    AlignedOperation::new(data_addr, data_len, SECTOR_SIZE 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 as *mut libc::c_void
             };
@@ -639,31 +622,17 @@ impl Request {
         Ok(ret)
     }
 
-    pub fn complete_async(&mut self) -> result::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::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

@@ -461,7 +461,9 @@ Setting device status to 'NEEDS_RESET' and stopping processing queues until rese
 
             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);