pvmfw: Stop confusing MEM_SHARE granule & DMA size

The VirtIO HAL assumes that shared allocations are aligned to a size
that is a multiple of the DMA alignment required by virtio devices. This
is wrong because 1. the allocator isn't forced to align the region it
returns to the memory granule (e.g. if 2 regions are requested that fit
in a single granule, it should be allowed to allocate only one granule)
and 2. nothing guarantees the granule to be a multiple of the DMA
alignment.

Therefore, modify the {de,}alloc_shared API to more closely follow
Rust-standard APIs (e.g. GlobalAllocator or LockedHeap) by taking a
Layout, allowing the caller to pass the alignment it requires.

Use that in the virtio_drivers HAL to request:

    - PAGE_SIZE-aligned regions for DMA;
    - size_of::<u128>()-aligned regions for bounce buffers

Keep the memory.rs code functionally unchanged for now, so that
allocations are still granule-aligned after this patch.

Bug: 280644106
Test: atest MicrodroidTests
Change-Id: I2b22e2095cde48e59e7303efeed1e3c09ee5ad5b
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index 16c1a37..7e8423a 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -16,7 +16,7 @@
 
 #![deny(unsafe_op_in_unsafe_fn)]
 
-use crate::helpers::{self, align_down, align_up, page_4kb_of, RangeExt, SIZE_4KB, SIZE_4MB};
+use crate::helpers::{self, align_down, page_4kb_of, RangeExt, SIZE_4KB, SIZE_4MB};
 use crate::mmu;
 use alloc::alloc::alloc_zeroed;
 use alloc::alloc::dealloc;
@@ -346,13 +346,12 @@
 /// host. Returns a pointer to the buffer.
 ///
 /// It will be aligned to the memory sharing granule size supported by the hypervisor.
-pub fn alloc_shared(size: usize) -> hyp::Result<NonNull<u8>> {
-    let layout = shared_buffer_layout(size)?;
-    let granule = layout.align();
-
+pub fn alloc_shared(layout: Layout) -> hyp::Result<NonNull<u8>> {
+    assert_ne!(layout.size(), 0);
+    let granule = get_hypervisor().memory_protection_granule()?;
+    let layout = layout.align_to(granule).unwrap().pad_to_align();
     if let Some(shared_pool) = SHARED_POOL.get() {
-        // Safe because `shared_buffer_layout` panics if the size is 0, so the
-        // layout must have a non-zero size.
+        // SAFETY - layout has a non-zero size.
         let buffer = unsafe { shared_pool.alloc_zeroed(layout) };
 
         let Some(buffer) = NonNull::new(buffer) else {
@@ -363,8 +362,7 @@
         return Ok(buffer);
     }
 
-    // Safe because `shared_buffer_layout` panics if the size is 0, so the layout must have a
-    // non-zero size.
+    // SAFETY - layout has a non-zero size.
     let buffer = unsafe { alloc_zeroed(layout) };
 
     let Some(buffer) = NonNull::new(buffer) else {
@@ -388,10 +386,9 @@
 ///
 /// The memory must have been allocated by `alloc_shared` with the same size, and not yet
 /// deallocated.
-pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, size: usize) -> hyp::Result<()> {
-    let layout = shared_buffer_layout(size)?;
-    let granule = layout.align();
-
+pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, layout: Layout) -> hyp::Result<()> {
+    let granule = get_hypervisor().memory_protection_granule()?;
+    let layout = layout.align_to(granule).unwrap().pad_to_align();
     if let Some(shared_pool) = SHARED_POOL.get() {
         // Safe because the memory was allocated by `alloc_shared` above using
         // the same allocator, and the layout is the same as was used then.
@@ -411,20 +408,6 @@
     Ok(())
 }
 
-/// Returns the layout to use for allocating a buffer of at least the given size shared with the
-/// host.
-///
-/// It will be aligned to the memory sharing granule size supported by the hypervisor.
-///
-/// Panics if `size` is 0.
-fn shared_buffer_layout(size: usize) -> hyp::Result<Layout> {
-    assert_ne!(size, 0);
-    let granule = get_hypervisor().memory_protection_granule()?;
-    let allocated_size =
-        align_up(size, granule).expect("Memory protection granule was not a power of two");
-    Ok(Layout::from_size_align(allocated_size, granule).unwrap())
-}
-
 /// Returns an iterator which yields the base address of each 4 KiB page within the given range.
 fn page_iterator(range: &MemoryRange) -> impl Iterator<Item = usize> {
     (page_4kb_of(range.start)..range.end).step_by(SIZE_4KB)
diff --git a/pvmfw/src/virtio/hal.rs b/pvmfw/src/virtio/hal.rs
index 41d0f24..ec0b9d8 100644
--- a/pvmfw/src/virtio/hal.rs
+++ b/pvmfw/src/virtio/hal.rs
@@ -17,6 +17,8 @@
 use super::pci::PCI_INFO;
 use crate::helpers::RangeExt as _;
 use crate::memory::{alloc_shared, dealloc_shared, phys_to_virt, virt_to_phys};
+use core::alloc::Layout;
+use core::mem::size_of;
 use core::ptr::{copy_nonoverlapping, NonNull};
 use log::trace;
 use virtio_drivers::{BufferDirection, Hal, PhysAddr, PAGE_SIZE};
@@ -39,17 +41,15 @@
     ///  block of memory using `alloc_shared` and returning a non-null pointer to it that is
     ///  aligned to `PAGE_SIZE`.
     fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
-        let size = pages * PAGE_SIZE;
-        let vaddr =
-            alloc_shared(size).expect("Failed to allocate and share VirtIO DMA range with host");
+        let vaddr = alloc_shared(dma_layout(pages))
+            .expect("Failed to allocate and share VirtIO DMA range with host");
         let paddr = virt_to_phys(vaddr);
         (paddr, vaddr)
     }
 
     unsafe fn dma_dealloc(_paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
-        let size = pages * PAGE_SIZE;
         // SAFETY - Memory was allocated by `dma_alloc` using `alloc_shared` with the same size.
-        unsafe { dealloc_shared(vaddr, size) }
+        unsafe { dealloc_shared(vaddr, dma_layout(pages)) }
             .expect("Failed to unshare VirtIO DMA range with host");
         0
     }
@@ -85,7 +85,7 @@
 
         // TODO: Copy to a pre-shared region rather than allocating and sharing each time.
         // Allocate a range of pages, copy the buffer if necessary, and share the new range instead.
-        let bounce = alloc_shared(size)
+        let bounce = alloc_shared(bb_layout(size))
             .expect("Failed to allocate and share VirtIO bounce buffer with host");
         let paddr = virt_to_phys(bounce);
         if direction == BufferDirection::DriverToDevice {
@@ -109,7 +109,18 @@
         }
 
         // SAFETY - Memory was allocated by `share` using `alloc_shared` with the same size.
-        unsafe { dealloc_shared(bounce, size) }
+        unsafe { dealloc_shared(bounce, bb_layout(size)) }
             .expect("Failed to unshare and deallocate VirtIO bounce buffer");
     }
 }
+
+fn dma_layout(pages: usize) -> Layout {
+    let size = pages.checked_mul(PAGE_SIZE).unwrap();
+    Layout::from_size_align(size, PAGE_SIZE).unwrap()
+}
+
+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()
+}