Ensure that VirtIO buffers are aligned to PAGE_SIZE as required.

The existing code assumed that the hypervisor granule was greater than
virtio_driver::PAGE_SIZE, which is true in practice so far but not
guaranteed.

Removed redundant documentation from the Hal trait implementation, and
made the implementation safety comments more explicit.

Bug: 274732281
Bug: 237249346
Test: Build virt APEX
Change-Id: I776dcda7f5660b6f6e696c6834ac03f8e4023543
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index cfd6b0a..b0a9406 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -395,10 +395,8 @@
     }
 }
 
-/// Allocates a memory range of at least the given size that is shared with
-/// host. Returns a pointer to the buffer.
-///
-/// It will be aligned to the memory sharing granule size supported by the hypervisor.
+/// Allocates a memory range of at least the given size and alignment that is shared with the host.
+/// Returns a pointer to the buffer.
 pub fn alloc_shared(layout: Layout) -> hyp::Result<NonNull<u8>> {
     assert_ne!(layout.size(), 0);
     let Some(buffer) = try_shared_alloc(layout) else {
@@ -424,11 +422,11 @@
 
 /// Unshares and deallocates a memory range which was previously allocated by `alloc_shared`.
 ///
-/// The size passed in must be the size passed to the original `alloc_shared` call.
+/// The layout passed in must be the same layout passed to the original `alloc_shared` call.
 ///
 /// # Safety
 ///
-/// The memory must have been allocated by `alloc_shared` with the same size, and not yet
+/// The memory must have been allocated by `alloc_shared` with the same layout, and not yet
 /// deallocated.
 pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, layout: Layout) -> hyp::Result<()> {
     SHARED_POOL.get().unwrap().lock().dealloc(vaddr, layout);
diff --git a/pvmfw/src/virtio/hal.rs b/pvmfw/src/virtio/hal.rs
index 51567cd..becc263 100644
--- a/pvmfw/src/virtio/hal.rs
+++ b/pvmfw/src/virtio/hal.rs
@@ -23,23 +23,25 @@
 use log::trace;
 use virtio_drivers::{BufferDirection, Hal, PhysAddr, PAGE_SIZE};
 
+/// The alignment to use for the temporary buffers allocated by `HalImpl::share`. There doesn't seem
+/// to be any particular alignment required by VirtIO for these, so 16 bytes should be enough to
+/// allow appropriate alignment for whatever fields are accessed. `alloc_shared` will increase the
+/// alignment to the memory sharing granule size anyway.
+const SHARED_BUFFER_ALIGNMENT: usize = size_of::<u128>();
+
 pub struct HalImpl;
 
-/// Implements the `Hal` trait for `HalImpl`.
-///
 /// # Safety
 ///
-/// Callers of this implementatation must follow the safety requirements documented in the `Hal`
-/// trait for the unsafe methods.
+/// See the 'Implementation Safety' comments on methods below for how they fulfill the safety
+/// requirements of the unsafe `Hal` trait.
 unsafe impl Hal for HalImpl {
-    /// Allocates the given number of contiguous physical pages of DMA memory for VirtIO use.
-    ///
     /// # Implementation Safety
     ///
     /// `dma_alloc` ensures the returned DMA buffer is not aliased with any other allocation or
-    ///  reference in the program until it is deallocated by `dma_dealloc` by allocating a unique
-    ///  block of memory using `alloc_shared` and returning a non-null pointer to it that is
-    ///  aligned to `PAGE_SIZE`.
+    /// reference in the program until it is deallocated by `dma_dealloc` by allocating a unique
+    /// block of memory using `alloc_shared`, which is guaranteed to allocate valid and unique
+    /// memory. We request an alignment of at least `PAGE_SIZE` from `alloc_shared`.
     fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
         let vaddr = alloc_shared(dma_layout(pages))
             .expect("Failed to allocate and share VirtIO DMA range with host");
@@ -51,18 +53,18 @@
     }
 
     unsafe fn dma_dealloc(_paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
-        // SAFETY - Memory was allocated by `dma_alloc` using `alloc_shared` with the same size.
+        // SAFETY - Memory was allocated by `dma_alloc` using `alloc_shared` with the same layout.
         unsafe { dealloc_shared(vaddr, dma_layout(pages)) }
             .expect("Failed to unshare VirtIO DMA range with host");
         0
     }
 
-    /// Converts a physical address used for MMIO to a virtual address which the driver can access.
-    ///
     /// # Implementation Safety
     ///
-    /// `mmio_phys_to_virt` satisfies the requirement by checking that the mapped memory region
-    /// is within the PCI MMIO range.
+    /// The returned pointer must be valid because the `paddr` describes a valid MMIO region, we
+    /// check that it is within the PCI MMIO range, and we previously mapped the entire PCI MMIO
+    /// range. It can't alias any other allocations because we previously validated in
+    /// `map_mmio_range` that the PCI MMIO range didn't overlap with any other memory ranges.
     unsafe fn mmio_phys_to_virt(paddr: PhysAddr, size: usize) -> NonNull<u8> {
         let pci_info = PCI_INFO.get().expect("VirtIO HAL used before PCI_INFO was initialised");
         let bar_range = {
@@ -109,7 +111,7 @@
             unsafe { copy_nonoverlapping(bounce.as_ptr(), dest, size) };
         }
 
-        // SAFETY - Memory was allocated by `share` using `alloc_shared` with the same size.
+        // SAFETY - Memory was allocated by `share` using `alloc_shared` with the same layout.
         unsafe { dealloc_shared(bounce, bb_layout(size)) }
             .expect("Failed to unshare and deallocate VirtIO bounce buffer");
     }
@@ -121,7 +123,5 @@
 }
 
 fn bb_layout(size: usize) -> Layout {
-    // In theory, it would be legal to align to 1-byte but use a larger alignment for good measure.
-    const VIRTIO_BOUNCE_BUFFER_ALIGN: usize = size_of::<u128>();
-    Layout::from_size_align(size, VIRTIO_BOUNCE_BUFFER_ALIGN).unwrap()
+    Layout::from_size_align(size, SHARED_BUFFER_ALIGNMENT).unwrap()
 }