[virtio] Move virtio module from pvmfw to vmbase

for reuse in both pvmfw and rialo.

No code change in this cl.

Bug: 284462758
Test: m pvmfw_img
Change-Id: I225b495a2f8c1a49cb63ea69d60f01949bab3f05
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index 6420b2b..6db9ff8 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -65,12 +65,14 @@
     rustlibs: [
         "libaarch64_paging",
         "libbuddy_system_allocator",
+        "libfdtpci",
         "libhyp",
         "liblog_rust_nostd",
         "libonce_cell_nostd",
         "libsmccc",
         "libspin_nostd",
         "libtinyvec_nostd",
+        "libvirtio_drivers",
         "libzeroize_nostd",
     ],
     whole_static_libs: [
diff --git a/vmbase/src/lib.rs b/vmbase/src/lib.rs
index 80cdf4e..ebb3707 100644
--- a/vmbase/src/lib.rs
+++ b/vmbase/src/lib.rs
@@ -29,6 +29,7 @@
 pub mod power;
 pub mod uart;
 pub mod util;
+pub mod virtio;
 
 pub use bionic::STACK_CHK_GUARD;
 
diff --git a/vmbase/src/virtio/hal.rs b/vmbase/src/virtio/hal.rs
new file mode 100644
index 0000000..ac5b967
--- /dev/null
+++ b/vmbase/src/virtio/hal.rs
@@ -0,0 +1,128 @@
+// Copyright 2023, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! HAL for the virtio_drivers crate.
+
+use super::pci::PCI_INFO;
+use crate::memory::{alloc_shared, dealloc_shared, phys_to_virt, virt_to_phys};
+use crate::util::RangeExt as _;
+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};
+
+/// 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>();
+
+/// HAL implementation for the virtio_drivers crate.
+pub struct HalImpl;
+
+/// # Safety
+///
+/// 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 {
+    /// # 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`, 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");
+        // TODO(ptosi): Move this zeroing to virtio_drivers, if it silently wants a zeroed region.
+        // SAFETY - vaddr points to a region allocated for the caller so is safe to access.
+        unsafe { core::ptr::write_bytes(vaddr.as_ptr(), 0, dma_layout(pages).size()) };
+        let paddr = virt_to_phys(vaddr);
+        (paddr, vaddr)
+    }
+
+    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 layout.
+        unsafe { dealloc_shared(vaddr, dma_layout(pages)) }
+            .expect("Failed to unshare VirtIO DMA range with host");
+        0
+    }
+
+    /// # Implementation Safety
+    ///
+    /// 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 = {
+            let start = pci_info.bar_range.start.try_into().unwrap();
+            let end = pci_info.bar_range.end.try_into().unwrap();
+
+            start..end
+        };
+        let mmio_range = paddr..paddr.checked_add(size).expect("PCI MMIO region end overflowed");
+
+        // Check that the region is within the PCI MMIO range that we read from the device tree. If
+        // not, the host is probably trying to do something malicious.
+        assert!(
+            mmio_range.is_within(&bar_range),
+            "PCI MMIO region was outside of expected BAR range.",
+        );
+
+        phys_to_virt(paddr)
+    }
+
+    unsafe fn share(buffer: NonNull<[u8]>, direction: BufferDirection) -> PhysAddr {
+        let size = buffer.len();
+
+        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 {
+            let src = buffer.cast::<u8>().as_ptr().cast_const();
+            trace!("VirtIO bounce buffer at {bounce:?} (PA:{paddr:#x}) initialized from {src:?}");
+            // SAFETY - Both regions are valid, properly aligned, and don't overlap.
+            unsafe { copy_nonoverlapping(src, bounce.as_ptr(), size) };
+        }
+
+        paddr
+    }
+
+    unsafe fn unshare(paddr: PhysAddr, buffer: NonNull<[u8]>, direction: BufferDirection) {
+        let bounce = phys_to_virt(paddr);
+        let size = buffer.len();
+        if direction == BufferDirection::DeviceToDriver {
+            let dest = buffer.cast::<u8>().as_ptr();
+            trace!("VirtIO bounce buffer at {bounce:?} (PA:{paddr:#x}) copied back to {dest:?}");
+            // SAFETY - Both regions are valid, properly aligned, and don't overlap.
+            unsafe { copy_nonoverlapping(bounce.as_ptr(), dest, 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");
+    }
+}
+
+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 {
+    Layout::from_size_align(size, SHARED_BUFFER_ALIGNMENT).unwrap()
+}
diff --git a/vmbase/src/virtio/mod.rs b/vmbase/src/virtio/mod.rs
new file mode 100644
index 0000000..df916bc
--- /dev/null
+++ b/vmbase/src/virtio/mod.rs
@@ -0,0 +1,18 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Modules for working with VirtIO devices.
+
+mod hal;
+pub mod pci;
diff --git a/vmbase/src/virtio/pci.rs b/vmbase/src/virtio/pci.rs
new file mode 100644
index 0000000..b54f7d1
--- /dev/null
+++ b/vmbase/src/virtio/pci.rs
@@ -0,0 +1,128 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Functions to scan the PCI bus for VirtIO devices.
+
+use super::hal::HalImpl;
+use crate::memory::{MemoryTracker, MemoryTrackerError};
+use alloc::boxed::Box;
+use core::fmt;
+use fdtpci::PciInfo;
+use log::debug;
+use once_cell::race::OnceBox;
+use virtio_drivers::{
+    device::blk,
+    transport::{
+        pci::{
+            bus::{BusDeviceIterator, PciRoot},
+            virtio_device_type, PciTransport,
+        },
+        DeviceType, Transport,
+    },
+};
+
+pub(super) static PCI_INFO: OnceBox<PciInfo> = OnceBox::new();
+
+/// PCI errors.
+#[derive(Debug, Clone)]
+pub enum PciError {
+    /// Attempted to initialize the PCI more than once.
+    DuplicateInitialization,
+    /// Failed to map PCI CAM.
+    CamMapFailed(MemoryTrackerError),
+    /// Failed to map PCI BAR.
+    BarMapFailed(MemoryTrackerError),
+}
+
+impl fmt::Display for PciError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::DuplicateInitialization => {
+                write!(f, "Attempted to initialize the PCI more than once.")
+            }
+            Self::CamMapFailed(e) => write!(f, "Failed to map PCI CAM: {e}"),
+            Self::BarMapFailed(e) => write!(f, "Failed to map PCI BAR: {e}"),
+        }
+    }
+}
+
+/// Prepares to use VirtIO PCI devices.
+///
+/// In particular:
+///
+/// 1. Maps the PCI CAM and BAR range in the page table and MMIO guard.
+/// 2. Stores the `PciInfo` for the VirtIO HAL to use later.
+/// 3. Creates and returns a `PciRoot`.
+///
+/// This must only be called once; it will panic if it is called a second time.
+pub fn initialise(pci_info: PciInfo, memory: &mut MemoryTracker) -> Result<PciRoot, PciError> {
+    PCI_INFO.set(Box::new(pci_info.clone())).map_err(|_| PciError::DuplicateInitialization)?;
+
+    memory.map_mmio_range(pci_info.cam_range.clone()).map_err(PciError::CamMapFailed)?;
+    let bar_range = pci_info.bar_range.start as usize..pci_info.bar_range.end as usize;
+    memory.map_mmio_range(bar_range).map_err(PciError::BarMapFailed)?;
+
+    // Safety: This is the only place where we call make_pci_root, and `PCI_INFO.set` above will
+    // panic if it is called a second time.
+    Ok(unsafe { pci_info.make_pci_root() })
+}
+
+/// Virtio Block device.
+pub type VirtIOBlk = blk::VirtIOBlk<HalImpl, PciTransport>;
+
+/// Virtio Block device iterator.
+pub struct VirtIOBlkIterator<'a> {
+    pci_root: &'a mut PciRoot,
+    bus: BusDeviceIterator,
+}
+
+impl<'a> VirtIOBlkIterator<'a> {
+    /// Creates a new iterator.
+    pub fn new(pci_root: &'a mut PciRoot) -> Self {
+        let bus = pci_root.enumerate_bus(0);
+        Self { pci_root, bus }
+    }
+}
+
+impl<'a> Iterator for VirtIOBlkIterator<'a> {
+    type Item = VirtIOBlk;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let (device_function, info) = self.bus.next()?;
+            let (status, command) = self.pci_root.get_status_command(device_function);
+            debug!(
+                "Found PCI device {} at {}, status {:?} command {:?}",
+                info, device_function, status, command
+            );
+
+            let Some(virtio_type) = virtio_device_type(&info) else {
+                continue;
+            };
+            debug!("  VirtIO {:?}", virtio_type);
+
+            let mut transport =
+                PciTransport::new::<HalImpl>(self.pci_root, device_function).unwrap();
+            debug!(
+                "Detected virtio PCI device with device type {:?}, features {:#018x}",
+                transport.device_type(),
+                transport.read_device_features(),
+            );
+
+            if virtio_type == DeviceType::Block {
+                return Some(Self::Item::new(transport).expect("failed to create blk driver"));
+            }
+        }
+    }
+}