scsi: support multiple units per single enclosure

Cargo.lock

@@ -2681,6 +2681,8 @@ name = "ygg_driver_net_loopback"
 version = "0.1.0"
 dependencies = [
  "bytemuck",
+ "device-api",
+ "libk",
  "libk-mm",
  "libk-util",
  "ygg_driver_net_core",
@@ -2794,6 +2796,7 @@ version = "0.1.0"
 dependencies = [
  "bitflags 2.8.0",
  "device-api",
+ "libk",
  "libk-mm",
  "libk-util",
  "log",
@@ -2824,6 +2827,7 @@ dependencies = [
  "bitflags 2.8.0",
  "bytemuck",
  "device-api",
+ "libk",
  "libk-mm",
  "libk-util",
  "log",

kernel/driver/block/ahci/src/command.rs

@@ -1,12 +1,10 @@
 use core::mem::{size_of, MaybeUninit};
 
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox, PageSlice,
-};
+use device_api::dma::DmaAllocator;
+use libk::dma::{BusAddress, DmaBuffer, DmaSliceMut};
 use tock_registers::register_structs;
 
-use crate::{data::AtaString, error::AhciError, MAX_PRD_SIZE, SECTOR_SIZE};
+use crate::{data::AtaString, error::AhciError, MAX_PRD_SIZE};
 
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 #[repr(u8)]
@@ -22,7 +20,7 @@ pub trait AtaCommand {
 
     fn lba(&self) -> u64;
     fn sector_count(&self) -> usize;
-    fn buffer(&self) -> Option<(PhysicalAddress, usize)>;
+    fn buffer(&self) -> Option<(BusAddress, usize)>;
     unsafe fn into_response(self) -> Self::Response;
 
     fn prd_count(&self) -> usize {
@@ -64,44 +62,41 @@ register_structs! {
 }
 
 pub struct AtaIdentify {
-    buffer: PageBox<MaybeUninit<AtaIdentifyResponse>>,
+    buffer: DmaBuffer<MaybeUninit<AtaIdentifyResponse>>,
 }
 
 pub struct AtaReadDmaEx {
     lba: u64,
     sector_count: usize,
-    buffer_base: PhysicalAddress,
+    buffer_base: BusAddress,
     buffer_size: usize,
 }
 
 impl AtaIdentify {
-    pub fn create() -> Result<Self, AhciError> {
-        PageBox::new_uninit()
+    pub fn create(dma: &dyn DmaAllocator) -> Result<Self, AhciError> {
+        DmaBuffer::new_uninit(dma)
             .map(Self::with_data)
             .map_err(AhciError::MemoryError)
     }
 
-    pub fn with_data(buffer: PageBox<MaybeUninit<AtaIdentifyResponse>>) -> Self {
+    pub fn with_data(buffer: DmaBuffer<MaybeUninit<AtaIdentifyResponse>>) -> Self {
         Self { buffer }
     }
 }
 
 impl AtaReadDmaEx {
-    pub fn new(lba: u64, sector_count: usize, buffer: &PageSlice<MaybeUninit<u8>>) -> Self {
-        assert_eq!(buffer.len() % SECTOR_SIZE, 0);
-        assert_ne!(buffer.len(), 0);
-
+    pub fn new(lba: u64, sector_count: usize, buffer: DmaSliceMut<MaybeUninit<u8>>) -> Self {
         Self {
             lba,
             sector_count,
-            buffer_base: unsafe { buffer.as_physical_address() },
+            buffer_base: buffer.bus_address(),
             buffer_size: buffer.len(),
         }
     }
 }
 
 impl AtaCommand for AtaIdentify {
-    type Response = PageBox<AtaIdentifyResponse>;
+    type Response = DmaBuffer<AtaIdentifyResponse>;
 
     const COMMAND_ID: AtaCommandId = AtaCommandId::Identify;
 
@@ -113,14 +108,14 @@ impl AtaCommand for AtaIdentify {
         0
     }
 
-    fn buffer(&self) -> Option<(PhysicalAddress, usize)> {
-        let base = unsafe { self.buffer.as_physical_address() };
+    fn buffer(&self) -> Option<(BusAddress, usize)> {
+        let base = self.buffer.bus_address();
         let size = size_of::<AtaIdentifyResponse>();
         Some((base, size))
     }
 
     unsafe fn into_response(self) -> Self::Response {
-        self.buffer.assume_init()
+        DmaBuffer::assume_init(self.buffer)
     }
 }
 
@@ -137,7 +132,7 @@ impl AtaCommand for AtaReadDmaEx {
         self.sector_count
     }
 
-    fn buffer(&self) -> Option<(PhysicalAddress, usize)> {
+    fn buffer(&self) -> Option<(BusAddress, usize)> {
         Some((self.buffer_base, self.buffer_size))
     }
 

kernel/driver/block/ahci/src/data.rs

@@ -2,7 +2,7 @@ use core::mem::size_of;
 
 use alloc::string::String;
 use bytemuck::{Pod, Zeroable};
-use libk_mm::address::PhysicalAddress;
+use libk::dma::BusAddress;
 use libk_util::{ConstAssert, IsTrue};
 use static_assertions::const_assert_eq;
 
@@ -174,7 +174,7 @@ impl CommandTable {
 }
 
 impl CommandListEntry {
-    pub fn new(command_table_entry: PhysicalAddress, prd_count: usize) -> Result<Self, AhciError> {
+    pub fn new(command_table_entry: BusAddress, prd_count: usize) -> Result<Self, AhciError> {
         if prd_count > 0xFFFF {
             todo!()
         }
@@ -183,7 +183,7 @@ impl CommandListEntry {
             attr: (size_of::<RegisterHostToDeviceFis>() / size_of::<u32>()) as _,
             prdtl: prd_count as _,
             prdbc: 0,
-            ctba: command_table_entry.into(),
+            ctba: command_table_entry.into_u64(),
             _0: [0; 4],
         })
     }
@@ -201,18 +201,14 @@ unsafe impl Zeroable for CommandTable {
 }
 
 impl PhysicalRegionDescriptor {
-    pub fn new(
-        address: PhysicalAddress,
-        byte_count: usize,
-        is_last: bool,
-    ) -> Result<Self, AhciError> {
+    pub fn new(address: BusAddress, byte_count: usize, is_last: bool) -> Result<Self, AhciError> {
         if byte_count > MAX_PRD_SIZE {
             return Err(AhciError::RegionTooLarge);
         }
 
         let dbc_mask = (is_last as u32) << 31;
         Ok(Self {
-            buffer_address: address.into(),
+            buffer_address: address.into_u64(),
             _0: 0,
             dbc: ((byte_count as u32 - 1) << 1) | 1 | dbc_mask,
         })

kernel/driver/block/ahci/src/lib.rs

@@ -8,12 +8,13 @@ use alloc::{format, sync::Arc, vec::Vec};
 use bytemuck::Zeroable;
 use data::ReceivedFis;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
+    dma::DmaAllocator,
     interrupt::{InterruptAffinity, InterruptHandler, IrqVector},
 };
 use error::AhciError;
-use libk::{device::manager::probe_partitions, fs::devfs, task::runtime};
-use libk_mm::{address::AsPhysicalAddress, device::DeviceMemoryIo, PageBox};
+use libk::{device::manager::probe_partitions, dma::DmaBuffer, fs::devfs, task::runtime};
+use libk_mm::device::DeviceMemoryIo;
 use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
 use port::AhciPort;
 use regs::{PortRegs, Regs};
@@ -40,8 +41,9 @@ const MAX_DRIVES: usize = (b'z' - b'a') as usize;
 
 pub struct AhciController {
     regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
+    dma: Arc<dyn DmaAllocator>,
     ports: OneTimeInit<Vec<Arc<AhciPort>>>,
-    received_fis_buffers: OneTimeInit<[Option<PageBox<ReceivedFis>>; 16]>,
+    received_fis_buffers: OneTimeInit<[Option<DmaBuffer<ReceivedFis>>; 16]>,
 
     version: Version,
     max_port_count: usize,
@@ -81,8 +83,9 @@ impl AhciController {
             let regs = self.regs.lock();
             let port = &regs.PORTS[i];
 
-            let buffer = PageBox::new(ReceivedFis::zeroed()).map_err(AhciError::MemoryError)?;
-            port.set_received_fis_address_64(unsafe { buffer.as_physical_address() });
+            let buffer = DmaBuffer::new(&*self.dma, ReceivedFis::zeroed())
+                .map_err(AhciError::MemoryError)?;
+            port.set_received_fis_address_64(buffer.bus_address());
             *fis_buffer_slot = Some(buffer);
         }
 
@@ -181,7 +184,7 @@ impl InterruptHandler for AhciController {
 }
 
 impl Device for AhciController {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         // Do the init in background
         runtime::spawn(self.late_init())?;
         Ok(())
@@ -239,7 +242,7 @@ pci_driver! {
             "ahci"
         }
 
-        fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
+        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
             let bar5 = info.config_space.bar(5).ok_or(Error::InvalidOperation)?;
             let bar5 = bar5.as_memory().ok_or(Error::InvalidOperation)?;
 
@@ -268,6 +271,7 @@ pci_driver! {
 
             let ahci = Arc::new(AhciController {
                 regs: IrqSafeSpinlock::new(regs),
+                dma: dma.clone(),
                 ports: OneTimeInit::new(),
                 received_fis_buffers: OneTimeInit::new(),
                 version,

kernel/driver/block/ahci/src/port.rs

@@ -8,14 +8,15 @@ use core::{
 use alloc::{boxed::Box, string::String, sync::Arc};
 use async_trait::async_trait;
 use bytemuck::Zeroable;
-use device_api::device::Device;
+use device_api::{device::Device, dma::DmaAllocator};
 use futures_util::task::AtomicWaker;
-use libk::{device::block::BlockDevice, error::Error};
+use libk::{
+    device::block::BlockDevice,
+    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
+    error::Error,
+};
 use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    device::DeviceMemoryIo,
-    table::MapAttributes,
-    PageBox, PageProvider, PageSlice,
+    address::PhysicalAddress, device::DeviceMemoryIo, table::MapAttributes, PageProvider,
 };
 use libk_util::{sync::IrqSafeSpinlock, waker::QueueWaker, OneTimeInit};
 use tock_registers::interfaces::{Readable, Writeable};
@@ -37,8 +38,8 @@ struct PortInner {
     regs: DeviceMemoryIo<'static, PortRegs>,
 
     #[allow(unused)]
-    received_fis: PageBox<ReceivedFis>,
-    command_list: PageBox<[CommandListEntry]>,
+    received_fis: DmaBuffer<ReceivedFis>,
+    command_list: DmaBuffer<[CommandListEntry]>,
 }
 
 pub struct PortInfo {
@@ -90,18 +91,16 @@ impl Drop for SubmittedCommand<'_> {
 impl PortInner {
     fn submit_command<C: AtaCommand>(
         &mut self,
+        dma: &dyn DmaAllocator,
         index: usize,
         command: &C,
     ) -> Result<(), AhciError> {
         let list_entry = &mut self.command_list[index];
         let mut table_entry =
-            PageBox::new(CommandTable::zeroed()).map_err(AhciError::MemoryError)?;
+            DmaBuffer::new(dma, CommandTable::zeroed()).map_err(AhciError::MemoryError)?;
 
         table_entry.setup_command(command)?;
-        *list_entry = CommandListEntry::new(
-            unsafe { table_entry.as_physical_address() },
-            command.prd_count(),
-        )?;
+        *list_entry = CommandListEntry::new(table_entry.bus_address(), command.prd_count())?;
 
         // Sync before send
         // XXX do this properly
@@ -137,12 +136,14 @@ impl AhciPort {
             return Err(AhciError::DeviceError);
         }
 
-        let received_fis = PageBox::new(ReceivedFis::zeroed()).map_err(AhciError::MemoryError)?;
-        let command_list = PageBox::new_slice(CommandListEntry::zeroed(), COMMAND_LIST_LENGTH)
+        let received_fis =
+            DmaBuffer::new(&*ahci.dma, ReceivedFis::zeroed()).map_err(AhciError::MemoryError)?;
+        let command_list =
+            DmaBuffer::new_slice(&*ahci.dma, CommandListEntry::zeroed(), COMMAND_LIST_LENGTH)
                 .map_err(AhciError::MemoryError)?;
 
-        regs.set_received_fis_address_64(unsafe { received_fis.as_physical_address() });
-        regs.set_command_list_address_64(unsafe { command_list.as_physical_address() });
+        regs.set_received_fis_address_64(received_fis.bus_address());
+        regs.set_command_list_address_64(command_list.bus_address());
 
         regs.IE.write(
             IE::DPE::SET
@@ -182,7 +183,9 @@ impl AhciPort {
     }
 
     pub async fn init_inner(&self) -> Result<(), AhciError> {
-        let identify = self.perform_command(AtaIdentify::create()?).await?;
+        let identify = self
+            .perform_command(AtaIdentify::create(&*self.ahci.dma)?)
+            .await?;
 
         let model = identify.model_number.to_string();
         let serial = identify.serial_number.to_string();
@@ -237,7 +240,11 @@ impl AhciPort {
             return Err(AhciError::RegionTooLarge);
         }
         let index = self.allocate_command().await;
-        if let Err(error) = self.inner.lock().submit_command(index, command) {
+        if let Err(error) = self
+            .inner
+            .lock()
+            .submit_command(&*self.ahci.dma, index, command)
+        {
             self.free_command(index);
             return Err(error);
         }
@@ -302,28 +309,38 @@ impl AhciPort {
 
 #[async_trait]
 impl BlockDevice for AhciPort {
+    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        DmaBuffer::new_uninit_slice(&*self.ahci.dma, size)
+    }
+
     async fn read_aligned(
         &self,
         position: u64,
-        buffer: &mut PageSlice<MaybeUninit<u8>>,
+        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
     ) -> Result<(), Error> {
-        if position % SECTOR_SIZE as u64 != 0 {
+        if buffer.len() % SECTOR_SIZE != 0 {
+            log::warn!("ahci: misaligned buffer size: {}", buffer.len());
             return Err(Error::InvalidOperation);
         }
-        if buffer.len() % SECTOR_SIZE != 0 {
+        if position % SECTOR_SIZE as u64 != 0 {
+            log::warn!("ahci: misaligned read");
             return Err(Error::InvalidOperation);
         }
 
         let lba = position / SECTOR_SIZE as u64;
-        let command = AtaReadDmaEx::new(lba, buffer.len() / SECTOR_SIZE, buffer);
-        self.submit(&command)
-            .await?
-            .wait_for_completion()
-            .await
-            .map_err(AhciError::into)
+        let lba_count = buffer.len() / SECTOR_SIZE;
+        if lba + lba_count as u64 >= self.block_count() {
+            log::warn!("ahci: read crosses medium end");
+            return Err(Error::InvalidOperation);
         }
 
-    async fn write_aligned(&self, _position: u64, _buffer: &PageSlice<u8>) -> Result<(), Error> {
+        let command = AtaReadDmaEx::new(lba, lba_count, buffer);
+        self.submit(&command).await?.wait_for_completion().await?;
+
+        Ok(())
+    }
+
+    async fn write_aligned(&self, _position: u64, _buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
         // TODO AtaWriteDmaEx
         Err(Error::NotImplemented)
     }

kernel/driver/block/ahci/src/regs.rs

@@ -1,4 +1,4 @@
-use libk_mm::address::PhysicalAddress;
+use libk::dma::BusAddress;
 use tock_registers::{
     interfaces::{ReadWriteable, Readable, Writeable},
     register_bitfields, register_structs,
@@ -141,14 +141,14 @@ impl PortRegs {
         Ok(())
     }
 
-    pub fn set_received_fis_address_64(&self, address: PhysicalAddress) {
-        let address: u64 = address.into();
+    pub fn set_received_fis_address_64(&self, address: BusAddress) {
+        let address: u64 = address.into_u64();
         self.FB.set(address as u32);
         self.FBU.set((address >> 32) as u32);
     }
 
-    pub fn set_command_list_address_64(&self, address: PhysicalAddress) {
-        let address: u64 = address.into();
+    pub fn set_command_list_address_64(&self, address: BusAddress) {
+        let address: u64 = address.into_u64();
         self.CLB.set(address as u32);
         self.CLBU.set((address >> 32) as u32);
     }

kernel/driver/block/nvme/src/command.rs

@@ -2,7 +2,7 @@
 
 use core::fmt::{self, Write};
 
-use libk_mm::address::PhysicalAddress;
+use libk::dma::BusAddress;
 use tock_registers::{interfaces::Readable, register_structs, registers::ReadOnly, UIntLike};
 
 use crate::queue::PhysicalRegionPage;
@@ -74,7 +74,7 @@ pub struct CreateIoCompletionQueue {
     pub id: u32,
     pub size: usize,
     pub vector: u32,
-    pub data: PhysicalAddress,
+    pub data: BusAddress,
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -82,7 +82,7 @@ pub struct CreateIoSubmissionQueue {
     pub id: u32,
     pub cq_id: u32,
     pub size: usize,
-    pub data: PhysicalAddress,
+    pub data: BusAddress,
 }
 
 // Replies

kernel/driver/block/nvme/src/drive.rs

@@ -3,7 +3,11 @@ use core::mem::MaybeUninit;
 use alloc::{boxed::Box, sync::Arc};
 use async_trait::async_trait;
 use device_api::device::Device;
-use libk::{device::block::BlockDevice, error::Error};
+use libk::{
+    device::block::BlockDevice,
+    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
+    error::Error,
+};
 use libk_mm::{
     address::{AsPhysicalAddress, PhysicalAddress},
     table::MapAttributes,
@@ -30,7 +34,9 @@ impl NvmeNamespace {
         max_transfer_size: usize,
     ) -> Result<Arc<NvmeNamespace>, NvmeError> {
         let admin_q = controller.admin_q.get();
-        let identify = admin_q.request(IdentifyNamespaceRequest { nsid }).await?;
+        let identify = admin_q
+            .request(&*controller.dma, IdentifyNamespaceRequest { nsid })
+            .await?;
 
         let current_lba_format_idx = identify.current_lba_fmt_idx();
         let current_lba_format = identify.lba_fmt(current_lba_format_idx).unwrap();
@@ -76,17 +82,27 @@ impl Device for NvmeNamespace {
 
 #[async_trait]
 impl BlockDevice for NvmeNamespace {
+    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        DmaBuffer::new_uninit_slice(&*self.controller.dma, size)
+    }
+
+    // TODO read directly to cache
     async fn read_aligned(
         &self,
         position: u64,
-        buffer: &mut PageSlice<MaybeUninit<u8>>,
+        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
     ) -> Result<(), Error> {
-        debug_assert_eq!(position % self.block_size() as u64, 0);
+        if position % self.block_size() as u64 != 0 {
+            return Err(Error::InvalidOperation);
+        }
+        if buffer.len() % self.block_size() != 0 {
+            return Err(Error::InvalidOperation);
+        }
         let lba = position / self.block_size() as u64;
-        debug_assert_eq!(buffer.len() % self.block_size(), 0);
-        let buffer_address = unsafe { buffer.as_physical_address() };
-        debug_assert_eq!(buffer_address.into_u64() % self.block_size() as u64, 0);
         let lba_count = buffer.len() / self.block_size();
+        if lba + lba_count as u64 > self.block_count() {
+            return Err(Error::InvalidOperation);
+        }
 
         let result = self
             .controller
@@ -94,24 +110,29 @@ impl BlockDevice for NvmeNamespace {
                 self.nsid,
                 lba,
                 lba_count,
-                buffer_address,
+                buffer.bus_address(),
                 buffer.len(),
                 IoDirection::Read,
             )
             .await;
 
-        log::info!(target: "io", "read #{lba}, {lba_count} blocks -> {result:?} @ {buffer_address:#x}");
+        log::info!("read #{lba}, {lba_count} blocks -> {result:?} @ {buffer:p}");
 
         result.map_err(NvmeError::into)
     }
 
-    async fn write_aligned(&self, position: u64, buffer: &PageSlice<u8>) -> Result<(), Error> {
-        debug_assert_eq!(position % self.block_size() as u64, 0);
+    async fn write_aligned(&self, position: u64, buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
+        if position % self.block_size() as u64 != 0 {
+            return Err(Error::InvalidOperation);
+        }
+        if buffer.len() % self.block_size() != 0 {
+            return Err(Error::InvalidOperation);
+        }
         let lba = position / self.block_size() as u64;
-        debug_assert_eq!(buffer.len() % self.block_size(), 0);
-        let buffer_address = unsafe { buffer.as_physical_address() };
-        debug_assert_eq!(buffer_address.into_u64() % self.block_size() as u64, 0);
         let lba_count = buffer.len() / self.block_size();
+        if lba + lba_count as u64 > self.block_count() {
+            return Err(Error::InvalidOperation);
+        }
 
         // TODO ArchitectureImpl::flush_data_cache()
         #[cfg(target_arch = "x86_64")]
@@ -125,13 +146,13 @@ impl BlockDevice for NvmeNamespace {
                 self.nsid,
                 lba,
                 lba_count,
-                buffer_address,
+                buffer.bus_address(),
                 buffer.len(),
                 IoDirection::Write,
             )
             .await;
 
-        log::info!(target: "io", "write -> #{lba}, {lba_count} blocks -> {result:?} @ {buffer_address:#x}");
+        log::info!(target: "io", "write -> #{lba}, {lba_count} blocks -> {result:?} @ {buffer:p}");
 
         result.map_err(NvmeError::into)
     }

kernel/driver/block/nvme/src/lib.rs

@@ -15,12 +15,14 @@ use core::{
 use alloc::{collections::BTreeMap, format, sync::Arc, vec::Vec};
 use command::{IdentifyActiveNamespaceIdListRequest, IdentifyControllerRequest};
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
+    dma::DmaAllocator,
     interrupt::{InterruptAffinity, InterruptHandler, IrqVector},
 };
 use drive::NvmeNamespace;
 use libk::{
     device::manager::probe_partitions,
+    dma::BusAddress,
     fs::devfs,
     task::{cpu_count, cpu_index, runtime},
 };
@@ -137,6 +139,7 @@ pub struct NvmeController {
     controller_id: OneTimeInit<u32>,
 
     pci: PciDeviceInfo,
+    dma: Arc<dyn DmaAllocator>,
 
     doorbell_shift: usize,
     min_page_size: usize,
@@ -183,7 +186,14 @@ impl NvmeController {
             let id = i as u32;
 
             let (sq_doorbell, cq_doorbell) = unsafe { self.doorbell_pair(i) };
-            let queue = QueuePair::new(id, i, Self::IO_QUEUE_SIZE, sq_doorbell, cq_doorbell)
+            let queue = QueuePair::new(
+                &*self.dma,
+                id,
+                i,
+                Self::IO_QUEUE_SIZE,
+                sq_doorbell,
+                cq_doorbell,
+            )
             .map_err(NvmeError::MemoryError)?;
 
             admin_q
@@ -191,7 +201,7 @@ impl NvmeController {
                     id,
                     vector: id,
                     size: Self::IO_QUEUE_SIZE,
-                    data: queue.cq_physical_pointer(),
+                    data: queue.cq_bus_pointer(),
                 })
                 .await?;
 
@@ -200,7 +210,7 @@ impl NvmeController {
                     id,
                     cq_id: id,
                     size: Self::IO_QUEUE_SIZE,
-                    data: queue.sq_physical_pointer(),
+                    data: queue.sq_bus_pointer(),
                 })
                 .await?;
 
@@ -233,7 +243,9 @@ impl NvmeController {
         let admin_q = self.admin_q.get();
 
         // Identify the controller
-        let identify = admin_q.request(IdentifyControllerRequest).await?;
+        let identify = admin_q
+            .request(&*self.dma, IdentifyControllerRequest)
+            .await?;
 
         let max_transfer_size = if identify.mdts == 0 {
             // Pick some sane default value
@@ -257,7 +269,10 @@ impl NvmeController {
         let admin_q = self.admin_q.get();
 
         let namespaces = admin_q
-            .request(IdentifyActiveNamespaceIdListRequest { start_id: 0 })
+            .request(
+                &*self.dma,
+                IdentifyActiveNamespaceIdListRequest { start_id: 0 },
+            )
             .await?;
 
         let count = namespaces.entries.iter().position(|&x| x == 0).unwrap();
@@ -282,11 +297,11 @@ impl NvmeController {
         nsid: u32,
         lba: u64,
         lba_count: usize,
-        buffer_address: PhysicalAddress,
+        buffer_address: BusAddress,
         transfer_size: usize,
         direction: IoDirection,
     ) -> Result<(), NvmeError> {
-        let prp_list = PrpList::from_buffer(buffer_address, transfer_size)?;
+        let prp_list = PrpList::from_buffer(&*self.dma, buffer_address, transfer_size)?;
 
         let _guard = IrqGuard::acquire();
         let cpu_index = cpu_index();
@@ -345,7 +360,7 @@ impl InterruptHandler for NvmeController {
 }
 
 impl Device for NvmeController {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let regs = self.regs.lock();
 
         let timeout = Duration::from_millis(regs.CAP.read(CAP::TO) * 500);
@@ -367,6 +382,7 @@ impl Device for NvmeController {
         let admin_cq_doorbell = unsafe { regs.doorbell_ptr(self.doorbell_shift, true, 0) };
         log::debug!("sq_doorbell for adminq = {:p}", admin_sq_doorbell);
         let admin_q = QueuePair::new(
+            &*self.dma,
             0,
             0,
             Self::ADMIN_QUEUE_SIZE,
@@ -379,8 +395,8 @@ impl Device for NvmeController {
             AQA::ASQS.val(Self::ADMIN_QUEUE_SIZE as u32 - 1)
                 + AQA::ACQS.val(Self::ADMIN_QUEUE_SIZE as u32 - 1),
         );
-        regs.ASQ.set(admin_q.sq_physical_pointer().into());
-        regs.ACQ.set(admin_q.cq_physical_pointer().into());
+        regs.ASQ.set(admin_q.sq_bus_pointer().into_u64());
+        regs.ACQ.set(admin_q.cq_bus_pointer().into_u64());
 
         // Configure the controller
         const IOSQES: u32 = size_of::<SubmissionQueueEntry>().ilog2();
@@ -464,7 +480,7 @@ pci_driver! {
             "nvme"
         }
 
-        fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
+        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
             let bar0 = info
                 .config_space
                 .bar(0)
@@ -500,6 +516,7 @@ pci_driver! {
                 controller_id: OneTimeInit::new(),
 
                 pci: info.clone(),
+                dma: dma.clone(),
 
                 io_queue_count: AtomicUsize::new(1),
                 doorbell_shift,

kernel/driver/block/nvme/src/queue.rs

@@ -2,10 +2,9 @@ use core::{future::poll_fn, mem::size_of, ptr::null_mut, task::Poll};
 
 use alloc::collections::{BTreeMap, BTreeSet};
 use bytemuck::{Pod, Zeroable};
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox,
-};
+use device_api::dma::DmaAllocator;
+use libk::dma::{BusAddress, DmaBuffer};
+use libk_mm::address::AsPhysicalAddress;
 use libk_util::{sync::IrqSafeSpinlock, waker::QueueWaker};
 use static_assertions::const_assert;
 use yggdrasil_abi::error::Error;
@@ -58,7 +57,7 @@ pub struct CompletionQueueEntry {
 }
 
 pub struct Queue<T> {
-    data: PageBox<[T]>,
+    data: DmaBuffer<[T]>,
     mask: usize,
     head: usize,
     tail: usize,
@@ -82,8 +81,8 @@ pub struct QueuePair {
     #[allow(unused)]
     vector: usize,
 
-    sq_base: PhysicalAddress,
-    cq_base: PhysicalAddress,
+    sq_base: BusAddress,
+    cq_base: BusAddress,
 
     pub completion_notify: QueueWaker,
 
@@ -94,7 +93,7 @@ pub struct PrpList {
     prp1: PhysicalRegionPage,
     prp2: PhysicalRegionPage,
     #[allow(unused)]
-    list: Option<PageBox<[PhysicalAddress]>>,
+    list: Option<DmaBuffer<[BusAddress]>>,
 }
 
 impl PrpList {
@@ -106,7 +105,11 @@ impl PrpList {
         }
     }
 
-    pub fn from_buffer(base: PhysicalAddress, size: usize) -> Result<Self, NvmeError> {
+    pub fn from_buffer(
+        dma: &dyn DmaAllocator,
+        base: BusAddress,
+        size: usize,
+    ) -> Result<Self, NvmeError> {
         // TODO hardcoded page size
         if base.into_u64() % 0x1000 != 0 {
             todo!();
@@ -126,12 +129,13 @@ impl PrpList {
             }),
             _ => {
                 let count = (size + 0xFFF) / 0x1000;
-                let list = PageBox::new_slice_with(|i| base.add((i + 1) * 0x1000), count - 1)
+                let list =
+                    DmaBuffer::new_slice_with(dma, |i| base.add((i + 1) * 0x1000), count - 1)
                         .map_err(NvmeError::MemoryError)?;
 
                 Ok(Self {
                     prp1: PhysicalRegionPage::with_addr(base),
-                    prp2: PhysicalRegionPage::with_addr(unsafe { list.as_physical_address() }),
+                    prp2: PhysicalRegionPage::with_addr(list.bus_address()),
                     list: Some(list),
                 })
             }
@@ -146,7 +150,7 @@ impl PhysicalRegionPage {
         Self(0)
     }
 
-    pub const fn with_addr(address: PhysicalAddress) -> Self {
+    pub const fn with_addr(address: BusAddress) -> Self {
         Self(address.into_u64())
     }
 }
@@ -197,7 +201,7 @@ impl CompletionQueueEntry {
 
 impl<T> Queue<T> {
     pub fn new(
-        data: PageBox<[T]>,
+        data: DmaBuffer<[T]>,
         head_doorbell: *mut u32,
         tail_doorbell: *mut u32,
         phase: bool,
@@ -277,17 +281,18 @@ impl<T> Queue<T> {
 
 impl QueuePair {
     pub fn new(
+        dma: &dyn DmaAllocator,
         id: u32,
         vector: usize,
         capacity: usize,
         sq_doorbell: *mut u32,
         cq_doorbell: *mut u32,
     ) -> Result<Self, Error> {
-        let sq_data = PageBox::new_slice(SubmissionQueueEntry::zeroed(), capacity)?;
-        let cq_data = PageBox::new_slice(CompletionQueueEntry::zeroed(), capacity)?;
+        let sq_data = DmaBuffer::new_slice(dma, SubmissionQueueEntry::zeroed(), capacity)?;
+        let cq_data = DmaBuffer::new_slice(dma, CompletionQueueEntry::zeroed(), capacity)?;
 
-        let sq_base = unsafe { sq_data.as_physical_address() };
-        let cq_base = unsafe { cq_data.as_physical_address() };
+        let sq_base = sq_data.bus_address();
+        let cq_base = cq_data.bus_address();
 
         log::debug!("Allocated queue pair: sq={:p}, cq={:p}", sq_data, cq_data);
 
@@ -313,12 +318,12 @@ impl QueuePair {
     }
 
     #[inline]
-    pub fn sq_physical_pointer(&self) -> PhysicalAddress {
+    pub fn sq_bus_pointer(&self) -> BusAddress {
         self.sq_base
     }
 
     #[inline]
-    pub fn cq_physical_pointer(&self) -> PhysicalAddress {
+    pub fn cq_bus_pointer(&self) -> BusAddress {
         self.cq_base
     }
 
@@ -385,16 +390,17 @@ impl QueuePair {
 
     pub async fn request<'r, R: Request>(
         &'r self,
+        dma: &dyn DmaAllocator,
         req: R,
-    ) -> Result<PageBox<R::Response>, NvmeError>
+    ) -> Result<DmaBuffer<R::Response>, NvmeError>
     where
         R::Response: 'r,
     {
-        let response = PageBox::new_uninit().map_err(NvmeError::MemoryError)?;
-        let list = PrpList::from_buffer(unsafe { response.as_physical_address() }, size_of::<R>())?;
+        let response = DmaBuffer::new_uninit(dma).map_err(NvmeError::MemoryError)?;
+        let list = PrpList::from_buffer(dma, response.bus_address(), size_of::<R>())?;
         let command_id = self.submit(req, &list, true)?;
         let result = self.wait_for_completion(command_id, response).await?;
-        Ok(unsafe { result.assume_init() })
+        Ok(unsafe { DmaBuffer::assume_init(result) })
     }
 
     pub fn process_completions(&self) -> usize {

kernel/driver/block/scsi/src/lib.rs

@@ -2,7 +2,11 @@
 #![allow(incomplete_features)]
 #![no_std]
 
-use core::{mem::MaybeUninit, time::Duration};
+use core::{
+    mem::MaybeUninit,
+    sync::atomic::{AtomicBool, Ordering},
+    time::Duration,
+};
 
 use alloc::{
     boxed::Box, collections::btree_map::BTreeMap, format, string::String, sync::Arc, vec::Vec,
@@ -11,12 +15,14 @@ use async_trait::async_trait;
 use command::{ScsiReadCapacity, ScsiRequestSense, ScsiTestUnitReady};
 use device_api::device::Device;
 use libk::{
+    block,
     device::{block::BlockDevice, manager::probe_partitions},
+    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
     error::Error,
     fs::devfs,
-    task::runtime,
+    task::{runtime, sync::AsyncMutex},
 };
-use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageProvider, PageSlice};
+use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageProvider};
 use libk_util::{
     sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
     OneTimeInit,
@@ -30,26 +36,69 @@ pub mod command;
 pub mod device;
 pub mod transport;
 
-// TODO SCSI detach
-pub struct ScsiDevice {
-    transport: IrqSafeSpinlock<ScsiTransportWrapper>,
+pub struct ScsiEnclosure {
+    transport: AsyncMutex<ScsiTransportWrapper>,
+    units: Vec<IrqSafeRwLock<Option<Arc<ScsiUnit>>>>,
+    index: OneTimeInit<u32>,
+    shutdown: AtomicBool,
+}
+
+pub struct ScsiUnit {
+    enclosure: Arc<ScsiEnclosure>,
+    lun: u8,
     lba_count: u64,
     lba_size: usize,
-    index: OneTimeInit<u32>,
+    max_lba_per_request: usize,
     names: IrqSafeRwLock<Vec<String>>,
 }
 
-impl ScsiDevice {
-    // TODO support LUNs other than 0
-    pub async fn setup<T: ScsiTransport + 'static>(transport: T) -> Result<Arc<Self>, Error> {
-        let mut transport = ScsiTransportWrapper::new(transport);
+impl ScsiEnclosure {
+    pub async fn setup(
+        transport: Box<dyn ScsiTransport>,
+        lun_count: usize,
+    ) -> Result<Arc<Self>, Error> {
+        let transport = AsyncMutex::new(ScsiTransportWrapper::new(transport));
+        let units = (0..lun_count).map(|_| IrqSafeRwLock::new(None)).collect();
+        let this = Arc::new(Self {
+            transport,
+            units,
+            index: OneTimeInit::new(),
+            shutdown: AtomicBool::new(false),
+        });
+        register_enclosure(this.clone())?;
 
-        let mut attempts = 5;
-        let mut timeout = 100;
+        // Probe LUNs
+        for i in 0..lun_count {
+            if this.probe_lun(i as u8).await {
+                if let Ok(unit) = ScsiUnit::setup(this.clone(), i as u8).await {
+                    *this.units[i].write() = Some(unit);
+                }
+            }
+        }
+
+        // Start enclosure poll task
+        let enclosure = this.clone();
+        runtime::spawn(async move {
+            while !enclosure.shutdown.load(Ordering::Acquire) {
+                enclosure.poll().await;
+                runtime::sleep(Duration::from_millis(100)).await;
+            }
+        })
+        .ok();
+
+        Ok(this)
+    }
+
+    async fn probe_lun(self: &Arc<Self>, lun: u8) -> bool {
+        let mut attempts = 3;
+        let mut timeout = 10;
+        // TODO get statuses to better see if there's a real error or the LUN is not present
         while attempts > 0 {
+            let mut transport = self.transport.lock().await;
+
             // TEST UNIT READY (6)
             if transport
-                .perform_command(0, ScsiTestUnitReady)
+                .perform_command(lun, ScsiTestUnitReady)
                 .await
                 .is_ok()
             {
@@ -57,82 +106,161 @@ impl ScsiDevice {
             }
 
             // If not, send a REQUEST SENSE (6)
-            if transport.perform_command(0, ScsiRequestSense).await.is_ok() {
-                break;
-            }
+            transport.perform_command(lun, ScsiRequestSense).await.ok();
 
-            log::warn!("scsi: unit not ready [{attempts}/5]");
+            drop(transport);
 
             runtime::sleep(Duration::from_millis(timeout)).await;
             timeout *= 2;
             attempts -= 1;
         }
+
         if attempts == 0 {
-            log::error!("scsi: unit not ready");
-            return Err(Error::DoesNotExist);
+            false
+        } else {
+            true
         }
+    }
+
+    async fn poll(self: &Arc<Self>) {
+        let index = *self.index.get();
+        for lun in 0..self.units.len() {
+            let mut slot = self.units[lun].write();
+            let present = self.probe_lun(lun as u8).await;
+
+            if let Some(unit) = slot.as_ref() {
+                if !present {
+                    log::warn!("scsi{index}u{lun} lost");
+                    unit.detach();
+                    *slot = None;
+                }
+            } else if present {
+                if let Ok(unit) = ScsiUnit::setup(self.clone(), lun as u8).await {
+                    log::info!("scsi{index}u{lun} attached");
+                    *slot = Some(unit);
+                } else {
+                    log::warn!("scsi{index}u{lun} attached, but could not setup");
+                }
+            }
+        }
+    }
+
+    pub fn detach(&self) {
+        self.shutdown.store(true, Ordering::Release);
+        let index = self.index.try_get().copied();
+
+        for unit in self.units.iter() {
+            if let Some(unit) = unit.write().take() {
+                unit.detach();
+            }
+        }
+
+        // Deregister the enclosure
+        if let Some(index) = index {
+            remove_enclosure(index);
+        }
+    }
+}
+
+impl ScsiUnit {
+    pub async fn setup(enclosure: Arc<ScsiEnclosure>, lun: u8) -> Result<Arc<Self>, Error> {
+        let enclosure_index = *enclosure.index.get();
+        let mut transport = enclosure.transport.lock().await;
 
         // TODO INQUIRY fails for real USB flash drives
         // transport.perform_command(0, ScsiInquiry).await?;
 
-        let capacity_info = transport.perform_command(0, ScsiReadCapacity).await?;
+        let capacity_info = transport.perform_command(lun, ScsiReadCapacity).await?;
+        let max_lba_per_request =
+            transport.max_bytes_per_request() / capacity_info.block_size as usize;
         log::info!(
-            "scsi: lba_size={}, lba_count={}",
+            "scsi{enclosure_index}u{lun}: lba_size={}, lba_count={}, max_lba_per_request={}",
             capacity_info.block_size,
-            capacity_info.block_count
+            capacity_info.block_count,
+            max_lba_per_request
         );
 
-        Ok(Arc::new(Self {
-            transport: IrqSafeSpinlock::new(transport),
+        drop(transport);
+
+        let unit = Arc::new(Self {
+            enclosure,
+            lun,
             lba_count: capacity_info.block_count.into(),
             lba_size: capacity_info.block_size as usize,
-            index: OneTimeInit::new(),
+            max_lba_per_request,
             names: IrqSafeRwLock::new(Vec::new()),
-        }))
+        });
+
+        register_unit(enclosure_index, lun, unit.clone());
+
+        Ok(unit)
     }
 
-    pub fn detach(&self) {
-        if let Some(&index) = self.index.try_get() {
-            detach(index);
+    fn detach(&self) {
+        let id = *self.enclosure.index.get();
+        log::info!("scsi{id}u{} detached", self.lun);
+        for name in self.names.read().iter() {
+            devfs::remove_node(name).ok();
         }
     }
 }
 
 #[async_trait]
-impl BlockDevice for ScsiDevice {
+impl BlockDevice for ScsiUnit {
+    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        block!(self.enclosure.transport.lock().await.allocate_buffer(size))?
+    }
+
     async fn read_aligned(
         &self,
         position: u64,
-        buffer: &mut PageSlice<MaybeUninit<u8>>,
+        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
     ) -> Result<(), Error> {
+        if position % self.lba_size as u64 != 0 {
+            log::warn!("scsi: misaligned read");
+            return Err(Error::InvalidArgument);
+        }
         if buffer.len() % self.lba_size != 0 {
+            log::warn!("scsi: misaligned buffer size");
             return Err(Error::InvalidArgument);
         }
         let lba_start = position / self.lba_size as u64;
         let lba_count = buffer.len() / self.lba_size;
+
         if lba_start.saturating_add(lba_count as u64) >= self.lba_count {
+            log::warn!("scsi: read beyond medium end");
             return Err(Error::InvalidArgument);
         }
-        log::info!("scsi: read lba={lba_start}, count={lba_count}");
-        let mut transport = self.transport.lock();
-        for i in 0..lba_count {
-            let lba = lba_start + i as u64;
-            let offset = self.lba_size * i;
-            let slice = unsafe {
-                MaybeUninit::slice_assume_init_mut(&mut buffer[offset..offset + self.lba_size])
-            };
 
-            let len = transport.read(0, lba, slice).await?;
-            if len != self.lba_size {
-                log::warn!("scsi: truncated read received at lba {lba}");
-                return Err(Error::InvalidOperation);
+        let lba_end = lba_start + lba_count as u64;
+
+        let mut transport = self.enclosure.transport.lock().await;
+
+        // TODO DmaSliceMut subslicing
+        let (buffer, range) = buffer.into_parts();
+        let mut offset = range.start;
+
+        for i in (0..lba_count).step_by(self.max_lba_per_request) {
+            let lba = lba_start + i as u64;
+            let end = (lba + self.max_lba_per_request as u64).min(lba_end);
+            let count = (end - lba) as usize;
+            let amount = count * self.lba_size;
+
+            let dst_slice = buffer.slice_mut(offset..offset + amount);
+            let len = transport
+                .read(self.lun, lba, count as u16, dst_slice)
+                .await?;
+            if len != amount {
+                return Err(Error::InvalidArgument);
             }
+
+            offset += amount;
         }
+
         Ok(())
     }
 
-    async fn write_aligned(&self, _position: u64, _buffer: &PageSlice<u8>) -> Result<(), Error> {
-        // TODO AtaWriteDmaEx
+    async fn write_aligned(&self, _position: u64, _buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
         Err(Error::NotImplemented)
     }
 
@@ -145,11 +273,11 @@ impl BlockDevice for ScsiDevice {
     }
 
     fn max_blocks_per_request(&self) -> usize {
-        65536 / self.lba_size
+        self.max_lba_per_request
     }
 }
 
-impl PageProvider for ScsiDevice {
+impl PageProvider for ScsiUnit {
     fn get_page(&self, _offset: u64) -> Result<PhysicalAddress, Error> {
         Err(Error::NotImplemented)
     }
@@ -168,53 +296,57 @@ impl PageProvider for ScsiDevice {
     }
 }
 
-impl Device for ScsiDevice {
+impl Device for ScsiUnit {
     fn display_name(&self) -> &str {
-        "SCSI Storage Device"
+        "SCSI Unit"
     }
 }
 
-impl Drop for ScsiDevice {
+impl Drop for ScsiUnit {
     fn drop(&mut self) {
-        if let Some(index) = self.index.try_get() {
-            log::info!("scsi{index} dropped");
+        if let Some(index) = self.enclosure.index.try_get() {
+            log::info!("scsi{index}u{} dropped", self.lun);
         }
     }
 }
 
 // TODO this is crap
-static SCSI_DEVICES: IrqSafeSpinlock<BTreeMap<u32, Arc<ScsiDevice>>> =
+static SCSI_ENCLOSURES: IrqSafeSpinlock<BTreeMap<u32, Arc<ScsiEnclosure>>> =
     IrqSafeSpinlock::new(BTreeMap::new());
 static SCSI_BITMAP: IrqSafeSpinlock<u32> = IrqSafeSpinlock::new(0);
 
-pub fn attach(device: Arc<ScsiDevice>) -> Result<(), Error> {
+fn register_enclosure(enclosure: Arc<ScsiEnclosure>) -> Result<(), Error> {
     let index = {
         let mut bitmap = SCSI_BITMAP.lock();
         let index = (0..8)
             .position(|p| *bitmap & (1 << p) == 0)
             .ok_or(Error::InvalidOperation)
-            .inspect_err(|_| log::warn!("Cannot attach SCSI device: too many of them"))?
+            .inspect_err(|_| log::warn!("Cannot attach SCSI enclosure: too many of them"))?
             as u32;
-        let mut devices = SCSI_DEVICES.lock();
+        let mut devices = SCSI_ENCLOSURES.lock();
         *bitmap |= 1 << index;
         assert!(!devices.contains_key(&index));
-        devices.insert(index, device.clone());
+        devices.insert(index, enclosure.clone());
         index
     };
 
-    let name = format!("scsi{index}");
-    device.index.init(index);
-    device.names.write().push(name.clone());
-    devfs::add_named_block_device(device.clone(), name.clone(), FileMode::new(0o600)).ok();
-    log::info!("{name} attached");
+    enclosure.index.init(index);
+
+    Ok(())
+}
+
+fn register_unit(enclosure_index: u32, lun: u8, unit: Arc<ScsiUnit>) {
+    let name = format!("scsi{enclosure_index}u{lun}");
+    unit.names.write().push(name.clone());
+    devfs::add_named_block_device(unit.clone(), name.clone(), FileMode::new(0o600)).ok();
 
     // TODO this code is repeated everywhere
     runtime::spawn(async move {
         let name = name;
-        probe_partitions(device.clone(), |index, partition| {
+        probe_partitions(unit.clone(), |index, partition| {
             let partition_name = format!("{name}p{}", index + 1);
             log::info!("{name}: partition {partition_name}");
-            device.names.write().push(partition_name.clone());
+            unit.names.write().push(partition_name.clone());
             devfs::add_named_block_device(
                 Arc::new(partition),
                 partition_name,
@@ -226,23 +358,13 @@ pub fn attach(device: Arc<ScsiDevice>) -> Result<(), Error> {
         .ok();
     })
     .ok();
-
-    Ok(())
 }
 
-pub fn detach(index: u32) {
-    let mut devices = SCSI_DEVICES.lock();
+fn remove_enclosure(index: u32) {
+    let mut devices = SCSI_ENCLOSURES.lock();
     let mut bitmap = SCSI_BITMAP.lock();
 
-    if let Some(device) = devices.remove(&index) {
-        {
-            let names = device.names.read();
-            for name in names.iter() {
-                devfs::remove_node(name).ok();
-            }
-        }
-
     *bitmap &= !(1 << index);
-        log::info!("scsi{index} detached");
-    }
+    devices.remove(&index);
+    log::info!("scsi: enclosure {index} detached");
 }

kernel/driver/block/scsi/src/transport.rs

@@ -1,18 +1,27 @@
+use core::{mem::MaybeUninit, ops::Deref};
+
 use alloc::boxed::Box;
 use async_trait::async_trait;
-use libk::error::Error;
+use libk::{
+    dma::{DmaBuffer, DmaSliceMut},
+    error::Error,
+};
 
 use crate::command::ScsiCommand;
 
 #[async_trait]
 pub trait ScsiTransport: Send + Sync {
+    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error>;
+
     /// Perform a no-data request
     async fn perform_request_raw(
         &mut self,
         lun: u8,
         request_data: &[u8],
-        response_buffer: &mut [u8],
+        response_buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
     ) -> Result<usize, Error>;
+
+    fn max_bytes_per_request(&self) -> usize;
 }
 
 pub struct ScsiTransportWrapper {
@@ -20,17 +29,22 @@ pub struct ScsiTransportWrapper {
 }
 
 impl ScsiTransportWrapper {
-    pub fn new<T: ScsiTransport + 'static>(inner: T) -> Self {
-        Self {
-            inner: Box::new(inner),
-        }
+    pub fn new(inner: Box<dyn ScsiTransport>) -> Self {
+        Self { inner }
     }
 
-    pub async fn read(&mut self, lun: u8, lba: u64, buffer: &mut [u8]) -> Result<usize, Error> {
-        if lba >= u32::MAX as u64 || buffer.len() > u16::MAX as usize {
+    pub async fn read(
+        &mut self,
+        lun: u8,
+        lba: u64,
+        lba_count: u16,
+        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
+    ) -> Result<usize, Error> {
+        if lba >= u32::MAX as u64 {
             return Err(Error::InvalidArgument);
         }
         let lba_bytes = (lba as u32).to_be_bytes();
+        let lba_count = (lba_count as u16).to_be_bytes();
         // Issue a READ (10) command
         let request_buffer = [
             0x28,
@@ -40,8 +54,8 @@ impl ScsiTransportWrapper {
             lba_bytes[2],
             lba_bytes[3],
             0x00,
-            0x00,
-            0x01,
+            lba_count[0],
+            lba_count[1],
             0x00,
         ];
 
@@ -59,14 +73,28 @@ impl ScsiTransportWrapper {
         [u8; R::RESPONSE_LEN]: Sized,
         [u8; R::REQUEST_LEN]: Sized,
     {
-        let mut response_buffer = [0; R::RESPONSE_LEN];
+        let mut response_buffer = self.allocate_buffer(R::RESPONSE_LEN)?;
         let request_buffer = request.into_bytes();
 
         let response_len = self
             .inner
-            .perform_request_raw(lun, &request_buffer, &mut response_buffer)
+            .perform_request_raw(
+                lun,
+                &request_buffer,
+                response_buffer.slice_mut(0..R::RESPONSE_LEN),
+            )
             .await?;
+        let response_bytes =
+            unsafe { MaybeUninit::slice_assume_init_ref(&response_buffer[..response_len]) };
 
-        R::parse_response(&response_buffer[..response_len])
+        R::parse_response(response_bytes)
+    }
+}
+
+impl Deref for ScsiTransportWrapper {
+    type Target = dyn ScsiTransport;
+
+    fn deref(&self) -> &Self::Target {
+        self.inner.as_ref()
     }
 }

kernel/driver/bus/pci/src/driver.rs

@@ -1,5 +1,5 @@
 use alloc::{sync::Arc, vec::Vec};
-use device_api::device::Device;
+use device_api::{device::Device, dma::DmaAllocator};
 use libk::error::Error;
 use libk_util::sync::spin_rwlock::IrqSafeRwLock;
 
@@ -17,7 +17,11 @@ pub struct PciDriverMatch {
 }
 
 pub trait PciDriver: Sync {
-    fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error>;
+    fn probe(
+        &self,
+        info: &PciDeviceInfo,
+        dma: &Arc<dyn DmaAllocator>,
+    ) -> Result<Arc<dyn Device>, Error>;
     fn driver_name(&self) -> &str;
 }
 

kernel/driver/bus/pci/src/lib.rs

@@ -13,8 +13,12 @@ use alloc::{format, sync::Arc, vec::Vec};
 
 use bitflags::bitflags;
 use device::{PciBusDevice, PciDeviceInfo};
+use device_api::{device::DeviceInitContext, dma::DmaAllocator};
 use interrupt::{PciInterruptMap, PciMsiMap};
-use libk::fs::sysfs::{self, object::KObject};
+use libk::{
+    dma::DummyDmaAllocator,
+    fs::sysfs::{self, object::KObject},
+};
 use libk_mm::address::PhysicalAddress;
 use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
 use space::legacy;
@@ -613,9 +617,14 @@ fn setup_bus_device(device: &mut PciBusDevice) -> Result<(), Error> {
 
     if let Some(driver) = driver::lookup_driver(&device.info) {
         log::info!("{} -> {}", device.info.address, driver.driver_name());
-        let instance = driver.probe(&device.info)?;
+        let dma: Arc<dyn DmaAllocator> = Arc::new(DummyDmaAllocator);
+        let cx = DeviceInitContext {
+            dma_allocator: dma.clone(),
+        };
 
-        unsafe { instance.clone().init() }?;
+        let instance = driver.probe(&device.info, &dma)?;
+
+        unsafe { instance.clone().init(cx) }?;
 
         device.device.replace(instance);
         device.driver_name.replace(driver.driver_name());

kernel/driver/bus/usb/src/class_driver/hid_keyboard.rs

@@ -2,7 +2,6 @@ use core::mem::MaybeUninit;
 
 use alloc::{boxed::Box, sync::Arc};
 use async_trait::async_trait;
-use libk_mm::PageBox;
 use yggdrasil_abi::io::{KeyboardKey, KeyboardKeyEvent};
 
 use crate::{device::UsbDeviceAccess, error::UsbError, info::UsbDeviceClass};
@@ -136,14 +135,14 @@ impl UsbDriver for UsbHidKeyboardDriver {
             .open_interrupt_in_pipe(1, config.endpoints[0].max_packet_size as u16)
             .await?;
 
-        let mut buffer = PageBox::new_slice(0, 8).map_err(UsbError::MemoryError)?;
+        let mut buffer = [0; 8];
         let mut state = KeyboardState::new();
         let mut events = [MaybeUninit::uninit(); 16];
 
         loop {
             let mut event_count = 0;
 
-            let len = pipe.read(buffer.as_slice_mut()).await?;
+            let len = pipe.read(&mut buffer).await?;
             if len < 8 {
                 continue;
             }

kernel/driver/bus/usb/src/class_driver/mass_storage.rs

@@ -1,9 +1,13 @@
+use core::mem::MaybeUninit;
+
 use alloc::{boxed::Box, sync::Arc};
 use async_trait::async_trait;
 use bytemuck::{Pod, Zeroable};
-use libk::error::Error;
-use libk_mm::PageBox;
-use ygg_driver_scsi::{transport::ScsiTransport, ScsiDevice};
+use libk::{
+    dma::{DmaBuffer, DmaSliceMut},
+    error::Error,
+};
+use ygg_driver_scsi::{transport::ScsiTransport, ScsiEnclosure};
 
 use crate::{
     communication::UsbDirection,
@@ -11,7 +15,7 @@ use crate::{
     error::UsbError,
     info::{UsbDeviceClass, UsbEndpointType},
     pipe::{
-        control::UsbClassSpecificRequest,
+        control::{ControlTransferSetup, UsbClassSpecificRequest},
         normal::{UsbBulkInPipeAccess, UsbBulkOutPipeAccess},
     },
 };
@@ -49,11 +53,10 @@ struct Bbb {
     device: Arc<UsbDeviceAccess>,
     in_pipe: UsbBulkInPipeAccess,
     out_pipe: UsbBulkOutPipeAccess,
-    buffer: PageBox<[u8]>,
     last_tag: u32,
 }
 
-struct DetachHandler(Arc<ScsiDevice>);
+struct DetachHandler(Arc<ScsiEnclosure>);
 
 impl Bbb {
     pub fn new(
@@ -61,12 +64,10 @@ impl Bbb {
         in_pipe: UsbBulkInPipeAccess,
         out_pipe: UsbBulkOutPipeAccess,
     ) -> Result<Self, UsbError> {
-        let buffer = PageBox::new_slice(0, 4096).map_err(UsbError::MemoryError)?;
         Ok(Self {
             device,
             in_pipe,
             out_pipe,
-            buffer,
             last_tag: 0,
         })
     }
@@ -82,36 +83,35 @@ impl Bbb {
     ) -> Result<u32, Error> {
         self.last_tag = self.last_tag.wrapping_add(1);
 
-        self.buffer[..size_of::<Cbw>()].fill(0);
-
-        let cbw = bytemuck::from_bytes_mut::<Cbw>(&mut self.buffer[..size_of::<Cbw>()]);
+        let flags = if !host_to_dev { 1 << 7 } else { 0 };
         let tag = self.last_tag;
+        let mut cbw_bytes = [0; 32];
+        let cbw = bytemuck::from_bytes_mut::<Cbw>(&mut cbw_bytes);
 
         cbw.signature = 0x43425355;
-        cbw.tag = tag;
         cbw.transfer_length = response_len as u32;
-        if !host_to_dev {
-            cbw.flags = 1 << 7;
-        }
+        cbw.flags = flags;
+        cbw.tag = tag;
         cbw.lun = lun;
         cbw.cb_length = command.len() as u8;
         cbw.cb_data[..command.len()].copy_from_slice(command);
 
         self.out_pipe
-            .write(self.buffer.as_slice().subslice(..31))
+            .write(&cbw_bytes[..31])
             .await
-            .inspect_err(|error| log::error!("msc: out pipe error: {error:?}"))?;
+            .inspect_err(|error| log::error!("msc: CBW send error: {error:?}"))?;
 
         Ok(tag)
     }
 
     async fn read_csw(&mut self, tag: u32) -> Result<(), Error> {
+        let mut csw_bytes = [0; 16];
         self.in_pipe
-            .read(self.buffer.as_slice_mut().subslice_mut(..13))
+            .read_exact(&mut csw_bytes[..13])
             .await
-            .inspect_err(|error| log::error!("msc: in pipe error: {error:?}"))?;
+            .inspect_err(|error| log::error!("msc: CSW receive error: {error:?}"))?;
+        let csw = bytemuck::from_bytes::<Csw>(&csw_bytes);
 
-        let csw = bytemuck::from_bytes::<Csw>(&self.buffer[..size_of::<Csw>()]);
         if csw.signature != 0x53425355 {
             log::warn!("msc: invalid csw signature");
             return Err(Error::InvalidArgument);
@@ -122,47 +122,54 @@ impl Bbb {
             return Err(Error::InvalidArgument);
         }
         if csw.status != 0x00 {
-            log::warn!("msc: csw error status {:#02x}", csw.status);
             return Err(Error::InvalidArgument);
         }
         Ok(())
     }
 
-    async fn read_response_data(&mut self, buffer: &mut [u8]) -> Result<usize, Error> {
-        // TODO limit by max_packet_size
+    async fn read_response_data(
+        &mut self,
+        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
+    ) -> Result<usize, Error> {
+        if buffer.len() == 0 {
+            return Ok(0);
+        }
         let len = self
             .in_pipe
-            .read(self.buffer.as_slice_mut().subslice_mut(..buffer.len()))
+            .read_dma(buffer)
             .await
-            .inspect_err(|error| log::error!("msc: in pipe error: {error:?}"))?;
-        buffer[..len].copy_from_slice(&self.buffer[..len]);
+            .inspect_err(|error| log::error!("msc: DMA read error: {error:?}"))?;
         Ok(len)
     }
 }
 
 #[async_trait]
 impl ScsiTransport for Bbb {
+    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        Ok(self.in_pipe.allocate_dma_buffer(size)?)
+    }
+
     async fn perform_request_raw(
         &mut self,
         lun: u8,
         request_data: &[u8],
-        response_buffer: &mut [u8],
+        response_buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
     ) -> Result<usize, Error> {
-        if request_data.len() > 16 || response_buffer.len() > self.buffer.len() {
+        if request_data.len() > 16 || response_buffer.len() > self.max_bytes_per_request() {
             return Err(Error::InvalidArgument);
         }
 
         let tag = self
             .send_cbw(lun, false, request_data, response_buffer.len())
             .await?;
-        let response_len = if response_buffer.is_empty() {
-            0
-        } else {
-            self.read_response_data(response_buffer).await?
-        };
+        let response_len = self.read_response_data(response_buffer).await?;
         self.read_csw(tag).await?;
         Ok(response_len)
     }
+
+    fn max_bytes_per_request(&self) -> usize {
+        32768
+    }
 }
 
 impl UsbDeviceDetachHandler for DetachHandler {
@@ -176,11 +183,20 @@ impl UsbDeviceDetachHandler for DetachHandler {
 #[repr(C)]
 pub struct BulkOnlyMassStorageReset;
 
+#[derive(Debug, Pod, Zeroable, Clone, Copy)]
+#[repr(C)]
+pub struct GetMaxLun;
+
 impl UsbClassSpecificRequest for BulkOnlyMassStorageReset {
     const BM_REQUEST_TYPE: u8 = 0b00100001;
     const B_REQUEST: u8 = 0b11111111;
 }
 
+impl UsbClassSpecificRequest for GetMaxLun {
+    const BM_REQUEST_TYPE: u8 = 0b10100001;
+    const B_REQUEST: u8 = 0b11111110;
+}
+
 #[async_trait]
 impl UsbDriver for UsbMassStorageDriverBulkOnly {
     async fn run(self: Arc<Self>, device: Arc<UsbDeviceAccess>) -> Result<(), UsbError> {
@@ -188,7 +204,6 @@ impl UsbDriver for UsbMassStorageDriverBulkOnly {
         let config = device.select_configuration(|_| true).await?.unwrap();
         // Bulk-in, bulk-out
         assert_eq!(config.endpoints.len(), 2);
-        // TODO those indices may be different
         let control_pipe = device.control_pipe();
         let (in_index, in_info) = config
             .find_endpoint(|ep| ep.is(UsbEndpointType::Bulk, UsbDirection::In))
@@ -206,21 +221,44 @@ impl UsbDriver for UsbMassStorageDriverBulkOnly {
         // Perform a Bulk-Only Mass Storage Reset
         // TODO interface id?
         control_pipe
-            .class_specific_request::<BulkOnlyMassStorageReset>(0, 0)
+            .control_transfer(ControlTransferSetup {
+                bm_request_type: BulkOnlyMassStorageReset::BM_REQUEST_TYPE,
+                b_request: BulkOnlyMassStorageReset::B_REQUEST,
+                w_value: 0,
+                w_index: 0,
+                w_length: 0,
+            })
             .await?;
 
-        // TODO get max LUN
+        // Get max LUN
+        // TODO on devices which do not support multiple LUNs, this command may STALL
+        let mut buffer = [MaybeUninit::uninit()];
+        let len = control_pipe
+            .control_transfer_in(
+                ControlTransferSetup {
+                    bm_request_type: GetMaxLun::BM_REQUEST_TYPE,
+                    b_request: GetMaxLun::B_REQUEST,
+                    w_value: 0,
+                    w_index: 0,
+                    w_length: 1,
+                },
+                &mut buffer,
+            )
+            .await?;
+        let max_lun = if len < 1 {
+            0
+        } else {
+            unsafe { buffer[0].assume_init() }
+        };
 
         let bbb = Bbb::new(device.clone(), in_pipe, out_pipe)?;
-        let scsi = ScsiDevice::setup(bbb)
+        let scsi = ScsiEnclosure::setup(Box::new(bbb), max_lun as usize + 1)
             .await
             .inspect_err(|error| log::error!("msc: scsi error {error:?}"))
             .map_err(|_| UsbError::DriverError)?;
         let detach = DetachHandler(scsi.clone());
         device.set_detach_handler(Arc::new(detach));
 
-        ygg_driver_scsi::attach(scsi).ok();
-
         Ok(())
     }
 

kernel/driver/bus/usb/src/device.rs

@@ -2,7 +2,6 @@ use core::{fmt, ops::Deref};
 
 use alloc::{boxed::Box, sync::Arc, vec::Vec};
 use async_trait::async_trait;
-use libk_mm::PageBox;
 use libk_util::sync::spin_rwlock::{IrqSafeRwLock, IrqSafeRwLockReadGuard};
 
 use crate::{
@@ -87,7 +86,6 @@ impl UsbDeviceAccess {
     /// * Device has been assigned a bus address
     pub async fn setup(raw: Arc<dyn UsbDevice>) -> Result<Self, UsbError> {
         let control = raw.control_pipe();
-        let mut string_buffer = PageBox::new_uninit().map_err(UsbError::MemoryError)?;
 
         let device_desc = control.query_device_descriptor().await?;
 
@@ -102,12 +100,8 @@ impl UsbDeviceAccess {
                 )
             })?;
 
-        let manufacturer = control
-            .query_string(device_desc.manufacturer_str, &mut string_buffer)
-            .await?;
-        let product = control
-            .query_string(device_desc.product_str, &mut string_buffer)
-            .await?;
+        let manufacturer = control.query_string(device_desc.manufacturer_str).await?;
+        let product = control.query_string(device_desc.product_str).await?;
 
         let info = UsbDeviceInfo {
             manufacturer,
@@ -208,12 +202,11 @@ impl UsbDeviceAccess {
         if index >= self.info.num_configurations {
             return Err(UsbError::InvalidConfiguration);
         }
-        let mut string_buffer = PageBox::new_uninit().map_err(UsbError::MemoryError)?;
         let control_pipe = self.control_pipe();
         let query = control_pipe.query_configuration_descriptor(index).await?;
 
         let configuration_name = control_pipe
-            .query_string(query.configuration().config_str, &mut string_buffer)
+            .query_string(query.configuration().config_str)
             .await?;
 
         let mut endpoints = Vec::new();
@@ -230,9 +223,7 @@ impl UsbDeviceAccess {
                     });
                 }
                 ConfigurationDescriptorEntry::Interface(iface) => {
-                    let name = control_pipe
-                        .query_string(iface.interface_str, &mut string_buffer)
-                        .await?;
+                    let name = control_pipe.query_string(iface.interface_str).await?;
                     interfaces.push(UsbInterfaceInfo {
                         name,
                         number: iface.interface_number,

kernel/driver/bus/usb/src/error.rs

@@ -43,7 +43,10 @@ impl From<UsbError> for Error {
     fn from(value: UsbError) -> Self {
         match value {
             UsbError::MemoryError(e) => e,
-            _ => Error::InvalidOperation,
+            e => {
+                log::warn!("usb: {e:?}");
+                Error::InvalidOperation
+            }
         }
     }
 }

kernel/driver/bus/usb/src/lib.rs

@@ -1,6 +1,13 @@
 #![no_std]
 #![allow(clippy::new_without_default)]
-#![feature(iter_array_chunks, maybe_uninit_slice)]
+#![feature(
+    iter_array_chunks,
+    maybe_uninit_slice,
+    maybe_uninit_as_bytes,
+    maybe_uninit_uninit_array,
+    maybe_uninit_write_slice,
+    maybe_uninit_fill
+)]
 
 extern crate alloc;
 

kernel/driver/bus/usb/src/pipe/control.rs

@@ -1,5 +1,4 @@
 use core::{
-    cmp::Ordering,
     mem::{size_of, MaybeUninit},
     ops::Deref,
 };
@@ -7,10 +6,9 @@ use core::{
 use alloc::{boxed::Box, string::String};
 use async_trait::async_trait;
 use bytemuck::{Pod, Zeroable};
-use libk_mm::{PageBox, PageSlice};
+use libk_mm::PageBox;
 
 use crate::{
-    communication::UsbDirection,
     descriptor::{
         UsbConfigurationDescriptor, UsbDeviceDescriptor, UsbDeviceQualifier, UsbEndpointDescriptor,
         UsbInterfaceDescriptor, UsbOtherSpeedConfiguration,
@@ -86,10 +84,16 @@ fn decode_usb_string(bytes: &[u8]) -> Result<String, UsbError> {
 
 #[async_trait]
 pub trait UsbControlPipe: Send + Sync {
-    async fn control_transfer(
+    async fn control_transfer(&self, setup: ControlTransferSetup) -> Result<(), UsbError>;
+    async fn control_transfer_in(
         &self,
         setup: ControlTransferSetup,
-        data: Option<(&mut PageSlice<MaybeUninit<u8>>, UsbDirection)>,
+        buffer: &mut [MaybeUninit<u8>],
+    ) -> Result<usize, UsbError>;
+    async fn control_transfer_out(
+        &self,
+        setup: ControlTransferSetup,
+        buffer: &[u8],
     ) -> Result<usize, UsbError>;
 }
 
@@ -173,28 +177,10 @@ impl ConfigurationDescriptorQuery {
 }
 
 impl UsbControlPipeAccess {
-    pub async fn query_device_descriptor2(&self) -> Result<PageBox<UsbDeviceDescriptor>, UsbError> {
-        let mut output = PageBox::new_uninit().map_err(UsbError::MemoryError)?;
+    pub async fn query_device_descriptor(&self) -> Result<UsbDeviceDescriptor, UsbError> {
+        let mut buffer = MaybeUninit::uninit();
 
-        self.control_transfer(
-            ControlTransferSetup {
-                bm_request_type: 0b10000000,
-                b_request: 0xFF,
-                // b_request: 0x06,
-                w_value: 0x100,
-                w_index: 0,
-                w_length: size_of::<UsbDeviceDescriptor>() as _,
-            },
-            Some((PageBox::as_bytes_mut(&mut output), UsbDirection::In)),
-        )
-        .await?;
-
-        Ok(unsafe { output.assume_init() })
-    }
-    pub async fn query_device_descriptor(&self) -> Result<PageBox<UsbDeviceDescriptor>, UsbError> {
-        let mut output = PageBox::new_uninit().map_err(UsbError::MemoryError)?;
-
-        self.control_transfer(
+        self.control_transfer_in(
             ControlTransferSetup {
                 bm_request_type: 0b10000000,
                 b_request: 0x06,
@@ -202,19 +188,19 @@ impl UsbControlPipeAccess {
                 w_index: 0,
                 w_length: size_of::<UsbDeviceDescriptor>() as _,
             },
-            Some((PageBox::as_bytes_mut(&mut output), UsbDirection::In)),
+            buffer.as_bytes_mut(),
         )
         .await?;
 
-        Ok(unsafe { output.assume_init() })
+        Ok(unsafe { buffer.assume_init() })
     }
 
     async fn fill_configuation_descriptor(
         &self,
         index: u8,
-        buffer: &mut PageBox<[MaybeUninit<u8>]>,
+        buffer: &mut [MaybeUninit<u8>],
     ) -> Result<(), UsbError> {
-        self.control_transfer(
+        self.control_transfer_in(
             ControlTransferSetup {
                 bm_request_type: 0b10000000,
                 b_request: 0x06,
@@ -222,18 +208,17 @@ impl UsbControlPipeAccess {
                 w_index: 0,
                 w_length: buffer.len().try_into().unwrap(),
             },
-            Some((buffer.as_slice_mut(), UsbDirection::In)),
+            buffer,
         )
         .await?;
         Ok(())
     }
 
-    pub async fn query_string(
-        &self,
-        index: u8,
-        buffer: &mut PageBox<MaybeUninit<[u8; 4096]>>,
-    ) -> Result<String, UsbError> {
-        self.control_transfer(
+    pub async fn query_string(&self, index: u8) -> Result<String, UsbError> {
+        let mut buffer = MaybeUninit::uninit_array::<256>();
+
+        let len = self
+            .control_transfer_in(
                 ControlTransferSetup {
                     bm_request_type: 0b10000000,
                     b_request: 0x06,
@@ -241,11 +226,10 @@ impl UsbControlPipeAccess {
                     w_index: 0,
                     w_length: 4096,
                 },
-            Some((PageBox::as_bytes_mut(buffer), UsbDirection::In)),
+                &mut buffer[..],
             )
             .await?;
-        let data = unsafe { buffer.assume_init_ref() };
-
+        let data = unsafe { MaybeUninit::slice_assume_init_ref(&buffer[..len]) };
         let len = data[0] as usize;
 
         decode_usb_string(&data[2..len])
@@ -255,35 +239,20 @@ impl UsbControlPipeAccess {
         &self,
         index: u8,
     ) -> Result<ConfigurationDescriptorQuery, UsbError> {
-        // First, query the real length of the descriptor
-        let mut buffer = PageBox::new_uninit_slice(size_of::<UsbConfigurationDescriptor>())
-            .map_err(UsbError::MemoryError)?;
-        self.fill_configuation_descriptor(index, &mut buffer)
+        // 4K should be enough for a configuration descriptor
+        let mut buffer = PageBox::new_uninit_slice(4096).map_err(UsbError::MemoryError)?;
+
+        self.fill_configuation_descriptor(index, &mut buffer[..])
             .await?;
-        let buffer = unsafe { PageBox::assume_init_slice(buffer) };
 
-        let desc: &UsbConfigurationDescriptor = bytemuck::from_bytes(&buffer);
-        let total_len = desc.total_length as usize;
-
-        // Return if everything's ready at this point
-        match total_len.cmp(&size_of::<UsbConfigurationDescriptor>()) {
-            Ordering::Less => todo!(),
-            Ordering::Equal => return Ok(ConfigurationDescriptorQuery { buffer }),
-            _ => (),
-        }
-
-        // Otherwise, query the rest of the data
-        let mut buffer = PageBox::new_uninit_slice(total_len).map_err(UsbError::MemoryError)?;
-        self.fill_configuation_descriptor(index, &mut buffer)
-            .await?;
         let buffer = unsafe { PageBox::assume_init_slice(buffer) };
 
         let desc: &UsbConfigurationDescriptor =
             bytemuck::from_bytes(&buffer[..size_of::<UsbConfigurationDescriptor>()]);
         let total_len = desc.total_length as usize;
 
-        if total_len != buffer.len() {
-            todo!();
+        if total_len > 4096 {
+            unimplemented!("4KiB wasn't enough");
         }
 
         Ok(ConfigurationDescriptorQuery { buffer })
@@ -294,38 +263,30 @@ impl UsbControlPipeAccess {
         w_value: u16,
         w_index: u16,
     ) -> Result<(), UsbError> {
-        self.control_transfer(
-            ControlTransferSetup {
+        self.control_transfer(ControlTransferSetup {
             bm_request_type: D::BM_REQUEST_TYPE,
             b_request: D::B_REQUEST,
             w_value,
             w_index,
             w_length: 0,
-            },
-            None,
-        )
-        .await?;
-        Ok(())
+        })
+        .await
     }
 
-    pub async fn class_specific_request<D: UsbClassSpecificRequest>(
-        &self,
-        w_value: u16,
-        w_index: u16,
-    ) -> Result<(), UsbError> {
-        self.control_transfer(
-            ControlTransferSetup {
-                bm_request_type: D::BM_REQUEST_TYPE,
-                b_request: D::B_REQUEST,
-                w_value,
-                w_index,
-                w_length: 0,
-            },
-            None,
-        )
-        .await?;
-        Ok(())
-    }
+    // pub async fn class_specific_request<D: UsbClassSpecificRequest>(
+    //     &self,
+    //     w_value: u16,
+    //     w_index: u16,
+    // ) -> Result<usize, UsbError> {
+    //     self.control_transfer(ControlTransferSetup {
+    //         bm_request_type: D::BM_REQUEST_TYPE,
+    //         b_request: D::B_REQUEST,
+    //         w_value,
+    //         w_index,
+    //         w_length: 0,
+    //     })
+    //     .await
+    // }
 
     pub async fn set_configuration(&self, value: u16) -> Result<(), UsbError> {
         self.perform_action::<SetConfiguration>(value, 0).await

kernel/driver/bus/usb/src/pipe/mod.rs

@@ -1,2 +1,12 @@
+use core::mem::MaybeUninit;
+
+use libk::dma::DmaBuffer;
+
+use crate::error::UsbError;
+
 pub mod control;
 pub mod normal;
+
+pub trait UsbGenericPipe: Send + Sync {
+    fn allocate_dma_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, UsbError>;
+}

kernel/driver/bus/usb/src/pipe/normal.rs

@@ -1,15 +1,26 @@
-use core::ops::Deref;
+use core::{mem::MaybeUninit, ops::Deref};
 
 use alloc::boxed::Box;
 use async_trait::async_trait;
-use libk_mm::PageSlice;
+use libk::dma::{DmaBuffer, DmaSlice, DmaSliceMut};
 
 use crate::error::{TransferError, UsbError};
 
+use super::UsbGenericPipe;
+
 #[async_trait]
-pub trait UsbNormalPipeIn: Send + Sync {
-    async fn read(&self, buffer: &mut PageSlice<u8>) -> Result<usize, UsbError>;
-    async fn read_exact(&self, buffer: &mut PageSlice<u8>) -> Result<(), UsbError> {
+pub trait UsbNormalPipeIn: UsbGenericPipe {
+    async fn read_dma(&self, buffer: DmaSliceMut<'_, MaybeUninit<u8>>) -> Result<usize, UsbError>;
+
+    async fn read(&self, buffer: &mut [u8]) -> Result<usize, UsbError> {
+        let mut dma_buffer = self.allocate_dma_buffer(buffer.len())?;
+        let len = self.read_dma(dma_buffer.slice_mut(0..buffer.len())).await?;
+        let dma_slice = unsafe { MaybeUninit::slice_assume_init_ref(&dma_buffer[..len]) };
+        buffer[..len].copy_from_slice(dma_slice);
+        Ok(len)
+    }
+
+    async fn read_exact(&self, buffer: &mut [u8]) -> Result<(), UsbError> {
         match self.read(buffer).await {
             Ok(len) if len == buffer.len() => Ok(()),
             Ok(len) => Err(UsbError::TransferFailed(TransferError::ShortPacket(len))),
@@ -19,8 +30,16 @@ pub trait UsbNormalPipeIn: Send + Sync {
 }
 
 #[async_trait]
-pub trait UsbNormalPipeOut: Send + Sync {
-    async fn write(&self, buffer: &PageSlice<u8>) -> Result<usize, UsbError>;
+pub trait UsbNormalPipeOut: UsbGenericPipe {
+    async fn write_dma(&self, buffer: DmaSlice<'_, u8>) -> Result<usize, UsbError>;
+
+    async fn write(&self, buffer: &[u8]) -> Result<usize, UsbError> {
+        let mut dma_buffer = self.allocate_dma_buffer(buffer.len())?;
+        MaybeUninit::copy_from_slice(&mut dma_buffer, buffer);
+        let dma_buffer = unsafe { DmaBuffer::assume_init_slice(dma_buffer) };
+
+        self.write_dma(dma_buffer.slice(0..buffer.len())).await
+    }
 }
 
 pub struct UsbBulkInPipeAccess(pub Box<dyn UsbNormalPipeIn>);

kernel/driver/net/core/src/ethernet.rs

@@ -2,7 +2,7 @@ use core::{fmt, mem::size_of};
 
 use alloc::sync::Arc;
 use bytemuck::Pod;
-use libk_mm::PageBox;
+use libk::dma::DmaBuffer;
 use yggdrasil_abi::{
     error::Error,
     net::{
@@ -24,7 +24,7 @@ pub struct L2Packet {
     pub l2_offset: usize,
     pub l3_offset: usize,
 
-    pub data: Arc<PageBox<[u8]>>,
+    pub data: Arc<DmaBuffer<[u8]>>,
 }
 
 /// Defines an Ethernet link speed

kernel/driver/net/core/src/interface.rs

@@ -1,12 +1,12 @@
 use core::{
-    mem::size_of,
+    mem::{size_of, MaybeUninit},
     net::IpAddr,
     sync::atomic::{AtomicU32, AtomicUsize, Ordering},
 };
 
 use alloc::{boxed::Box, collections::BTreeMap, format, sync::Arc};
+use libk::dma::DmaBuffer;
 // TODO: link state management?
-use libk_mm::PageBox;
 use libk_util::{
     sync::spin_rwlock::{IrqSafeRwLock, IrqSafeRwLockReadGuard},
     OneTimeInit,
@@ -19,7 +19,9 @@ use yggdrasil_abi::{
 use crate::l3::{arp::ArpTable, Route};
 
 pub trait NetworkDevice: Sync + Send {
-    fn transmit(&self, packet: PageBox<[u8]>) -> Result<(), Error>;
+    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error>;
+    fn transmit_buffer(&self, buffer: DmaBuffer<[u8]>) -> Result<(), Error>;
+
     fn packet_prefix_size(&self) -> usize;
 
     fn read_hardware_address(&self) -> MacAddress;
@@ -103,12 +105,15 @@ impl NetworkInterface {
         let l2_offset = self.device.packet_prefix_size();
         let l2_data_offset = l2_offset + size_of::<EthernetFrame>();
 
-        let mut packet = PageBox::new_slice(0, l2_data_offset + l2_data.len())?;
+        let packet = self
+            .device
+            .allocate_transmit_buffer(l2_data_offset + l2_data.len())?;
+        let mut packet = unsafe { DmaBuffer::assume_init_slice(packet) };
 
         packet[l2_offset..l2_data_offset].copy_from_slice(bytemuck::bytes_of(l2_frame));
-        packet[l2_data_offset..].copy_from_slice(l2_data);
+        packet[l2_data_offset..l2_data_offset + l2_data.len()].copy_from_slice(l2_data);
 
-        self.device.transmit(packet)
+        self.device.transmit_buffer(packet)
     }
 }
 

kernel/driver/net/core/src/l3/mod.rs

@@ -6,7 +6,7 @@ use core::{
 
 use alloc::{sync::Arc, vec::Vec};
 use bytemuck::{Pod, Zeroable};
-use libk_mm::PageBox;
+use libk::dma::DmaBuffer;
 use libk_util::sync::spin_rwlock::{
     IrqSafeRwLock, IrqSafeRwLockReadGuard, IrqSafeRwLockWriteGuard,
 };
@@ -40,7 +40,7 @@ pub struct L3Packet {
     pub l4_offset: usize,
     pub data_length: usize,
 
-    pub data: Arc<PageBox<[u8]>>,
+    pub data: Arc<DmaBuffer<[u8]>>,
 }
 
 pub trait IpFrame: Pod {
@@ -221,7 +221,7 @@ pub async fn send_l4_ip_resolved(packet: &L4ResolvedPacket<'_, '_>) -> Result<()
     let l3_frame = packet.make_l3_frame()?;
 
     let mut builder = PacketBuilder::new(
-        packet.interface.device.packet_prefix_size(),
+        packet.interface.device.as_ref(),
         size_of::<EthernetFrame>() + size_of::<Ipv4Frame>() + packet.total_l4_len(),
     )?;
     builder.push(&EthernetFrame {
@@ -235,7 +235,7 @@ pub async fn send_l4_ip_resolved(packet: &L4ResolvedPacket<'_, '_>) -> Result<()
     builder.push_bytes(packet.l4_data)?;
 
     let (sent_packet, _len) = builder.finish();
-    packet.interface.device.transmit(sent_packet)
+    packet.interface.device.transmit_buffer(sent_packet)
 }
 
 pub async fn send_l4_ip(packet: &L4UnresolvedPacket<'_>) -> Result<(), Error> {

kernel/driver/net/core/src/lib.rs

@@ -9,9 +9,9 @@ use core::mem::size_of;
 use alloc::sync::Arc;
 use bytemuck::Pod;
 use ethernet::L2Packet;
+use interface::NetworkDevice;
 use l3::L3Packet;
-use libk::task::runtime;
-use libk_mm::PageBox;
+use libk::{dma::DmaBuffer, task::runtime};
 use libk_util::queue::UnboundedMpmcQueue;
 use yggdrasil_abi::{error::Error, net::protocols::EthernetFrame};
 
@@ -28,20 +28,22 @@ pub use interface::register_interface;
 
 pub struct Packet {
     // TODO info about "received" interface
-    buffer: PageBox<[u8]>,
+    buffer: DmaBuffer<[u8]>,
     offset: usize,
     iface: u32,
 }
 
 pub struct PacketBuilder {
-    data: PageBox<[u8]>,
+    data: DmaBuffer<[u8]>,
     pos: usize,
     len: usize,
 }
 
 impl PacketBuilder {
-    pub fn new(l2_offset: usize, l2_size: usize) -> Result<Self, Error> {
-        let data = PageBox::new_slice(0, l2_offset + l2_size)?;
+    pub fn new(nic: &dyn NetworkDevice, l2_size: usize) -> Result<Self, Error> {
+        let l2_offset = nic.packet_prefix_size();
+        let data = nic.allocate_transmit_buffer(l2_offset + l2_size)?;
+        let data = unsafe { DmaBuffer::assume_init_slice(data) };
         Ok(Self {
             data,
             pos: l2_offset,
@@ -63,14 +65,14 @@ impl PacketBuilder {
         Ok(())
     }
 
-    pub fn finish(self) -> (PageBox<[u8]>, usize) {
+    pub fn finish(self) -> (DmaBuffer<[u8]>, usize) {
         (self.data, self.len)
     }
 }
 
 impl Packet {
     #[inline]
-    pub fn new(buffer: PageBox<[u8]>, offset: usize, iface: u32) -> Self {
+    pub fn new(buffer: DmaBuffer<[u8]>, offset: usize, iface: u32) -> Self {
         Self {
             buffer,
             offset,

kernel/driver/net/core/src/socket/raw.rs

@@ -8,11 +8,11 @@ use core::{
 use alloc::{boxed::Box, collections::btree_map::BTreeMap, sync::Arc, vec::Vec};
 use async_trait::async_trait;
 use libk::{
+    dma::DmaBuffer,
     error::Error,
     task::runtime::maybe_timeout,
     vfs::{FileReadiness, Socket},
 };
-use libk_mm::PageBox;
 use libk_util::{queue::BoundedMpmcQueue, sync::spin_rwlock::IrqSafeRwLock};
 use yggdrasil_abi::{
     net::{
@@ -60,7 +60,6 @@ impl RawSocket {
         if let Some(ids) = bound_sockets.get(&packet.interface_id) {
             for id in ids {
                 let socket = raw_sockets.get(id).unwrap();
-                log::info!("Packet -> {id}");
                 socket.bound_packet_received(packet.clone());
             }
         }
@@ -172,9 +171,12 @@ impl Socket for RawSocket {
         if message.payload.len() > 4096 - l2_offset {
             return Err(Error::InvalidArgument);
         }
-        let mut packet = PageBox::new_slice(0, l2_offset + message.payload.len())?;
+        let packet = interface
+            .device
+            .allocate_transmit_buffer(l2_offset + message.payload.len())?;
+        let mut packet = unsafe { DmaBuffer::assume_init_slice(packet) };
         packet[l2_offset..l2_offset + message.payload.len()].copy_from_slice(message.payload);
-        interface.device.transmit(packet)?;
+        interface.device.transmit_buffer(packet)?;
         Ok(message.payload.len())
     }
     async fn send_message(

kernel/driver/net/loopback/Cargo.toml

@@ -5,8 +5,10 @@ edition = "2021"
 
 [dependencies]
 yggdrasil-abi.workspace = true
+device-api.workspace = true
 libk-util.workspace = true
 libk-mm.workspace = true
+libk.workspace = true
 
 ygg_driver_net_core = { path = "../../net/core" }
 

kernel/driver/net/loopback/src/lib.rs

@@ -2,10 +2,14 @@
 
 extern crate alloc;
 
-use core::net::{IpAddr, Ipv4Addr};
+use core::{
+    mem::MaybeUninit,
+    net::{IpAddr, Ipv4Addr},
+};
 
 use alloc::sync::Arc;
-use libk_mm::PageBox;
+use device_api::dma::DmaAllocator;
+use libk::dma::{DmaBuffer, DummyDmaAllocator};
 use libk_util::OneTimeInit;
 use ygg_driver_net_core::{
     interface::{NetworkDevice, NetworkInterfaceType},
@@ -13,11 +17,18 @@ use ygg_driver_net_core::{
 };
 use yggdrasil_abi::{error::Error, net::MacAddress};
 
-struct LoopbackDevice;
+struct LoopbackDevice {
+    allocator: Arc<dyn DmaAllocator>,
+}
 
 impl NetworkDevice for LoopbackDevice {
-    fn transmit(&self, packet: PageBox<[u8]>) -> Result<(), Error> {
-        ygg_driver_net_core::receive_packet(Packet::new(packet, 0, *LOOPBACK_ID.get()))
+    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        DmaBuffer::new_uninit_slice(&*self.allocator, len)
+    }
+
+    fn transmit_buffer(&self, buffer: DmaBuffer<[u8]>) -> Result<(), Error> {
+        let packet = Packet::new(buffer, 0, *LOOPBACK_ID.get());
+        ygg_driver_net_core::receive_packet(packet)
     }
 
     fn packet_prefix_size(&self) -> usize {
@@ -32,7 +43,9 @@ impl NetworkDevice for LoopbackDevice {
 static LOOPBACK_ID: OneTimeInit<u32> = OneTimeInit::new();
 
 pub fn init() {
-    let loopback = Arc::new(LoopbackDevice);
+    let loopback = Arc::new(LoopbackDevice {
+        allocator: Arc::new(DummyDmaAllocator),
+    });
     let interface =
         ygg_driver_net_core::register_interface(NetworkInterfaceType::Loopback, loopback);
 

kernel/driver/net/rtl81xx/src/lib.rs

@@ -1,7 +1,7 @@
 #![no_std]
 
 use alloc::sync::Arc;
-use device_api::device::Device;
+use device_api::{device::Device, dma::DmaAllocator};
 use libk::error::Error;
 use rtl8139::Rtl8139;
 use rtl8168::Rtl8168;
@@ -23,7 +23,7 @@ pci_driver! {
             "rtl8139"
         }
 
-        fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
+        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
             info.init_interrupts(PreferredInterruptMode::Msi(false))?;
 
             // Enable MMIO + interrupts + bus mastering
@@ -38,7 +38,7 @@ pci_driver! {
                 .and_then(PciBaseAddress::as_memory)
                 .ok_or(Error::InvalidArgument)?;
 
-            let device = Rtl8139::new(base, info.clone())?;
+            let device = Rtl8139::new(dma.clone(), base, info.clone())?;
             Ok(Arc::new(device))
         }
     }
@@ -51,7 +51,7 @@ pci_driver! {
             "rtl816x"
         }
 
-        fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
+        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
             let base = info
                 .config_space
                 .bar(2)
@@ -65,7 +65,7 @@ pci_driver! {
             // Enable MMIO + interrupts + bus mastering
             info.set_command(true, true, false, true);
 
-            let device = Rtl8168::new(base, info.clone())?;
+            let device = Rtl8168::new(dma.clone(), base, info.clone())?;
             Ok(Arc::new(device))
         }
     }

kernel/driver/net/rtl81xx/src/rtl8139.rs

@@ -2,15 +2,12 @@ use core::mem::MaybeUninit;
 
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
+    dma::DmaAllocator,
     interrupt::{InterruptHandler, IrqVector},
 };
-use libk::error::Error;
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    device::DeviceMemoryIo,
-    PageBox,
-};
+use libk::{dma::DmaBuffer, error::Error};
+use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
 use libk_util::{queue::BoundedQueue, sync::IrqSafeSpinlock, OneTimeInit};
 use tock_registers::{
     interfaces::{ReadWriteable, Readable, Writeable},
@@ -207,17 +204,17 @@ register_structs! {
 }
 
 struct Rx {
-    buffer: PageBox<[MaybeUninit<u8>]>,
+    buffer: DmaBuffer<[MaybeUninit<u8>]>,
     rd: usize,
 }
 
 // TODO place a secondary Tx queue here, to send the queued packets when more slots become
 // available
 struct Tx {
-    buffers: [Option<PageBox<[u8]>>; 4],
+    buffers: [Option<DmaBuffer<[u8]>>; 4],
     wr: usize,
     rd: usize,
-    queue: BoundedQueue<PageBox<[u8]>>,
+    queue: BoundedQueue<DmaBuffer<[u8]>>,
 }
 
 pub struct Rtl8139 {
@@ -228,13 +225,13 @@ pub struct Rtl8139 {
     nic: OneTimeInit<u32>,
     rx: OneTimeInit<IrqSafeSpinlock<Rx>>,
     tx: IrqSafeSpinlock<Tx>,
+
+    dma: Arc<dyn DmaAllocator>,
 }
 
 impl Tx {
-    pub fn tx_now(&mut self, regs: &Regs, packet: PageBox<[u8]>) -> Result<(), Error> {
-        let packet_address = unsafe { packet.as_physical_address() }
-            .try_into_u32()
-            .map_err(|_| Error::InvalidArgument)?;
+    pub fn tx_now(&mut self, regs: &Regs, packet: DmaBuffer<[u8]>) -> Result<(), Error> {
+        let packet_address = packet.bus_address().try_into_u32()?;
         let packet_len = packet.len();
         if packet_len > 1500 {
             return Err(Error::InvalidArgument);
@@ -259,7 +256,11 @@ impl Rtl8139 {
     const RX_BUFFER_LEN: usize = 8192;
     const RX_BUFFER_OVERFLOW: usize = 4096;
 
-    pub fn new(base: PhysicalAddress, info: PciDeviceInfo) -> Result<Self, Error> {
+    pub fn new(
+        dma: Arc<dyn DmaAllocator>,
+        base: PhysicalAddress,
+        info: PciDeviceInfo,
+    ) -> Result<Self, Error> {
         let regs = unsafe { DeviceMemoryIo::<Regs>::map(base, Default::default()) }?;
         let mac0 = regs.IDRn[0].get().to_le_bytes();
         let mac1 = regs.IDRn[1].get().to_le_bytes();
@@ -278,6 +279,8 @@ impl Rtl8139 {
                 rd: 0,
                 queue: BoundedQueue::new(64),
             }),
+
+            dma,
         })
     }
 }
@@ -318,9 +321,10 @@ impl InterruptHandler for Rtl8139 {
                 let rx_len = u16::from_le_bytes([rx_len_0, rx_len_1]) as usize;
 
                 if rx_len >= 16 {
-                    if let Ok(mut packet_buf) = PageBox::new_uninit_slice(rx_len) {
+                    if let Ok(mut packet_buf) = DmaBuffer::new_uninit_slice(&*self.dma, rx_len) {
                         packet_buf.copy_from_slice(&rx.buffer[rx_pos + 4..rx_pos + rx_len + 4]);
-                        let packet_buf = unsafe { packet_buf.assume_init_slice() };
+                        let packet_buf = unsafe { DmaBuffer::assume_init_slice(packet_buf) };
+                        // let packet_buf = unsafe { packet_buf.assume_init_slice() };
                         let packet = Packet::new(packet_buf, 0, nic);
                         ygg_driver_net_core::receive_packet(packet).ok();
                     }
@@ -342,15 +346,16 @@ impl InterruptHandler for Rtl8139 {
 }
 
 impl Device for Rtl8139 {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         log::info!("Initialize rtl8139 driver");
         log::info!("MAC: {}", self.mac);
 
         // Setup the initial Rx buffer
-        let rx_buffer = PageBox::new_uninit_slice(Self::RX_BUFFER_LEN + Self::RX_BUFFER_OVERFLOW)?;
-        let rx_buffer_address = unsafe { rx_buffer.as_physical_address() }
-            .try_into_u32()
-            .map_err(|_| Error::InvalidArgument)?;
+        let rx_buffer = DmaBuffer::new_uninit_slice(
+            &*self.dma,
+            Self::RX_BUFFER_LEN + Self::RX_BUFFER_OVERFLOW,
+        )?;
+        let rx_buffer_address = rx_buffer.bus_address().try_into_u32()?;
 
         self.pci.map_interrupt(Default::default(), self.clone())?;
 
@@ -399,7 +404,11 @@ impl Device for Rtl8139 {
 }
 
 impl NetworkDevice for Rtl8139 {
-    fn transmit(&self, packet: PageBox<[u8]>) -> Result<(), Error> {
+    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        DmaBuffer::new_uninit_slice(&*self.dma, len)
+    }
+
+    fn transmit_buffer(&self, packet: DmaBuffer<[u8]>) -> Result<(), Error> {
         let mut tx = self.tx.lock();
 
         // Buffer still in Tx, cannot send

kernel/driver/net/rtl81xx/src/rtl8168.rs

@@ -5,15 +5,15 @@ use core::{
 
 use alloc::{sync::Arc, vec::Vec};
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
+    dma::DmaAllocator,
     interrupt::{InterruptHandler, IrqVector},
 };
-use libk::error::Error;
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    device::DeviceMemoryIo,
-    PageBox,
+use libk::{
+    dma::{BusAddress, DmaBuffer},
+    error::Error,
 };
+use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo, L3_PAGE_SIZE};
 use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
 use tock_registers::{
     interfaces::{ReadWriteable, Readable, Writeable},
@@ -218,18 +218,18 @@ enum Revision {
 struct Descriptor {
     cmd: u32,
     vlan_cmd: u32,
-    address: PhysicalAddress,
+    address: BusAddress,
 }
 
 struct RxRing {
-    entries: PageBox<[Descriptor]>,
-    buffers: Vec<PageBox<[MaybeUninit<u8>]>>,
+    entries: DmaBuffer<[Descriptor]>,
+    buffers: Vec<DmaBuffer<[MaybeUninit<u8>]>>,
     rd: usize,
 }
 
 struct TxRing {
-    entries: PageBox<[Descriptor]>,
-    buffers: Vec<Option<PageBox<[u8]>>>,
+    entries: DmaBuffer<[Descriptor]>,
+    buffers: Vec<Option<DmaBuffer<[u8]>>>,
     wr: usize,
     rd: usize,
 }
@@ -242,20 +242,18 @@ pub struct Rtl8168 {
     nic: OneTimeInit<u32>,
     rx: OneTimeInit<IrqSafeSpinlock<RxRing>>,
     tx: OneTimeInit<IrqSafeSpinlock<TxRing>>,
+
+    dma: Arc<dyn DmaAllocator>,
 }
 
 impl RxRing {
-    pub fn with_capacity(capacity: usize) -> Result<Self, Error> {
+    pub fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
         let buffers = (0..capacity)
-            .map(|_| PageBox::new_uninit_slice(4096))
-            .collect::<Result<Vec<_>, _>>()?;
-        let entries = PageBox::new_slice_with(
-            |i| {
-                Descriptor::new_rx(
-                    unsafe { buffers[i].as_physical_address() },
-                    i == capacity - 1,
-                )
-            },
+            .map(|_| DmaBuffer::new_uninit_slice(dma, L3_PAGE_SIZE))
+            .collect::<Result<Vec<_>, Error>>()?;
+        let entries = DmaBuffer::new_slice_with(
+            dma,
+            |i| Descriptor::new_rx(buffers[i].bus_address(), i == capacity - 1),
             capacity,
         )?;
         Ok(Self {
@@ -265,17 +263,21 @@ impl RxRing {
         })
     }
 
-    pub fn consume<F: Fn(PageBox<[u8]>)>(&mut self, handler: F) -> Result<usize, Error> {
+    pub fn consume<F: Fn(DmaBuffer<[u8]>)>(
+        &mut self,
+        dma: &dyn DmaAllocator,
+        handler: F,
+    ) -> Result<usize, Error> {
         let mut count = 0;
         loop {
             let index = self.rd % self.entries.len();
             let entry = &self.entries[index];
 
             if entry.is_host_owned() {
-                let new_buffer = PageBox::new_uninit_slice(4096)?;
-                let new_buffer_address = unsafe { new_buffer.as_physical_address() };
+                let new_buffer = DmaBuffer::new_uninit_slice(dma, 4096)?;
+                let new_buffer_address = new_buffer.bus_address();
                 let buffer = mem::replace(&mut self.buffers[index], new_buffer);
-                let buffer = unsafe { buffer.assume_init_slice() };
+                let buffer = unsafe { DmaBuffer::assume_init_slice(buffer) };
                 handler(buffer);
 
                 self.entries[index].setup_rx(new_buffer_address);
@@ -289,15 +291,15 @@ impl RxRing {
         Ok(count)
     }
 
-    pub fn base_address(&self) -> PhysicalAddress {
-        unsafe { self.entries.as_physical_address() }
+    pub fn base_address(&self) -> BusAddress {
+        self.entries.bus_address()
     }
 }
 
 impl TxRing {
-    pub fn with_capacity(capacity: usize) -> Result<Self, Error> {
+    pub fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
         let entries =
-            PageBox::new_slice_with(|i| Descriptor::empty_tx(i == capacity - 1), capacity)?;
+            DmaBuffer::new_slice_with(dma, |i| Descriptor::empty_tx(i == capacity - 1), capacity)?;
         let buffers = (0..capacity).map(|_| None).collect();
 
         Ok(Self {
@@ -312,8 +314,8 @@ impl TxRing {
         self.wr.wrapping_add(1) % self.entries.len() == self.rd % self.entries.len()
     }
 
-    pub fn push(&mut self, packet: PageBox<[u8]>) -> Result<(), Error> {
-        let packet_base = unsafe { packet.as_physical_address() };
+    pub fn push(&mut self, packet: DmaBuffer<[u8]>) -> Result<(), Error> {
+        let packet_base = packet.bus_address();
         let packet_size = packet.len();
 
         // TODO packet size checks
@@ -348,8 +350,8 @@ impl TxRing {
         Ok(count)
     }
 
-    pub fn base_address(&self) -> PhysicalAddress {
-        unsafe { self.entries.as_physical_address() }
+    pub fn base_address(&self) -> BusAddress {
+        self.entries.bus_address()
     }
 }
 
@@ -363,7 +365,7 @@ impl Descriptor {
     // First segment of a packet
     const CMD_FS: u32 = 1 << 29;
 
-    pub fn new_rx(buffer: PhysicalAddress, last: bool) -> Self {
+    pub fn new_rx(buffer: BusAddress, last: bool) -> Self {
         let mut cmd = Self::CMD_OWN | 0x1000;
         if last {
             cmd |= Self::CMD_END;
@@ -380,11 +382,11 @@ impl Descriptor {
         Self {
             cmd,
             vlan_cmd: 0,
-            address: PhysicalAddress::ZERO,
+            address: BusAddress::ZERO,
         }
     }
 
-    pub fn setup_rx(&mut self, buffer: PhysicalAddress) {
+    pub fn setup_rx(&mut self, buffer: BusAddress) {
         let cmd = self.cmd;
         self.address = buffer;
         self.vlan_cmd = 0;
@@ -394,7 +396,7 @@ impl Descriptor {
         }
     }
 
-    pub fn setup_tx(&mut self, buffer: PhysicalAddress, size: usize) {
+    pub fn setup_tx(&mut self, buffer: BusAddress, size: usize) {
         let mut cmd = self.cmd;
         cmd |= Self::CMD_OWN | Self::CMD_FS | Self::CMD_LS | (size as u32);
         self.address = buffer;
@@ -527,7 +529,11 @@ impl Regs {
 }
 
 impl Rtl8168 {
-    pub fn new(base: PhysicalAddress, info: PciDeviceInfo) -> Result<Self, Error> {
+    pub fn new(
+        dma: Arc<dyn DmaAllocator>,
+        base: PhysicalAddress,
+        info: PciDeviceInfo,
+    ) -> Result<Self, Error> {
         let regs = unsafe { DeviceMemoryIo::<Regs>::map(base, Default::default()) }?;
         let mac0 = regs.IDRn[0].get().to_le_bytes();
         let mac1 = regs.IDRn[1].get().to_le_bytes();
@@ -540,6 +546,8 @@ impl Rtl8168 {
             nic: OneTimeInit::new(),
             rx: OneTimeInit::new(),
             tx: OneTimeInit::new(),
+
+            dma,
         })
     }
 }
@@ -564,7 +572,7 @@ impl InterruptHandler for Rtl8168 {
             let mut rx = self.rx.get().lock();
             let nic = *self.nic.get();
             let count = rx
-                .consume(|buffer| {
+                .consume(&*self.dma, |buffer| {
                     // TODO add packet len hint to packets
                     let packet = Packet::new(buffer, 0, nic);
                     ygg_driver_net_core::receive_packet(packet).ok();
@@ -592,12 +600,12 @@ impl InterruptHandler for Rtl8168 {
 }
 
 impl Device for Rtl8168 {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         log::info!("Initialize rtl8168");
         log::info!("MAC: {}", self.mac);
 
-        let rx_ring = RxRing::with_capacity(256)?;
-        let tx_ring = TxRing::with_capacity(256)?;
+        let rx_ring = RxRing::with_capacity(&*self.dma, 256)?;
+        let tx_ring = TxRing::with_capacity(&*self.dma, 256)?;
         let rx_ring_base = rx_ring.base_address().into_u64();
         let tx_ring_base = tx_ring.base_address().into_u64();
 
@@ -686,11 +694,15 @@ impl Device for Rtl8168 {
 }
 
 impl NetworkDevice for Rtl8168 {
-    fn transmit(&self, packet: PageBox<[u8]>) -> Result<(), Error> {
+    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        DmaBuffer::new_uninit_slice(&*self.dma, len)
+    }
+
+    fn transmit_buffer(&self, buffer: DmaBuffer<[u8]>) -> Result<(), Error> {
         let mut tx = self.tx.get().lock();
         let regs = self.regs.lock();
 
-        tx.push(packet)?;
+        tx.push(buffer)?;
         regs.TPPOLL.write(TPPOLL::NPQ::SET);
 
         Ok(())

kernel/driver/usb/xhci/src/context.rs

@@ -1,9 +1,7 @@
 use core::ops::{Deref, DerefMut};
 
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox,
-};
+use device_api::dma::DmaAllocator;
+use libk::dma::{BusAddress, DmaBuffer};
 use libk_util::sync::spin_rwlock::IrqSafeRwLock;
 use xhci_lib::context::{self, DeviceHandler, InputHandler};
 use ygg_driver_usb::error::UsbError;
@@ -11,22 +9,22 @@ use ygg_driver_usb::error::UsbError;
 use crate::regs::{ContextSize, PortNumber};
 
 pub enum XhciDeviceContext {
-    Context32(IrqSafeRwLock<PageBox<context::Device32Byte>>),
-    Context64(IrqSafeRwLock<PageBox<context::Device64Byte>>),
+    Context32(IrqSafeRwLock<DmaBuffer<context::Device32Byte>>),
+    Context64(IrqSafeRwLock<DmaBuffer<context::Device64Byte>>),
 }
 
 pub enum XhciInputContext {
-    Context32(PageBox<context::Input32Byte>),
-    Context64(PageBox<context::Input64Byte>),
+    Context32(DmaBuffer<context::Input32Byte>),
+    Context64(DmaBuffer<context::Input64Byte>),
 }
 
 impl XhciDeviceContext {
-    pub fn new(size: ContextSize) -> Result<Self, UsbError> {
+    pub fn new(dma: &dyn DmaAllocator, size: ContextSize) -> Result<Self, UsbError> {
         match size {
-            ContextSize::Context32 => PageBox::new(context::Device::new_32byte())
+            ContextSize::Context32 => DmaBuffer::new(dma, context::Device::new_32byte())
                 .map(IrqSafeRwLock::new)
                 .map(Self::Context32),
-            ContextSize::Context64 => PageBox::new(context::Device::new_64byte())
+            ContextSize::Context64 => DmaBuffer::new(dma, context::Device::new_64byte())
                 .map(IrqSafeRwLock::new)
                 .map(Self::Context64),
         }
@@ -40,35 +38,38 @@ impl XhciDeviceContext {
         }
     }
 
-    pub fn physical_address(&self) -> PhysicalAddress {
+    pub fn bus_address(&self) -> BusAddress {
         match self {
-            Self::Context32(cx) => unsafe { cx.read().as_physical_address() },
-            Self::Context64(cx) => unsafe { cx.read().as_physical_address() },
+            Self::Context32(cx) => cx.read().bus_address(),
+            Self::Context64(cx) => cx.read().bus_address(),
         }
     }
+    // pub fn physical_address(&self) -> PhysicalAddress {
+    // }
 }
 
 impl XhciInputContext {
-    pub fn new(size: ContextSize) -> Result<Self, UsbError> {
+    pub fn new(dma: &dyn DmaAllocator, size: ContextSize) -> Result<Self, UsbError> {
         match size {
             ContextSize::Context32 => {
-                PageBox::new(context::Input::new_32byte()).map(Self::Context32)
+                DmaBuffer::new(dma, context::Input::new_32byte()).map(Self::Context32)
             }
             ContextSize::Context64 => {
-                PageBox::new(context::Input::new_64byte()).map(Self::Context64)
+                DmaBuffer::new(dma, context::Input::new_64byte()).map(Self::Context64)
             }
         }
         .map_err(UsbError::MemoryError)
     }
 
     pub fn new_address_device(
+        dma: &dyn DmaAllocator,
         size: ContextSize,
         root_hub_port_number: PortNumber,
         max_packet_size: usize,
         speed: u8,
-        dequeue_pointer: PhysicalAddress,
+        dequeue_pointer: BusAddress,
     ) -> Result<Self, UsbError> {
-        let mut cx = Self::new(size)?;
+        let mut cx = Self::new(dma, size)?;
 
         {
             let control = cx.control_mut();
@@ -100,10 +101,16 @@ impl XhciInputContext {
         Ok(cx)
     }
 
-    pub fn physical_address(&self) -> PhysicalAddress {
+    // pub fn physical_address(&self) -> PhysicalAddress {
+    //     match self {
+    //         Self::Context32(cx) => unsafe { cx.as_physical_address() },
+    //         Self::Context64(cx) => unsafe { cx.as_physical_address() },
+    //     }
+    // }
+    pub fn bus_address(&self) -> BusAddress {
         match self {
-            Self::Context32(cx) => unsafe { cx.as_physical_address() },
-            Self::Context64(cx) => unsafe { cx.as_physical_address() },
+            Self::Context32(cx) => cx.bus_address(),
+            Self::Context64(cx) => cx.bus_address(),
         }
     }
 }

kernel/driver/usb/xhci/src/controller.rs

@@ -1,15 +1,19 @@
-use core::sync::atomic::{AtomicU8, Ordering};
+use core::{
+    mem::MaybeUninit,
+    sync::atomic::{AtomicU8, Ordering},
+};
 
 use alloc::{boxed::Box, collections::btree_map::BTreeMap, sync::Arc, vec::Vec};
 use async_trait::async_trait;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
+    dma::DmaAllocator,
     interrupt::{InterruptHandler, IrqVector},
 };
-use libk::{error::Error, task::runtime};
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox,
+use libk::{
+    dma::{BusAddress, DmaBuffer},
+    error::Error,
+    task::runtime,
 };
 use libk_util::{
     sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
@@ -50,8 +54,8 @@ use crate::{
 
 #[allow(unused)]
 struct ScratchpadArray {
-    buffers: Vec<PageBox<[u8]>>,
-    array: PageBox<[PhysicalAddress]>,
+    buffers: Vec<DmaBuffer<[MaybeUninit<u8>]>>,
+    array: DmaBuffer<[BusAddress]>,
 }
 
 struct RootHubPort {
@@ -63,8 +67,9 @@ pub struct Xhci {
     pub(crate) regs: Regs,
     #[allow(unused)]
     pci: PciDeviceInfo,
+    pub(crate) dma: Arc<dyn DmaAllocator>,
 
-    dcbaa: IrqSafeRwLock<PageBox<[PhysicalAddress]>>,
+    dcbaa: IrqSafeRwLock<DmaBuffer<[BusAddress]>>,
     #[allow(unused)]
     scratchpads: Option<ScratchpadArray>,
     pub(crate) command_ring: CommandRing,
@@ -80,31 +85,39 @@ pub struct Xhci {
 }
 
 impl ScratchpadArray {
-    pub fn new(capacity: usize, element_size: usize) -> Result<Self, UsbError> {
+    pub fn new(
+        dma: &dyn DmaAllocator,
+        capacity: usize,
+        element_size: usize,
+    ) -> Result<Self, UsbError> {
         let buffers = (0..capacity)
-            .map(|_| PageBox::new_slice(0, element_size))
+            .map(|_| DmaBuffer::new_zeroed_slice(dma, element_size))
             .collect::<Result<Vec<_>, _>>()
             .map_err(UsbError::MemoryError)?;
-        let array =
-            PageBox::new_slice_with(|i| unsafe { buffers[i].as_physical_address() }, capacity)
+        let array = DmaBuffer::new_slice_with(dma, |i| buffers[i].bus_address(), capacity)
             .map_err(UsbError::MemoryError)?;
         Ok(Self { buffers, array })
     }
 }
 
 impl Xhci {
-    pub fn new(pci: PciDeviceInfo, regs: Regs) -> Result<Self, UsbError> {
-        let mut dcbaa = PageBox::new_slice(PhysicalAddress::ZERO, regs.slot_count + 1)
+    pub fn new(
+        dma: Arc<dyn DmaAllocator>,
+        pci: PciDeviceInfo,
+        regs: Regs,
+    ) -> Result<Self, UsbError> {
+        let mut dcbaa = DmaBuffer::new_slice(&*dma, BusAddress::ZERO, regs.slot_count + 1)
             .map_err(UsbError::MemoryError)?;
-        let command_ring = CommandRing::new(128)?;
-        let event_ring = EventRing::new(128)?;
-        let erst = EventRingSegmentTable::for_event_rings(&[&event_ring])?;
+        let command_ring = CommandRing::new(&*dma, 128)?;
+        let event_ring = EventRing::new(&*dma, 128)?;
+        let erst = EventRingSegmentTable::for_event_rings(&*dma, &[&event_ring])?;
 
         // Setup scratch buffers
         // TODO: Linux seems to just ignore the PAGESIZE, it's (1 << 0) everywhere
         let scratchpads = if regs.scratch_count != 0 {
-            let array = ScratchpadArray::new(regs.scratch_count, 0x1000)?;
-            dcbaa[0] = unsafe { array.array.as_physical_address() };
+            let array = ScratchpadArray::new(&*dma, regs.scratch_count, 0x1000)?;
+            dcbaa[0] = array.array.bus_address();
+            // dcbaa[0] = unsafe { array.array.as_physical_address() };
             Some(array)
         } else {
             None
@@ -141,6 +154,7 @@ impl Xhci {
         Ok(Self {
             regs,
             pci,
+            dma,
 
             dcbaa: IrqSafeRwLock::new(dcbaa),
             scratchpads,
@@ -157,7 +171,7 @@ impl Xhci {
         })
     }
 
-    fn notify_endpoint(&self, slot_id: u8, endpoint_id: u8, address: PhysicalAddress, status: u32) {
+    fn notify_endpoint(&self, slot_id: u8, endpoint_id: u8, address: BusAddress, status: u32) {
         if let Some(endpoint) = self.endpoints.read().get(&(slot_id, endpoint_id)) {
             endpoint.notify(address, status);
         } else {
@@ -178,12 +192,12 @@ impl Xhci {
         slot_type: u8,
         speed: UsbSpeed,
     ) -> Result<(u8, Arc<XhciBusDevice>, Arc<TransferRing>), UsbError> {
-        let device_context = XhciDeviceContext::new(self.regs.context_size)?;
+        let device_context = XhciDeviceContext::new(&*self.dma, self.regs.context_size)?;
         let slot_id = self.command_ring.enable_slot(&**self, slot_type).await?;
         let (control_pipe, control_ring) = ControlPipe::new(self.clone(), slot_id, 1, 128)?;
         let control_pipe = UsbControlPipeAccess(Box::new(control_pipe));
 
-        self.dcbaa.write()[slot_id as usize] = device_context.physical_address();
+        self.dcbaa.write()[slot_id as usize] = device_context.bus_address();
         self.endpoints
             .write()
             .insert((slot_id, 1), control_ring.clone());
@@ -270,14 +284,16 @@ impl Xhci {
             .await?;
 
         let input_cx = XhciInputContext::new_address_device(
+            &*self.dma,
             self.regs.context_size,
             number,
             max_packet_size,
             speed,
-            control_ring.base(),
+            control_ring.bus_address(),
         )?;
+
         self.command_ring
-            .address_device(&**self, slot_id, input_cx.physical_address(), false)
+            .address_device(&**self, slot_id, &input_cx, false)
             .await
             .inspect_err(|error| {
                 log::error!("Port {number} Address Device TRB (BSR=0) failed: {error:?}")
@@ -426,7 +442,7 @@ impl CommandExecutor for Xhci {
         &self,
         slot_id: u8,
         endpoint_id: u8,
-        dequeue_pointer: PhysicalAddress,
+        dequeue_pointer: BusAddress,
         dequeue_cycle: bool,
     ) -> Result<(), UsbError> {
         log::warn!("xhci: reset stalled endpoint {slot_id}:{endpoint_id}");
@@ -443,13 +459,13 @@ impl CommandExecutor for Xhci {
 }
 
 impl Device for Xhci {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         self.regs.hc_reset(10000000)?;
         log::info!("xHC reset complete");
 
         // Configure the HC
-        let dcbaap = unsafe { self.dcbaa.read().as_physical_address() };
-        let cr_base = self.command_ring.base();
+        let dcbaap = self.dcbaa.read().bus_address();
+        let cr_base = self.command_ring.bus_base();
 
         let op = self.regs.operational.write();
         let rt = self.regs.runtime.write();

kernel/driver/usb/xhci/src/device.rs

@@ -124,10 +124,12 @@ impl XhciBusDevice {
             (UsbEndpointType::Interrupt, UsbDirection::Out) => context::EndpointType::InterruptOut,
             (UsbEndpointType::Bulk, UsbDirection::In) => context::EndpointType::BulkIn,
             (UsbEndpointType::Bulk, UsbDirection::Out) => context::EndpointType::BulkOut,
-            _ => todo!(),
+            (UsbEndpointType::Isochronous, UsbDirection::In) => context::EndpointType::IsochIn,
+            (UsbEndpointType::Isochronous, UsbDirection::Out) => context::EndpointType::IsochOut,
+            (UsbEndpointType::Control, _) => context::EndpointType::Control,
         };
 
-        let mut input_cx = XhciInputContext::new(self.xhci.regs.context_size)?;
+        let mut input_cx = XhciInputContext::new(&*self.xhci.dma, self.xhci.regs.context_size)?;
 
         {
             let control = input_cx.control_mut();
@@ -150,17 +152,13 @@ impl XhciBusDevice {
             // TODO Pick a better value here
             endpoint.set_interval(3);
             endpoint.set_max_packet_size(max_packet_size);
-            endpoint.set_tr_dequeue_pointer(ring.base().into_u64());
+            endpoint.set_tr_dequeue_pointer(ring.bus_address().into_u64());
             endpoint.set_dequeue_cycle_state();
         }
 
         self.xhci
             .command_ring
-            .configure_endpoint(
-                self.xhci.as_ref(),
-                self.slot_id,
-                input_cx.physical_address(),
-            )
+            .configure_endpoint(self.xhci.as_ref(), self.slot_id, &input_cx)
             .await?;
 
         self.endpoints.write().insert(dci, ring.clone());

kernel/driver/usb/xhci/src/lib.rs

@@ -1,12 +1,12 @@
 #![no_std]
 #![allow(clippy::new_without_default)]
-#![feature(iter_array_chunks, let_chains)]
+#![feature(iter_array_chunks, let_chains, maybe_uninit_slice)]
 
 extern crate alloc;
 
 use alloc::sync::Arc;
 use controller::Xhci;
-use device_api::{device::Device, interrupt::InterruptAffinity};
+use device_api::{device::Device, dma::DmaAllocator, interrupt::InterruptAffinity};
 use regs::Regs;
 use ygg_driver_pci::{
     capability::{DevicePowerState, PowerManagementCapability},
@@ -32,7 +32,7 @@ pci_driver! {
             "xhci"
         }
 
-        fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
+        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
             // TODO Chip Hardware Reset
             let bar0 = info
                 .config_space
@@ -60,7 +60,7 @@ pci_driver! {
             info.init_interrupts(PreferredInterruptMode::Msi(true))?;
 
             let regs = Regs::map(bar0)?;
-            let xhci = Arc::new(Xhci::new(info.clone(), regs)?);
+            let xhci = Arc::new(Xhci::new(dma.clone(), info.clone(), regs)?);
 
             info.map_interrupt(InterruptAffinity::Any, xhci.clone())?;
 

kernel/driver/usb/xhci/src/pipe.rs

@@ -2,17 +2,20 @@ use core::mem::MaybeUninit;
 
 use alloc::{boxed::Box, sync::Arc};
 use async_trait::async_trait;
-use libk_mm::{address::AsPhysicalAddress, PageSlice};
+use libk::dma::{DmaBuffer, DmaSlice, DmaSliceMut};
 use ygg_driver_usb::{
-    communication::UsbDirection,
     error::{TransferError, UsbError},
     pipe::{
         control::{ControlTransferSetup, UsbControlPipe},
         normal::{UsbNormalPipeIn, UsbNormalPipeOut},
+        UsbGenericPipe,
     },
 };
 
-use crate::{controller::Xhci, ring::transfer::TransferRing};
+use crate::{
+    controller::Xhci,
+    ring::transfer::{ControlDataStage, TransferRing},
+};
 
 pub struct ControlPipe {
     xhci: Arc<Xhci>,
@@ -36,7 +39,12 @@ impl ControlPipe {
         endpoint_id: u8,
         capacity: usize,
     ) -> Result<(Self, Arc<TransferRing>), UsbError> {
-        let ring = Arc::new(TransferRing::new(slot_id, endpoint_id, capacity)?);
+        let ring = Arc::new(TransferRing::new(
+            &*xhci.dma,
+            slot_id,
+            endpoint_id,
+            capacity,
+        )?);
         Ok((
             Self {
                 xhci,
@@ -54,7 +62,12 @@ impl NormalInPipe {
         endpoint_id: u8,
         capacity: usize,
     ) -> Result<(Self, Arc<TransferRing>), UsbError> {
-        let ring = Arc::new(TransferRing::new(slot_id, endpoint_id, capacity)?);
+        let ring = Arc::new(TransferRing::new(
+            &*xhci.dma,
+            slot_id,
+            endpoint_id,
+            capacity,
+        )?);
         Ok((
             Self {
                 xhci,
@@ -72,7 +85,12 @@ impl NormalOutPipe {
         endpoint_id: u8,
         capacity: usize,
     ) -> Result<(Self, Arc<TransferRing>), UsbError> {
-        let ring = Arc::new(TransferRing::new(slot_id, endpoint_id, capacity)?);
+        let ring = Arc::new(TransferRing::new(
+            &*xhci.dma,
+            slot_id,
+            endpoint_id,
+            capacity,
+        )?);
         Ok((
             Self {
                 xhci,
@@ -85,45 +103,87 @@ impl NormalOutPipe {
 
 #[async_trait]
 impl UsbControlPipe for ControlPipe {
-    async fn control_transfer(
+    async fn control_transfer(&self, setup: ControlTransferSetup) -> Result<(), UsbError> {
+        self.ring
+            .control_transfer(self.xhci.as_ref(), setup, ControlDataStage::None)
+            .await?;
+        Ok(())
+    }
+
+    async fn control_transfer_in(
         &self,
         setup: ControlTransferSetup,
-        data: Option<(&mut PageSlice<MaybeUninit<u8>>, UsbDirection)>,
+        buffer: &mut [MaybeUninit<u8>],
     ) -> Result<usize, UsbError> {
-        let data_len = data.as_ref().map_or(0, |(data, _)| data.len());
+        let data_len = buffer.len();
+        let mut dma_buffer = DmaBuffer::new_uninit_slice(&*self.xhci.dma, data_len)
+            .map_err(UsbError::MemoryError)?;
+
         let result = self
             .ring
-            .control_transfer(self.xhci.as_ref(), setup, data)
+            .control_transfer(
+                self.xhci.as_ref(),
+                setup,
+                ControlDataStage::In(&mut dma_buffer),
+            )
             .await;
-        allow_short_packet(data_len, result)
+        let result = allow_short_packet(data_len, result);
+
+        match result {
+            Ok(len) => {
+                buffer[..len].copy_from_slice(&dma_buffer[..len]);
+                Ok(len)
+            }
+            Err(error) => Err(error),
+        }
+    }
+
+    async fn control_transfer_out(
+        &self,
+        setup: ControlTransferSetup,
+        buffer: &[u8],
+    ) -> Result<usize, UsbError> {
+        let mut dma_buffer =
+            DmaBuffer::from_slice(&*self.xhci.dma, buffer).map_err(UsbError::MemoryError)?;
+
+        self.ring
+            .control_transfer(
+                self.xhci.as_ref(),
+                setup,
+                ControlDataStage::Out(&mut dma_buffer),
+            )
+            .await
+    }
+}
+
+impl UsbGenericPipe for NormalInPipe {
+    fn allocate_dma_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, UsbError> {
+        DmaBuffer::new_uninit_slice(&*self.xhci.dma, len).map_err(UsbError::MemoryError)
     }
 }
 
 #[async_trait]
 impl UsbNormalPipeIn for NormalInPipe {
-    async fn read(&self, buffer: &mut PageSlice<u8>) -> Result<usize, UsbError> {
-        let data_len = buffer.len();
+    async fn read_dma(&self, buffer: DmaSliceMut<'_, MaybeUninit<u8>>) -> Result<usize, UsbError> {
         let result = self
             .ring
-            .normal_transfer(
-                self.xhci.as_ref(),
-                unsafe { buffer.as_physical_address() },
-                buffer.len(),
-            )
+            .normal_transfer(self.xhci.as_ref(), buffer.bus_address(), buffer.len())
             .await;
-        allow_short_packet(data_len, result)
+        allow_short_packet(buffer.len(), result)
+    }
+}
+
+impl UsbGenericPipe for NormalOutPipe {
+    fn allocate_dma_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, UsbError> {
+        DmaBuffer::new_uninit_slice(&*self.xhci.dma, len).map_err(UsbError::MemoryError)
     }
 }
 
 #[async_trait]
 impl UsbNormalPipeOut for NormalOutPipe {
-    async fn write(&self, buffer: &PageSlice<u8>) -> Result<usize, UsbError> {
+    async fn write_dma(&self, buffer: DmaSlice<'_, u8>) -> Result<usize, UsbError> {
         self.ring
-            .normal_transfer(
-                self.xhci.as_ref(),
-                unsafe { buffer.as_physical_address() },
-                buffer.len(),
-            )
+            .normal_transfer(self.xhci.as_ref(), buffer.bus_address(), buffer.len())
             .await
     }
 }

kernel/driver/usb/xhci/src/regs/operational.rs

@@ -1,5 +1,4 @@
-use libk::error::Error;
-use libk_mm::address::PhysicalAddress;
+use libk::{dma::BusAddress, error::Error};
 use tock_registers::{
     fields::FieldValue,
     interfaces::{Readable, Writeable},
@@ -116,12 +115,12 @@ impl OperationalRegs {
         }
     }
 
-    pub fn set_dcbaap(&self, value: PhysicalAddress) {
+    pub fn set_dcbaap(&self, value: BusAddress) {
         let value = value.into_u64();
         self.DCBAAP.set(value);
     }
 
-    pub fn set_crcr(&self, value: PhysicalAddress, dcs: bool) {
+    pub fn set_crcr(&self, value: BusAddress, dcs: bool) {
         let mut value = value.into_u64();
         if dcs {
             value |= 1;

kernel/driver/usb/xhci/src/regs/runtime.rs

@@ -1,11 +1,11 @@
-use libk_mm::address::PhysicalAddress;
+use libk::dma::BusAddress;
 use tock_registers::{
     interfaces::{Readable, Writeable},
     register_bitfields, register_structs,
     registers::ReadWrite,
 };
 
-use crate::ring::{EventRing, EventRingSegmentTable};
+use crate::ring::{EventRing, EventRingSegmentTable, GenericRing};
 
 register_bitfields! {
     u32,
@@ -53,8 +53,8 @@ impl RuntimeRegs {
     ) {
         let interrupter = &self.IRn[index];
 
-        let erdp = event_ring.dequeue_pointer().into_u64();
-        let erstba = erst.physical_address().into_u64();
+        let erdp = event_ring.bus_base().into_u64();
+        let erstba = erst.bus_address().into_u64();
 
         interrupter.ERSTSZ.set(erst.capacity() as u32);
         interrupter.ERSTBA.set(erstba);
@@ -62,7 +62,7 @@ impl RuntimeRegs {
         interrupter.IMAN.write(IMAN::IE::SET);
     }
 
-    pub fn set_interrupter_dequeue_pointer(&self, index: usize, erdp: PhysicalAddress) {
+    pub fn set_interrupter_dequeue_pointer(&self, index: usize, erdp: BusAddress) {
         let _ = self.IRn[index].ERDP.get();
         self.IRn[index].ERDP.set(erdp.into_u64() | (1 << 3));
     }

kernel/driver/usb/xhci/src/ring/command.rs

@@ -7,10 +7,8 @@ use core::{
 
 use alloc::collections::BTreeMap;
 use bytemuck::{Pod, Zeroable};
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox,
-};
+use device_api::dma::DmaAllocator;
+use libk::dma::{BusAddress, DmaBuffer};
 use libk_util::{
     sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
     waker::QueueWaker,
@@ -18,10 +16,12 @@ use libk_util::{
 use ygg_driver_usb::error::UsbError;
 use yggdrasil_abi::define_bitfields;
 
+use crate::context::XhciInputContext;
+
 use super::{CommandExecutor, GenericRing, LinkTrb};
 
 struct CommandRingInner {
-    trbs: PageBox<[MaybeUninit<RawCommandTrb>]>,
+    trbs: DmaBuffer<[MaybeUninit<RawCommandTrb>]>,
     enqueue_index: usize,
     #[allow(unused)]
     dequeue_index: usize,
@@ -31,23 +31,18 @@ struct CommandRingInner {
 pub struct CommandRing {
     inner: IrqSafeSpinlock<CommandRingInner>,
     // TODO maybe use Vec of "slots"?
-    completions: IrqSafeRwLock<BTreeMap<PhysicalAddress, CommandReply>>,
+    completions: IrqSafeRwLock<BTreeMap<BusAddress, CommandReply>>,
     completion_notify: QueueWaker,
-    capacity: usize,
 }
 
 impl GenericRing for CommandRing {
-    fn base(&self) -> PhysicalAddress {
-        unsafe { self.inner.lock().trbs.as_physical_address() }
-    }
-
-    fn capacity(&self) -> usize {
-        self.capacity
+    fn bus_base(&self) -> BusAddress {
+        self.inner.lock().trbs.bus_address()
     }
 }
 
 impl CommandRingInner {
-    fn enqueue<C: CommandTrb>(&mut self, trb: C) -> PhysicalAddress {
+    fn enqueue<C: CommandTrb>(&mut self, trb: C) -> BusAddress {
         let mut raw: RawCommandTrb = bytemuck::cast(trb);
 
         raw.flags.set_ty(C::TRB_TYPE as _);
@@ -55,8 +50,12 @@ impl CommandRingInner {
 
         self.trbs[self.enqueue_index].write(raw);
 
-        let address = unsafe { self.trbs.as_physical_address() }
+        let address = self
+            .trbs
+            .bus_address()
             .add(self.enqueue_index * size_of::<RawCommandTrb>());
+        // let address = unsafe { self.trbs.as_physical_address() }
+        //     .add(self.enqueue_index * size_of::<RawCommandTrb>());
 
         // Move to the next TRB slot
         self.enqueue_index += 1;
@@ -72,7 +71,7 @@ impl CommandRingInner {
     }
 
     fn enqueue_link(&mut self) {
-        let base = unsafe { self.trbs.as_physical_address() };
+        let base = self.trbs.bus_address();
 
         let link = LinkTrb::new(base, self.cycle_bit);
         self.trbs[self.enqueue_index].write(bytemuck::cast(link));
@@ -88,8 +87,8 @@ impl CommandRingInner {
 }
 
 impl CommandRing {
-    pub fn new(capacity: usize) -> Result<Self, UsbError> {
-        let trbs = PageBox::new_zeroed_slice(capacity).map_err(UsbError::MemoryError)?;
+    pub fn new(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, UsbError> {
+        let trbs = DmaBuffer::new_zeroed_slice(dma, capacity).map_err(UsbError::MemoryError)?;
 
         Ok(Self {
             inner: IrqSafeSpinlock::new(CommandRingInner {
@@ -100,11 +99,10 @@ impl CommandRing {
             }),
             completions: IrqSafeRwLock::new(BTreeMap::new()),
             completion_notify: QueueWaker::new(),
-            capacity,
         })
     }
 
-    pub fn enqueue<C: CommandTrb>(&self, trb: C) -> PhysicalAddress {
+    pub fn enqueue<C: CommandTrb>(&self, trb: C) -> BusAddress {
         let mut inner = self.inner.lock();
         inner.enqueue(trb)
     }
@@ -113,12 +111,12 @@ impl CommandRing {
         &self,
         executor: &E,
         slot_id: u8,
-        cx_physical_address: PhysicalAddress,
+        input_context: &XhciInputContext,
         bsr: bool,
     ) -> Result<(), UsbError> {
         self.submit_and_wait(
             executor,
-            AddressDeviceCommandTrb::new(cx_physical_address, slot_id, bsr),
+            AddressDeviceCommandTrb::new(input_context.bus_address(), slot_id, bsr),
         )
         .await?;
         Ok(())
@@ -128,11 +126,11 @@ impl CommandRing {
         &self,
         executor: &E,
         slot_id: u8,
-        cx_physical_address: PhysicalAddress,
+        input_context: &XhciInputContext,
     ) -> Result<(), UsbError> {
         self.submit_and_wait(
             executor,
-            ConfigureEndpointCommandTrb::new(cx_physical_address, slot_id),
+            ConfigureEndpointCommandTrb::new(input_context.bus_address(), slot_id),
         )
         .await?;
         Ok(())
@@ -158,7 +156,7 @@ impl CommandRing {
         executor: &E,
         slot_id: u8,
         endpoint_id: u8,
-        dequeue_pointer: PhysicalAddress,
+        dequeue_pointer: BusAddress,
         dequeue_cycle: bool,
     ) -> Result<(), UsbError> {
         self.submit_and_wait(
@@ -219,11 +217,11 @@ impl CommandRing {
         .await
     }
 
-    pub fn get_completion(&self, address: PhysicalAddress) -> Option<CommandReply> {
+    pub fn get_completion(&self, address: BusAddress) -> Option<CommandReply> {
         self.completions.write().remove(&address)
     }
 
-    pub fn notify(&self, address: PhysicalAddress, reply: CommandReply) {
+    pub fn notify(&self, address: BusAddress, reply: CommandReply) {
         self.inner.lock().advance();
         self.completions.write().insert(address, reply);
         self.completion_notify.wake_all();
@@ -296,7 +294,7 @@ pub struct DisableSlotCommandTrb {
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C, align(16))]
 pub struct AddressDeviceCommandTrb {
-    pub input_context_address: PhysicalAddress,
+    pub input_context_address: BusAddress,
     _0: u32,
     pub flags: AddressDeviceCommandFlags,
 }
@@ -304,7 +302,7 @@ pub struct AddressDeviceCommandTrb {
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C, align(16))]
 pub struct ConfigureEndpointCommandTrb {
-    pub input_context_address: PhysicalAddress,
+    pub input_context_address: BusAddress,
     _0: u32,
     pub flags: ConfigureEndpointCommandFlags,
 }
@@ -361,7 +359,7 @@ impl DisableSlotCommandTrb {
 }
 
 impl AddressDeviceCommandTrb {
-    pub fn new(input_context_address: PhysicalAddress, slot_id: u8, bsr: bool) -> Self {
+    pub fn new(input_context_address: BusAddress, slot_id: u8, bsr: bool) -> Self {
         Self {
             input_context_address,
             _0: 0,
@@ -371,7 +369,7 @@ impl AddressDeviceCommandTrb {
 }
 
 impl ConfigureEndpointCommandTrb {
-    pub fn new(input_context_address: PhysicalAddress, slot_id: u8) -> Self {
+    pub fn new(input_context_address: BusAddress, slot_id: u8) -> Self {
         Self {
             input_context_address,
             _0: 0,
@@ -397,7 +395,7 @@ impl SetTrDequeuePointerCommandTrb {
     pub fn new(
         slot_id: u8,
         endpoint_id: u8,
-        dequeue_pointer: PhysicalAddress,
+        dequeue_pointer: BusAddress,
         dequeue_cycle: bool,
     ) -> Self {
         let mut value = dequeue_pointer.into_u64();

kernel/driver/usb/xhci/src/ring/event.rs

@@ -1,10 +1,8 @@
 use core::mem::{size_of, MaybeUninit};
 
 use bytemuck::{Pod, Zeroable};
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox,
-};
+use device_api::dma::DmaAllocator;
+use libk::dma::{BusAddress, DmaBuffer};
 use libk_util::sync::IrqSafeSpinlock;
 use ygg_driver_usb::error::UsbError;
 use yggdrasil_abi::define_bitfields;
@@ -14,11 +12,11 @@ use super::{command::CommandReply, GenericRing};
 pub enum Event {
     PortChange(usize),
     CommandCompletion {
-        address: PhysicalAddress,
+        address: BusAddress,
         reply: CommandReply,
     },
     Transfer {
-        address: PhysicalAddress,
+        address: BusAddress,
         slot_id: u8,
         endpoint_id: u8,
         status: u32,
@@ -27,7 +25,7 @@ pub enum Event {
 
 #[repr(C, align(16))]
 pub struct EventRingSegment {
-    address: PhysicalAddress,
+    address: BusAddress,
     // Number of TRBs supported by the ring segment. Valid values are 16 to 4096
     size: u16,
     _0: u16,
@@ -35,11 +33,11 @@ pub struct EventRingSegment {
 }
 
 pub struct EventRingSegmentTable {
-    entries: PageBox<[EventRingSegment]>,
+    entries: DmaBuffer<[EventRingSegment]>,
 }
 
 struct EventRingInner {
-    trbs: PageBox<[MaybeUninit<RawEventTrb>]>,
+    trbs: DmaBuffer<[MaybeUninit<RawEventTrb>]>,
     dequeue_index: usize,
     cycle_bit: bool,
 }
@@ -50,20 +48,24 @@ pub struct EventRing {
 }
 
 impl EventRingSegmentTable {
-    pub fn for_event_rings(rings: &[&EventRing]) -> Result<Self, UsbError> {
-        let entries = PageBox::from_iter_exact(rings.iter().map(|ring| EventRingSegment {
-            address: ring.base(),
-            size: ring.capacity().try_into().unwrap(),
+    pub fn for_event_rings(dma: &dyn DmaAllocator, rings: &[&EventRing]) -> Result<Self, UsbError> {
+        let entries = DmaBuffer::new_slice_with(
+            dma,
+            |i| EventRingSegment {
+                address: rings[i].bus_base(),
+                size: rings[i].capacity.try_into().unwrap(),
                 _0: 0,
                 _1: 0,
-        }))
+            },
+            rings.len(),
+        )
         .map_err(UsbError::MemoryError)?;
 
         Ok(Self { entries })
     }
 
-    pub fn physical_address(&self) -> PhysicalAddress {
-        unsafe { self.entries.as_physical_address() }
+    pub fn bus_address(&self) -> BusAddress {
+        self.entries.bus_address()
     }
 
     pub fn capacity(&self) -> usize {
@@ -72,12 +74,8 @@ impl EventRingSegmentTable {
 }
 
 impl GenericRing for EventRing {
-    fn base(&self) -> PhysicalAddress {
-        unsafe { self.inner.lock().trbs.as_physical_address() }
-    }
-
-    fn capacity(&self) -> usize {
-        self.capacity
+    fn bus_base(&self) -> BusAddress {
+        self.inner.lock().trbs.bus_address()
     }
 }
 
@@ -103,8 +101,8 @@ impl EventRingInner {
 }
 
 impl EventRing {
-    pub fn new(capacity: usize) -> Result<Self, UsbError> {
-        let trbs = PageBox::new_zeroed_slice(capacity).map_err(UsbError::MemoryError)?;
+    pub fn new(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, UsbError> {
+        let trbs = DmaBuffer::new_zeroed_slice(dma, capacity).map_err(UsbError::MemoryError)?;
 
         Ok(Self {
             inner: IrqSafeSpinlock::new(EventRingInner {
@@ -120,9 +118,11 @@ impl EventRing {
         self.inner.lock().try_dequeue()
     }
 
-    pub fn dequeue_pointer(&self) -> PhysicalAddress {
+    pub fn dequeue_pointer(&self) -> BusAddress {
         let i = self.inner.lock();
-        unsafe { i.trbs.as_physical_address() }.add(i.dequeue_index * size_of::<RawEventTrb>())
+        i.trbs
+            .bus_address()
+            .add(i.dequeue_index * size_of::<RawEventTrb>())
     }
 }
 
@@ -171,7 +171,7 @@ define_bitfields! {
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C, align(16))]
 pub struct TransferEventTrb {
-    pub address: PhysicalAddress,
+    pub address: BusAddress,
     pub status: TransferEventStatus,
     pub flags: TransferEventFlags,
 }
@@ -179,7 +179,7 @@ pub struct TransferEventTrb {
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C, align(16))]
 pub struct CommandCompletionEventTrb {
-    pub address: PhysicalAddress,
+    pub address: BusAddress,
     pub status: CommandCompletionEventStatus,
     pub flags: CommandCompletionEventFlags,
 }

kernel/driver/usb/xhci/src/ring/mod.rs

@@ -1,11 +1,7 @@
-//use bytemuck::{Pod, Zeroable};
-//use libk_mm::address::PhysicalAddress;
-//use yggdrasil_abi::define_bitfields;
-
 use alloc::boxed::Box;
 use async_trait::async_trait;
 use bytemuck::{Pod, Zeroable};
-use libk_mm::address::PhysicalAddress;
+use libk::dma::BusAddress;
 use ygg_driver_usb::error::UsbError;
 use yggdrasil_abi::define_bitfields;
 
@@ -23,14 +19,13 @@ pub trait CommandExecutor {
         &self,
         slot_id: u8,
         endpoint_id: u8,
-        dequeue_pointer: PhysicalAddress,
+        dequeue_pointer: BusAddress,
         dequeue_cycle: bool,
     ) -> Result<(), UsbError>;
 }
 
 pub trait GenericRing {
-    fn capacity(&self) -> usize;
-    fn base(&self) -> PhysicalAddress;
+    fn bus_base(&self) -> BusAddress;
 }
 
 define_bitfields! {
@@ -45,13 +40,13 @@ define_bitfields! {
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C, align(16))]
 pub struct LinkTrb {
-    pub address: PhysicalAddress,
+    pub address: BusAddress,
     _0: u32,
     pub flags: LinkTrbFlags,
 }
 
 impl LinkTrb {
-    pub fn new(address: PhysicalAddress, cycle_bit: bool) -> Self {
+    pub fn new(address: BusAddress, cycle_bit: bool) -> Self {
         Self {
             address,
             _0: 0,

kernel/driver/usb/xhci/src/ring/transfer.rs

@@ -10,11 +10,9 @@ use alloc::{
     vec::Vec,
 };
 use bytemuck::{Pod, Zeroable};
+use device_api::dma::DmaAllocator;
 use futures_util::task::AtomicWaker;
-use libk_mm::{
-    address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox, PageSlice,
-};
+use libk::dma::{BusAddress, DmaBuffer};
 use libk_util::{
     queue::BoundedQueue,
     sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock, IrqSafeSpinlockGuard},
@@ -29,7 +27,7 @@ use yggdrasil_abi::define_bitfields;
 use super::{CommandExecutor, LinkTrb};
 
 struct TransferRingInner {
-    trbs: PageBox<[MaybeUninit<RawTransferTrb>]>,
+    trbs: DmaBuffer<[MaybeUninit<RawTransferTrb>]>,
     enqueue_index: usize,
     dequeue_index: usize,
     cycle_bit: bool,
@@ -37,7 +35,7 @@ struct TransferRingInner {
 
 pub struct TransferRing {
     inner: IrqSafeSpinlock<TransferRingInner>,
-    base: PhysicalAddress,
+    bus_base: BusAddress,
     capacity: usize,
 
     slot_id: u8,
@@ -50,7 +48,7 @@ pub struct TransferRing {
 pub struct TransactionBuilder<'a> {
     inner: IrqSafeSpinlockGuard<'a, TransferRingInner>,
     ring: &'a Arc<TransferRing>,
-    pending: Vec<PhysicalAddress>,
+    pending: Vec<BusAddress>,
 }
 
 pub struct Transaction {
@@ -66,15 +64,26 @@ pub enum TransactionEvent {
     Shutdown,
 }
 
+pub enum ControlDataStage<'a> {
+    None,
+    In(&'a mut DmaBuffer<[MaybeUninit<u8>]>),
+    Out(&'a DmaBuffer<[u8]>),
+}
+
 impl TransferRing {
-    pub fn new(slot_id: u8, endpoint_id: u8, capacity: usize) -> Result<Self, UsbError> {
-        let inner = TransferRingInner::new(capacity)?;
-        let base = unsafe { inner.trbs.as_physical_address() };
+    pub fn new(
+        dma: &dyn DmaAllocator,
+        slot_id: u8,
+        endpoint_id: u8,
+        capacity: usize,
+    ) -> Result<Self, UsbError> {
+        let inner = TransferRingInner::new(dma, capacity)?;
+        let bus_base = inner.trbs.bus_address();
         let transactions = (0..capacity).map(|_| None).collect();
 
         Ok(Self {
             inner: IrqSafeSpinlock::new(inner),
-            base,
+            bus_base,
             capacity,
 
             slot_id,
@@ -105,7 +114,7 @@ impl TransferRing {
     ) {
         if let Err(TransferError::Stall) = result {
             let dequeue = self
-                .base
+                .bus_base
                 .add(transaction.next_dequeue * size_of::<RawTransferTrb>());
             if let Err(rerror) = executor
                 .reset_endpoint(
@@ -130,7 +139,7 @@ impl TransferRing {
     pub async fn normal_transfer<E: CommandExecutor>(
         self: &Arc<Self>,
         executor: &E,
-        buffer: PhysicalAddress,
+        buffer: BusAddress,
         length: usize,
     ) -> Result<usize, UsbError> {
         if length == 0 {
@@ -154,12 +163,12 @@ impl TransferRing {
         self: &Arc<Self>,
         executor: &E,
         setup: ControlTransferSetup,
-        buffer: Option<(&mut PageSlice<MaybeUninit<u8>>, UsbDirection)>,
+        data: ControlDataStage<'_>,
     ) -> Result<usize, UsbError> {
         let mut builder = self.transaction_builder()?;
 
-        let data_len = buffer.as_ref().map_or(0, |(buffer, _)| buffer.len());
-        let (setup, data, status) = builder.enqueue_control(setup, buffer)?;
+        let data_len = data.len();
+        let (setup, data, status) = builder.enqueue_control(setup, data)?;
 
         let transaction = builder.submit(executor);
 
@@ -188,18 +197,19 @@ impl TransferRing {
         }
     }
 
-    pub fn notify(&self, address: PhysicalAddress, status: u32) {
+    pub fn notify(&self, address: BusAddress, status: u32) {
         if status == 0 {
             return;
         }
 
-        if address < self.base || address - self.base >= size_of::<RawTransferTrb>() * self.capacity
+        if address < self.bus_base
+            || address - self.bus_base >= size_of::<RawTransferTrb>() * self.capacity
         {
             log::warn!("xhci: event outside of trb array: {address:#x}");
             return;
         }
 
-        let index = (address - self.base) / size_of::<RawTransferTrb>();
+        let index = (address - self.bus_base) / size_of::<RawTransferTrb>();
         if let Some(tx) = self.transactions.write()[index]
             .take()
             .and_then(|tx| tx.upgrade())
@@ -210,8 +220,8 @@ impl TransferRing {
         }
     }
 
-    pub fn base(&self) -> PhysicalAddress {
-        self.base
+    pub fn bus_address(&self) -> BusAddress {
+        self.bus_base
     }
 }
 
@@ -221,14 +231,10 @@ impl TransactionBuilder<'_> {
     pub fn enqueue<C: TransferTrb>(&mut self, trb: C, ioc: bool) -> Result<usize, UsbError> {
         let address = self.inner.enqueue(trb, ioc)?;
         self.pending.push(address);
-        Ok((address - self.ring.base) / size_of::<RawTransferTrb>())
+        Ok((address - self.ring.bus_base) / size_of::<RawTransferTrb>())
     }
 
-    pub fn enqueue_normal(
-        &mut self,
-        buffer: PhysicalAddress,
-        length: usize,
-    ) -> Result<usize, UsbError> {
+    pub fn enqueue_normal(&mut self, buffer: BusAddress, length: usize) -> Result<usize, UsbError> {
         let trb_count = length.div_ceil(Self::TRB_SIZE_LIMIT);
         if self.inner.free_capacity() <= trb_count || trb_count == 0 {
             return Err(UsbError::DeviceBusy);
@@ -253,11 +259,11 @@ impl TransactionBuilder<'_> {
     pub fn enqueue_control(
         &mut self,
         setup: ControlTransferSetup,
-        buffer: Option<(&mut PageSlice<MaybeUninit<u8>>, UsbDirection)>,
+        buffer: ControlDataStage,
     ) -> Result<(usize, Option<usize>, usize), UsbError> {
         // Check ring capacity first
         // TODO larger DATA stages
-        let trb_count = 2 + if buffer.is_some() { 1 } else { 0 };
+        let trb_count = 2 + if buffer.len() != 0 { 1 } else { 0 };
         if self.inner.free_capacity() <= trb_count {
             return Err(UsbError::DeviceBusy);
         }
@@ -266,21 +272,33 @@ impl TransactionBuilder<'_> {
         let setup_stage = self
             .enqueue(ControlTransferSetupTrb::new(setup), true)
             .unwrap();
-        let data_stage = if let Some((buffer, direction)) = buffer {
-            Some(
-                self.enqueue(
+
+        let data_stage = match buffer {
+            ControlDataStage::None => None,
+            ControlDataStage::In(buffer) => {
+                let index = self.enqueue(
                     ControlTransferDataTrb::new(
-                        unsafe { buffer.as_physical_address() },
+                        buffer.bus_address(),
                         buffer.len(),
-                        direction,
+                        UsbDirection::In,
                     ),
                     true,
-                )
-                .unwrap(),
-            )
-        } else {
-            None
+                )?;
+                Some(index)
+            }
+            ControlDataStage::Out(buffer) => {
+                let index = self.enqueue(
+                    ControlTransferDataTrb::new(
+                        buffer.bus_address(),
+                        buffer.len(),
+                        UsbDirection::Out,
+                    ),
+                    true,
+                )?;
+                Some(index)
+            }
         };
+
         let status_stage = self
             .enqueue(ControlTransferStatusTrb::new(UsbDirection::In), true)
             .unwrap();
@@ -298,7 +316,7 @@ impl TransactionBuilder<'_> {
 
         let mut transactions = self.ring.transactions.write();
         for &pending in self.pending.iter() {
-            let index = (pending - self.ring.base) / size_of::<RawTransferTrb>();
+            let index = (pending - self.ring.bus_base) / size_of::<RawTransferTrb>();
             transactions[index] = Some(Arc::downgrade(&transaction));
         }
 
@@ -313,8 +331,9 @@ impl TransactionBuilder<'_> {
 }
 
 impl TransferRingInner {
-    fn new(capacity: usize) -> Result<Self, UsbError> {
-        let trbs = PageBox::new_zeroed_slice(capacity).map_err(UsbError::MemoryError)?;
+    fn new(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, UsbError> {
+        let trbs = DmaBuffer::new_zeroed_slice(dma, capacity).map_err(UsbError::MemoryError)?;
+
         Ok(Self {
             trbs,
             enqueue_index: 0,
@@ -323,7 +342,7 @@ impl TransferRingInner {
         })
     }
 
-    fn enqueue<C: TransferTrb>(&mut self, trb: C, ioc: bool) -> Result<PhysicalAddress, UsbError> {
+    fn enqueue<C: TransferTrb>(&mut self, trb: C, ioc: bool) -> Result<BusAddress, UsbError> {
         if (self.enqueue_index + 1) % (self.trbs.len() - 1) == self.dequeue_index {
             log::warn!("xhci: transfer ring full");
             return Err(UsbError::DeviceBusy);
@@ -337,7 +356,9 @@ impl TransferRingInner {
 
         self.trbs[self.enqueue_index].write(raw);
 
-        let address = unsafe { self.trbs.as_physical_address() }
+        let address = self
+            .trbs
+            .bus_address()
             .add(self.enqueue_index * size_of::<RawTransferTrb>());
 
         self.enqueue_index += 1;
@@ -352,7 +373,7 @@ impl TransferRingInner {
     }
 
     fn enqueue_link(&mut self) {
-        let base = unsafe { self.trbs.as_physical_address() };
+        let base = self.trbs.bus_address();
 
         let link = LinkTrb::new(base, self.cycle_bit);
         self.trbs[self.enqueue_index].write(bytemuck::cast(link));
@@ -363,6 +384,16 @@ impl TransferRingInner {
     }
 }
 
+impl ControlDataStage<'_> {
+    pub fn len(&self) -> usize {
+        match self {
+            Self::None => 0,
+            Self::In(buf) => buf.len(),
+            Self::Out(buf) => buf.len(),
+        }
+    }
+}
+
 impl Transaction {
     pub fn notify(&self, trb_index: usize, status: u32) {
         self.event_queue
@@ -505,7 +536,7 @@ define_bitfields! {
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C, align(16))]
 pub struct NormalTransferTrb {
-    pub buffer: PhysicalAddress,
+    pub buffer: BusAddress,
     pub flags: NormalTransferFlags,
 }
 
@@ -519,7 +550,7 @@ pub struct ControlTransferSetupTrb {
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C, align(16))]
 pub struct ControlTransferDataTrb {
-    pub buffer: PhysicalAddress,
+    pub buffer: BusAddress,
     pub flags: ControlTransferDataFlags,
 }
 
@@ -542,7 +573,7 @@ pub trait TransferTrb: Pod {
 }
 
 impl NormalTransferTrb {
-    pub fn new(buffer: PhysicalAddress, length: usize) -> Self {
+    pub fn new(buffer: BusAddress, length: usize) -> Self {
         Self {
             buffer,
             flags: NormalTransferFlags::new(length.try_into().unwrap()),
@@ -566,7 +597,7 @@ impl ControlTransferSetupTrb {
 }
 
 impl ControlTransferDataTrb {
-    pub fn new(buffer: PhysicalAddress, length: usize, direction: UsbDirection) -> Self {
+    pub fn new(buffer: BusAddress, length: usize, direction: UsbDirection) -> Self {
         Self {
             buffer,
             flags: ControlTransferDataFlags::new(

kernel/driver/virtio/core/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 yggdrasil-abi.workspace = true
 libk-util.workspace = true
 libk-mm.workspace = true
+libk.workspace = true
 device-api = { workspace = true, features = ["derive"] }
 
 ygg_driver_pci = { path = "../../bus/pci", optional = true }

kernel/driver/virtio/core/src/queue.rs

@@ -8,14 +8,15 @@ use core::{
     sync::atomic::{fence, Ordering},
 };
 
-use libk_mm::{address::AsPhysicalAddress, PageBox};
+use device_api::dma::DmaAllocator;
+use libk::dma::{BusAddress, DmaBuffer};
 
 use crate::{error::Error, transport::Transport};
 
 #[derive(Debug)]
 #[repr(C)]
 struct Descriptor {
-    address: u64,
+    address: BusAddress,
     len: u32,
     flags: u16,
     next: u16,
@@ -29,7 +30,7 @@ struct Descriptor {
 //      used_event: u16
 //  }
 struct AvailableRing {
-    data: PageBox<[MaybeUninit<u16>]>,
+    data: DmaBuffer<[MaybeUninit<u16>]>,
 }
 
 // Layout:
@@ -41,13 +42,12 @@ struct AvailableRing {
 //      _pad: u16
 //  }
 struct UsedRing {
-    data: PageBox<[MaybeUninit<u32>]>,
-
+    data: DmaBuffer<[MaybeUninit<u32>]>,
     used_count: usize,
 }
 
 pub struct VirtQueue {
-    descriptor_table: PageBox<[MaybeUninit<Descriptor>]>,
+    descriptor_table: DmaBuffer<[MaybeUninit<Descriptor>]>,
     available: AvailableRing,
     used: UsedRing,
 
@@ -63,8 +63,12 @@ pub struct VirtQueue {
 }
 
 impl AvailableRing {
-    pub fn with_capacity(no_irq: bool, capacity: usize) -> Result<Self, Error> {
-        let mut data = PageBox::new_zeroed_slice(capacity + 3)?;
+    pub fn with_capacity(
+        dma: &dyn DmaAllocator,
+        no_irq: bool,
+        capacity: usize,
+    ) -> Result<Self, Error> {
+        let mut data = DmaBuffer::new_zeroed_slice(dma, capacity + 3)?;
 
         if no_irq {
             data[0].write(1);
@@ -85,8 +89,8 @@ impl AvailableRing {
 }
 
 impl UsedRing {
-    pub fn with_capacity(capacity: usize) -> Result<Self, Error> {
-        let mut data = PageBox::new_zeroed_slice(capacity * 2 + 2)?;
+    pub fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
+        let mut data = DmaBuffer::new_zeroed_slice(dma, capacity * 2 + 2)?;
 
         data[0].write(0);
 
@@ -110,6 +114,7 @@ impl UsedRing {
 impl VirtQueue {
     pub fn with_capacity<T: Transport>(
         transport: &mut T,
+        dma: &dyn DmaAllocator,
         index: u16,
         capacity: usize,
         msix_vector: Option<u16>,
@@ -127,16 +132,16 @@ impl VirtQueue {
             return Err(Error::QueueTooLarge);
         }
 
-        let descriptor_table = PageBox::new_zeroed_slice(capacity)?;
-        let available = AvailableRing::with_capacity(no_avail_irq, capacity)?;
-        let used = UsedRing::with_capacity(capacity)?;
+        let descriptor_table = DmaBuffer::new_zeroed_slice(dma, capacity)?;
+        let available = AvailableRing::with_capacity(dma, no_avail_irq, capacity)?;
+        let used = UsedRing::with_capacity(dma, capacity)?;
 
         transport.set_queue(
             index,
             capacity as u16,
-            unsafe { descriptor_table.as_physical_address() },
-            unsafe { available.data.as_physical_address() },
-            unsafe { used.data.as_physical_address() },
+            descriptor_table.bus_address(),
+            available.data.bus_address(),
+            used.data.bus_address(),
             msix_vector,
         );
 
@@ -163,6 +168,7 @@ impl VirtQueue {
 
     pub fn with_max_capacity<T: Transport>(
         transport: &mut T,
+        dma: &dyn DmaAllocator,
         index: u16,
         capacity: usize,
         msix_vector: Option<u16>,
@@ -171,16 +177,16 @@ impl VirtQueue {
         let max_capacity = transport.max_queue_size(index);
         let capacity = capacity.min(max_capacity as usize);
 
-        Self::with_capacity(transport, index, capacity, msix_vector, no_avail_irq)
+        Self::with_capacity(transport, dma, index, capacity, msix_vector, no_avail_irq)
     }
 
     /// # Safety
     ///
-    /// Invariants: PageBox remains valid and allocated until it is properly dequeued.
+    /// Invariants: DmaBuffer remains valid and allocated until it is properly dequeued.
     pub unsafe fn add<'a, 'b>(
         &mut self,
-        input: &'a [&'b mut PageBox<[u8]>],
-        output: &'a [&'b PageBox<[u8]>],
+        input: &'a [&'b mut DmaBuffer<[MaybeUninit<u8>]>],
+        output: &'a [&'b DmaBuffer<[u8]>],
     ) -> Result<u16, Error> {
         if input.is_empty() && output.is_empty() {
             return Err(Error::EmptyTransaction);
@@ -209,8 +215,8 @@ impl VirtQueue {
 
     unsafe fn add_direct<'a, 'b>(
         &mut self,
-        input: &'a [&'b mut PageBox<[u8]>],
-        output: &'a [&'b PageBox<[u8]>],
+        input: &'a [&'b mut DmaBuffer<[MaybeUninit<u8>]>],
+        output: &'a [&'b DmaBuffer<[u8]>],
     ) -> u16 {
         let head = self.free_head;
         let mut last = self.free_head;
@@ -221,7 +227,7 @@ impl VirtQueue {
             let next = (self.free_head + 1) % self.capacity as u16;
 
             desc.write(Descriptor {
-                address: item.as_physical_address().into(),
+                address: item.bus_address(),
                 len: item.len().try_into().unwrap(),
                 // TODO
                 flags: (1 << 0),
@@ -238,7 +244,7 @@ impl VirtQueue {
             let next = (self.free_head + 1) % self.capacity as u16;
 
             desc.write(Descriptor {
-                address: item.as_physical_address().into(),
+                address: item.bus_address(),
                 len: item.len().try_into().unwrap(),
                 // TODO MAGIC
                 flags: (1 << 0) | (1 << 1),
@@ -265,8 +271,8 @@ impl VirtQueue {
 
     pub fn add_notify_wait_pop<'a, 'b, T: Transport>(
         &mut self,
-        input: &'a [&'b mut PageBox<[u8]>],
-        output: &'a [&'b PageBox<[u8]>],
+        input: &'a [&'b mut DmaBuffer<[MaybeUninit<u8>]>],
+        output: &'a [&'b DmaBuffer<[u8]>],
         transport: &mut T,
     ) -> Result<u32, Error> {
         let token = unsafe { self.add(input, output) }?;
@@ -333,7 +339,7 @@ impl VirtQueue {
             assert_ne!(current.len, 0);
             let next_head = (current.flags & (1 << 0) != 0).then_some(current.next);
 
-            current.address = 0;
+            current.address = BusAddress::ZERO;
             current.flags = 0;
             current.next = 0;
             current.len = 0;

kernel/driver/virtio/core/src/transport/mod.rs

@@ -1,6 +1,7 @@
 use core::mem::size_of;
 
-use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
+use libk::dma::BusAddress;
+use libk_mm::device::DeviceMemoryIo;
 use tock_registers::{
     interfaces::{Readable, Writeable},
     registers::WriteOnly,
@@ -56,17 +57,17 @@ pub trait Transport: Send {
         &mut self,
         queue: u16,
         capacity: u16,
-        descriptor_table_phys: PhysicalAddress,
-        available_ring_phys: PhysicalAddress,
-        used_ring_phys: PhysicalAddress,
+        descriptor_table_phys: BusAddress,
+        available_ring_phys: BusAddress,
+        used_ring_phys: BusAddress,
         msix_vector: Option<u16>,
     ) {
         let cfg = self.common_cfg();
         cfg.queue_select.set(queue);
         cfg.queue_size.set(capacity);
-        cfg.queue_desc.set(descriptor_table_phys.into());
-        cfg.queue_driver.set(available_ring_phys.into());
-        cfg.queue_device.set(used_ring_phys.into());
+        cfg.queue_desc.set(descriptor_table_phys.into_u64());
+        cfg.queue_driver.set(available_ring_phys.into_u64());
+        cfg.queue_device.set(used_ring_phys.into_u64());
         if self.supports_msix() {
             cfg.queue_msix_vector.set(msix_vector.unwrap_or(0xFFFF));
         } else {

kernel/driver/virtio/gpu/src/command.rs

@@ -1,6 +1,12 @@
+use core::mem::MaybeUninit;
+
 use bytemuck::{Pod, Zeroable};
-use libk::{device::display::PixelFormat, error::Error};
-use libk_mm::{address::PhysicalAddress, PageBox};
+use device_api::dma::DmaAllocator;
+use libk::{
+    device::display::PixelFormat,
+    dma::{BusAddress, DmaBuffer},
+    error::Error,
+};
 use libk_util::sync::IrqSafeSpinlockGuard;
 use ygg_driver_virtio_core::{queue::VirtQueue, transport::Transport};
 
@@ -102,10 +108,12 @@ impl<'a, T: Transport> ControlLock<'a, T> {
 
     fn send_recv<'r, Req: Pod>(
         &mut self,
+        dma: &dyn DmaAllocator,
         req: &Req,
-        buffer: &'r mut PageBox<[u8]>,
+        buffer: &'r mut DmaBuffer<[MaybeUninit<u8>]>,
     ) -> Result<(&'r ControlHeader, &'r [u8]), Error> {
-        let mut request = PageBox::new_slice(0u8, size_of::<Req>())?;
+        let request = DmaBuffer::new_uninit_slice(dma, size_of::<Req>())?;
+        let mut request = unsafe { DmaBuffer::assume_init_slice(request) };
         request.copy_from_slice(bytemuck::bytes_of(req));
 
         let len = self
@@ -121,18 +129,20 @@ impl<'a, T: Transport> ControlLock<'a, T> {
             return Err(Error::InvalidArgument);
         }
 
-        let header = bytemuck::from_bytes(&buffer[..size_of::<ControlHeader>()]);
-        let data = &buffer[size_of::<ControlHeader>()..len];
+        let payload = unsafe { MaybeUninit::slice_assume_init_ref(&buffer[..len]) };
+        let header = bytemuck::from_bytes(&payload[..size_of::<ControlHeader>()]);
+        let data = &payload[size_of::<ControlHeader>()..len];
 
         Ok((header, data))
     }
 
     fn send_recv_no_data<Req: Pod>(
         &mut self,
+        dma: &dyn DmaAllocator,
         req: &Req,
-        buffer: &mut PageBox<[u8]>,
+        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
     ) -> Result<(), Error> {
-        let (response, _) = self.send_recv(req, buffer)?;
+        let (response, _) = self.send_recv(dma, req, buffer)?;
         if response.ty == 0x1100 {
             Ok(())
         } else {
@@ -142,8 +152,9 @@ impl<'a, T: Transport> ControlLock<'a, T> {
 
     pub fn query_scanouts<'r>(
         &mut self,
+        dma: &dyn DmaAllocator,
         max_scanouts: usize,
-        buffer: &'r mut PageBox<[u8]>,
+        buffer: &'r mut DmaBuffer<[MaybeUninit<u8>]>,
     ) -> Result<&'r [ScanoutInfo], Error> {
         let request = ControlHeader {
             ty: 0x0100,
@@ -154,7 +165,7 @@ impl<'a, T: Transport> ControlLock<'a, T> {
             _0: [0; 3],
         };
 
-        let (response, data) = self.send_recv(&request, buffer)?;
+        let (response, data) = self.send_recv(dma, &request, buffer)?;
         if response.ty != 0x1101 {
             log::warn!("virtio-gpu: query_scanouts returned {:#x}", response.ty);
             return Err(Error::InvalidArgument);
@@ -169,7 +180,8 @@ impl<'a, T: Transport> ControlLock<'a, T> {
 
     pub fn create_resource_2d(
         &mut self,
-        buffer: &mut PageBox<[u8]>,
+        dma: &dyn DmaAllocator,
+        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
         width: u32,
         height: u32,
         pixel_format: PixelFormat,
@@ -189,16 +201,17 @@ impl<'a, T: Transport> ControlLock<'a, T> {
             format,
         };
 
-        self.send_recv_no_data(&request, buffer)?;
+        self.send_recv_no_data(dma, &request, buffer)?;
 
         Ok(1)
     }
 
     pub fn attach_backing(
         &mut self,
-        buffer: &mut PageBox<[u8]>,
+        dma: &dyn DmaAllocator,
+        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
         resource_id: u32,
-        base: PhysicalAddress,
+        base: BusAddress,
         length: u32,
     ) -> Result<(), Error> {
         let request = ResourceAttachBacking {
@@ -215,12 +228,13 @@ impl<'a, T: Transport> ControlLock<'a, T> {
             },
         };
 
-        self.send_recv_no_data(&request, buffer)
+        self.send_recv_no_data(dma, &request, buffer)
     }
 
     pub fn set_scanout(
         &mut self,
-        buffer: &mut PageBox<[u8]>,
+        dma: &dyn DmaAllocator,
+        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
         scanout_id: u32,
         resource_id: u32,
         width: u32,
@@ -241,12 +255,13 @@ impl<'a, T: Transport> ControlLock<'a, T> {
             resource_id,
         };
 
-        self.send_recv_no_data(&request, buffer)
+        self.send_recv_no_data(dma, &request, buffer)
     }
 
     pub fn transfer_to_host_2d(
         &mut self,
-        buffer: &mut PageBox<[u8]>,
+        dma: &dyn DmaAllocator,
+        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
         resource_id: u32,
         r: Rect,
     ) -> Result<(), Error> {
@@ -261,12 +276,13 @@ impl<'a, T: Transport> ControlLock<'a, T> {
             _0: 0,
         };
 
-        self.send_recv_no_data(&request, buffer)
+        self.send_recv_no_data(dma, &request, buffer)
     }
 
     pub fn resource_flush(
         &mut self,
-        buffer: &mut PageBox<[u8]>,
+        dma: &dyn DmaAllocator,
+        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
         resource_id: u32,
         r: Rect,
     ) -> Result<(), Error> {
@@ -280,6 +296,6 @@ impl<'a, T: Transport> ControlLock<'a, T> {
             _0: 0,
         };
 
-        self.send_recv_no_data(&request, buffer)
+        self.send_recv_no_data(dma, &request, buffer)
     }
 }

kernel/driver/virtio/gpu/src/lib.rs

@@ -1,3 +1,4 @@
+#![feature(maybe_uninit_slice)]
 #![no_std]
 
 extern crate alloc;
@@ -6,18 +7,23 @@ use core::mem::MaybeUninit;
 
 use alloc::{sync::Arc, vec::Vec};
 use command::{ControlLock, ScanoutInfo};
-use device_api::device::Device;
-use libk::device::{
+use device_api::{
+    device::{Device, DeviceInitContext},
+    dma::DmaAllocator,
+};
+use libk::{
+    device::{
         display::{
             DisplayDevice, DisplayMode, DisplayOwner, DriverFlags, FramebufferInfo, PixelFormat,
         },
         manager::DEVICE_REGISTRY,
+    },
+    dma::DmaBuffer,
 };
 use libk_mm::{
-    address::{PhysicalAddress, Virtualize},
-    phys,
+    address::{AsPhysicalAddress, PhysicalAddress},
     table::MapAttributes,
-    PageBox, PageProvider, L3_PAGE_SIZE,
+    PageProvider, L3_PAGE_SIZE,
 };
 use libk_util::{
     sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
@@ -42,9 +48,7 @@ struct Framebuffer {
     resource_id: u32,
 
     double: bool,
-    base: PhysicalAddress,
-    page_count: usize,
-    kernel_base: usize,
+    dma_buffer: DmaBuffer<[MaybeUninit<u8>]>,
     stride: usize,
     size: usize,
 }
@@ -54,7 +58,7 @@ struct Config {
     framebuffer: Option<Framebuffer>,
 
     owner: DisplayOwner,
-    response: PageBox<[u8]>,
+    response: DmaBuffer<[MaybeUninit<u8>]>,
 }
 
 pub struct VirtioGpu<T: Transport> {
@@ -65,11 +69,17 @@ pub struct VirtioGpu<T: Transport> {
     queues: OneTimeInit<Queues>,
     config: IrqSafeRwLock<Config>,
 
+    dma: Arc<dyn DmaAllocator>,
+
     num_scanouts: usize,
 }
 
 impl<T: Transport + 'static> VirtioGpu<T> {
-    pub fn new(transport: T, info: Option<PciDeviceInfo>) -> Result<Self, Error> {
+    pub fn new(
+        dma: Arc<dyn DmaAllocator>,
+        transport: T,
+        info: Option<PciDeviceInfo>,
+    ) -> Result<Self, Error> {
         // Read num-scanouts from device config
         let Some(device_cfg) = transport.device_cfg() else {
             log::error!("virtio-gpu must have device-specific configuration section");
@@ -93,10 +103,12 @@ impl<T: Transport + 'static> VirtioGpu<T> {
             config: IrqSafeRwLock::new(Config {
                 scanouts: Vec::new(),
                 framebuffer: None,
-                response: PageBox::new_slice(0, 4096)?,
+                response: DmaBuffer::new_uninit_slice(&*dma, L3_PAGE_SIZE)?,
                 owner: DisplayOwner::None,
             }),
 
+            dma,
+
             num_scanouts: num_scanouts as usize,
         })
     }
@@ -141,7 +153,7 @@ impl<T: Transport + 'static> VirtioGpu<T> {
         // TODO cursorq
         let mut transport = self.transport.lock();
 
-        let control = VirtQueue::with_max_capacity(&mut *transport, 0, 128, None, true)
+        let control = VirtQueue::with_max_capacity(&mut *transport, &*self.dma, 0, 128, None, true)
             .map_err(|_| Error::InvalidArgument)?;
 
         self.queues.init(Queues {
@@ -163,7 +175,8 @@ impl<T: Transport + 'static> VirtioGpu<T> {
         let mut control = self.control();
         let mut config = self.config.write();
 
-        let scanouts = control.query_scanouts(self.num_scanouts, &mut config.response)?;
+        let scanouts =
+            control.query_scanouts(&*self.dma, self.num_scanouts, &mut config.response)?;
         for (i, scanout) in scanouts.iter().enumerate() {
             log::info!(
                 "virtio-gpu: [{i}] {}x{} + {},{}",
@@ -199,21 +212,32 @@ impl<T: Transport + 'static> VirtioGpu<T> {
         let stride = w as usize * size_of::<u32>();
         let size = stride * h as usize;
 
-        let page_count = size.div_ceil(L3_PAGE_SIZE);
-        let base = phys::alloc_pages_contiguous(page_count)?;
-        let kernel_base = base.virtualize();
+        let dma_buffer = DmaBuffer::new_uninit_slice(&*self.dma, size)?;
 
         let mut control = self.control();
 
-        let resource_id =
-            control.create_resource_2d(&mut config.response, w, h, PixelFormat::R8G8B8A8)?;
+        let resource_id = control.create_resource_2d(
+            &*self.dma,
+            &mut config.response,
+            w,
+            h,
+            PixelFormat::R8G8B8A8,
+        )?;
         control.attach_backing(
+            &*self.dma,
             &mut config.response,
             resource_id,
-            base,
+            dma_buffer.bus_address(),
             size.try_into().unwrap(),
         )?;
-        control.set_scanout(&mut config.response, index as u32, resource_id, w, h)?;
+        control.set_scanout(
+            &*self.dma,
+            &mut config.response,
+            index as u32,
+            resource_id,
+            w,
+            h,
+        )?;
 
         config.framebuffer = Some(Framebuffer {
             scanout_index: index,
@@ -221,10 +245,9 @@ impl<T: Transport + 'static> VirtioGpu<T> {
 
             resource_id,
             size,
-            page_count,
             stride,
-            base,
-            kernel_base,
+
+            dma_buffer,
         });
 
         Ok(())
@@ -244,8 +267,8 @@ impl<T: Transport + 'static> VirtioGpu<T> {
         } else {
             let resource_id = framebuffer.resource_id;
 
-            control.transfer_to_host_2d(&mut config.response, resource_id, r)?;
-            control.resource_flush(&mut config.response, resource_id, r)?;
+            control.transfer_to_host_2d(&*self.dma, &mut config.response, resource_id, r)?;
+            control.resource_flush(&*self.dma, &mut config.response, resource_id, r)?;
 
             Ok(())
         }
@@ -253,7 +276,7 @@ impl<T: Transport + 'static> VirtioGpu<T> {
 }
 
 impl<T: Transport + 'static> Device for VirtioGpu<T> {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let status = self.begin_init()?;
         self.setup_queues()?;
         self.finish_init(status);
@@ -278,7 +301,7 @@ impl<T: Transport + 'static> PageProvider for VirtioGpu<T> {
         // TODO check that the page is mapped by framebuffer owner
         let config = self.config.read();
         let framebuffer = config.framebuffer.as_ref().ok_or(Error::DoesNotExist)?;
-        if offset as usize + L3_PAGE_SIZE > framebuffer.page_count * L3_PAGE_SIZE {
+        if offset as usize + L3_PAGE_SIZE > framebuffer.dma_buffer.page_count() * L3_PAGE_SIZE {
             log::warn!(
                 "virtio-gpu: offset {:#x} outside of framebuffer bounds {:#x}",
                 offset,
@@ -286,7 +309,7 @@ impl<T: Transport + 'static> PageProvider for VirtioGpu<T> {
             );
             return Err(Error::InvalidMemoryOperation);
         }
-        let phys = framebuffer.base.add(offset as usize);
+        let phys = unsafe { framebuffer.dma_buffer.as_physical_address() }.add(offset as usize);
 
         Ok(phys)
     }
@@ -358,8 +381,8 @@ impl<T: Transport + 'static> DisplayDevice for VirtioGpu<T> {
             }
 
             output[0].write(FramebufferInfo {
-                base: framebuffer.base,
-                kernel_base: Some(framebuffer.kernel_base),
+                base: unsafe { framebuffer.dma_buffer.as_physical_address() },
+                kernel_base: Some(framebuffer.dma_buffer.as_ptr().addr()),
                 stride: framebuffer.stride,
                 size: framebuffer.size,
             });
@@ -396,7 +419,7 @@ pci_driver! {
             "virtio-gpu"
         }
 
-        fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
+        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
             let space = &info.config_space;
 
             let transport = PciTransport::from_config_space(space)
@@ -404,9 +427,8 @@ pci_driver! {
                     log::error!("Couldn't set up PCI virtio transport: {error:?}");
                 })
                 .map_err(|_| Error::InvalidArgument)?;
-            let device = VirtioGpu::new(transport, Some(info.clone()))?;
+            let device = VirtioGpu::new(dma.clone(), transport, Some(info.clone()))?;
             let device = Arc::new(device);
-            // let device = Box::leak(Box::new(device));
 
             Ok(device)
         }

kernel/driver/virtio/net/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 yggdrasil-abi.workspace = true
 libk-util.workspace = true
 libk-mm.workspace = true
+libk.workspace = true
 device-api = { workspace = true, features = ["derive"] }
 
 ygg_driver_virtio_core = { path = "../core" }

kernel/driver/virtio/net/src/lib.rs

@@ -3,15 +3,16 @@
 
 extern crate alloc;
 
-use core::mem::size_of;
+use core::mem::{size_of, MaybeUninit};
 
 use alloc::{collections::BTreeMap, sync::Arc};
 use bytemuck::{Pod, Zeroable};
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
+    dma::DmaAllocator,
     interrupt::{InterruptAffinity, InterruptHandler, IrqVector},
 };
-use libk_mm::PageBox;
+use libk::dma::DmaBuffer;
 use libk_util::{
     sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock, IrqSafeSpinlockGuard},
     OneTimeInit,
@@ -43,7 +44,8 @@ pub struct VirtioNet<T: Transport> {
 
     mac: IrqSafeRwLock<MacAddress>,
 
-    pending_packets: IrqSafeRwLock<BTreeMap<u16, PageBox<[u8]>>>,
+    pending_packets: IrqSafeRwLock<BTreeMap<u16, DmaBuffer<[MaybeUninit<u8>]>>>,
+    dma: Arc<dyn DmaAllocator>,
 
     pci_device_info: Option<PciDeviceInfo>,
 }
@@ -71,7 +73,11 @@ impl Queues {
 impl<T: Transport + 'static> VirtioNet<T> {
     const PACKET_SIZE: usize = 4096;
 
-    pub fn new(transport: T, pci_device_info: Option<PciDeviceInfo>) -> Self {
+    pub fn new(
+        dma: Arc<dyn DmaAllocator>,
+        transport: T,
+        pci_device_info: Option<PciDeviceInfo>,
+    ) -> Self {
         // Read MAC from device config
         let device_cfg = transport
             .device_cfg()
@@ -90,6 +96,7 @@ impl<T: Transport + 'static> VirtioNet<T> {
             pending_packets: IrqSafeRwLock::new(BTreeMap::new()),
 
             pci_device_info,
+            dma,
         }
     }
 
@@ -99,7 +106,7 @@ impl<T: Transport + 'static> VirtioNet<T> {
         let mut packets = self.pending_packets.write();
 
         for _ in 0..buffers {
-            let mut packet = PageBox::new_slice(0, Self::PACKET_SIZE).unwrap();
+            let mut packet = DmaBuffer::new_uninit_slice(&*self.dma, Self::PACKET_SIZE).unwrap();
             let token = unsafe { queue.add(&[&mut packet], &[]).unwrap() };
             packets.insert(token, packet);
         }
@@ -117,11 +124,12 @@ impl<T: Transport + 'static> VirtioNet<T> {
             let mut pending_packets = self.pending_packets.write();
             let packet = pending_packets.remove(&token).unwrap();
 
-            let mut buffer = PageBox::new_slice(0, Self::PACKET_SIZE).unwrap();
+            let mut buffer = DmaBuffer::new_uninit_slice(&*self.dma, Self::PACKET_SIZE).unwrap();
 
             let token = unsafe { queue.add(&[&mut buffer], &[]).unwrap() };
             pending_packets.insert(token, buffer);
 
+            let packet = unsafe { DmaBuffer::assume_init_slice(packet) };
             let packet = Packet::new(packet, size_of::<VirtioPacketHeader>(), interface_id);
             ygg_driver_net_core::receive_packet(packet).unwrap();
             count += 1
@@ -191,10 +199,17 @@ impl<T: Transport + 'static> VirtioNet<T> {
         let mut transport = self.transport.lock();
 
         // Setup the virtqs
-        let rx = VirtQueue::with_max_capacity(&mut *transport, 0, 128, receive_vector, false)
+        let rx = VirtQueue::with_max_capacity(
+            &mut *transport,
+            &*self.dma,
+            0,
+            128,
+            receive_vector,
+            false,
+        )
+        .map_err(cvt_error)?;
+        let tx = VirtQueue::with_max_capacity(&mut *transport, &*self.dma, 1, 128, None, true)
             .map_err(cvt_error)?;
-        let tx =
-            VirtQueue::with_max_capacity(&mut *transport, 1, 128, None, true).map_err(cvt_error)?;
 
         self.queues.init(Queues {
             receive: IrqSafeSpinlock::new(rx),
@@ -206,7 +221,11 @@ impl<T: Transport + 'static> VirtioNet<T> {
 }
 
 impl<T: Transport + 'static> NetworkDevice for VirtioNet<T> {
-    fn transmit(&self, mut packet: PageBox<[u8]>) -> Result<(), Error> {
+    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        DmaBuffer::new_uninit_slice(&*self.dma, len)
+    }
+
+    fn transmit_buffer(&self, mut packet: DmaBuffer<[u8]>) -> Result<(), Error> {
         let queues = self.queues.get();
         let mut tx = queues.transmit.lock();
         let mut transport = self.transport.lock();
@@ -214,7 +233,6 @@ impl<T: Transport + 'static> NetworkDevice for VirtioNet<T> {
         let _len = tx
             .add_notify_wait_pop(&[], &[&packet], &mut *transport)
             .unwrap();
-
         Ok(())
     }
 
@@ -253,7 +271,7 @@ impl<T: Transport + 'static> Device for VirtioNet<T> {
         "VirtIO Network Device"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let status = self.begin_init()?;
 
         // TODO multiqueue
@@ -289,11 +307,11 @@ pci_driver! {
             "virtio-net"
         }
 
-        fn probe(&self, info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
+        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
             let space = &info.config_space;
 
             let transport = PciTransport::from_config_space(space).unwrap();
-            let device = VirtioNet::new(transport, Some(info.clone()));
+            let device = VirtioNet::new(dma.clone(), transport, Some(info.clone()));
 
             let device = Arc::new(device);
 

kernel/lib/device-api/src/device.rs

@@ -1,7 +1,11 @@
 use alloc::sync::Arc;
 use yggdrasil_abi::error::Error;
 
-use crate::{bus::Bus, clock::ClockController};
+use crate::{bus::Bus, clock::ClockController, dma::DmaAllocator};
+
+pub struct DeviceInitContext {
+    pub dma_allocator: Arc<dyn DmaAllocator>,
+}
 
 pub trait Device: Sync + Send {
     fn display_name(&self) -> &str;
@@ -13,7 +17,8 @@ pub trait Device: Sync + Send {
     /// # Safety
     ///
     /// The caller must make sure the function is only called once.
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, cx: DeviceInitContext) -> Result<(), Error> {
+        let _ = cx;
         Ok(())
     }
 

kernel/lib/device-api/src/dma.rs

@@ -0,0 +1,15 @@
+use core::{alloc::Layout, ptr::NonNull};
+
+use yggdrasil_abi::error::Error;
+
+#[derive(Debug)]
+pub struct DmaAllocation {
+    pub host_virtual: NonNull<()>,
+    pub host_physical: u64,
+    pub bus_address: u64,
+    pub page_count: usize,
+}
+
+pub trait DmaAllocator: Send + Sync {
+    fn allocate(&self, layout: Layout) -> Result<DmaAllocation, Error>;
+}

kernel/lib/device-api/src/lib.rs

@@ -11,6 +11,8 @@ pub mod interrupt;
 pub mod serial;
 pub mod timer;
 
+pub mod dma;
+
 use device::Device;
 use yggdrasil_abi::error::Error;
 

kernel/lib/device-tree/src/driver/tree.rs

@@ -8,10 +8,11 @@ use alloc::{
 use device_api::{
     bus::Bus,
     clock::ClockController,
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{ExternalInterruptController, FullIrq, MessageInterruptController},
 };
 use fdt_rs::spec::Phandle;
+use libk::dma::DummyDmaAllocator;
 use libk_mm::address::PhysicalAddress;
 use libk_util::OneTimeInit;
 use yggdrasil_abi::error::Error;
@@ -188,7 +189,8 @@ impl Node {
     pub fn lazy_init(self: Arc<Self>) -> Option<Result<(), Error>> {
         let device = self.clone().probe()?;
         let result = self.init_token.or_try_init_with(|| {
-            unsafe { device.init() }?;
+            let cx = self.make_init_context();
+            unsafe { device.init(cx) }?;
             Ok(())
         });
         match result {
@@ -233,12 +235,21 @@ impl Node {
             .inspect_err(|_| log::error!("Does not exist: probe({:?})", self.name))?;
 
         self.init_token.try_init_with_opt(|| {
-            unsafe { device.init() }?;
+            let cx = self.make_init_context();
+            unsafe { device.init(cx) }?;
             Ok(())
         })?;
         Ok(())
     }
 
+    fn make_init_context(&self) -> DeviceInitContext {
+        let cx = DeviceInitContext {
+            dma_allocator: Arc::new(DummyDmaAllocator),
+        };
+
+        cx
+    }
+
     /// Returns an iterator over the node's children
     pub fn children(&self) -> impl Iterator<Item = &Arc<Node>> {
         self.children.get().iter()

kernel/libk/src/device/block/cache.rs

@@ -1,48 +1,37 @@
 #![allow(clippy::missing_transmute_annotations)]
 use core::{
-    marker::PhantomData,
+    mem::MaybeUninit,
     ops::{Deref, DerefMut},
 };
 
 use alloc::sync::Arc;
-use kernel_arch::mem::PhysicalMemoryAllocator;
-use libk_mm::{address::PhysicalAddress, phys::GlobalPhysicalAllocator, PageBox};
 use libk_util::{lru_hash_table::LruCache, sync::spin_rwlock::IrqSafeRwLock};
 use yggdrasil_abi::error::Error;
 
-use crate::task::sync::AsyncMutex;
+use crate::{dma::DmaBuffer, task::sync::AsyncMutex};
 
 use super::BlockDevice;
 
-pub struct CachedSegment<
-    A: PhysicalMemoryAllocator<Address = PhysicalAddress> = GlobalPhysicalAllocator,
-> {
-    data: PageBox<[u8], A>,
+pub struct CachedSegment {
+    data: DmaBuffer<[u8]>,
     dirty: bool,
 }
 
-pub struct UncachedCache<
-    A: PhysicalMemoryAllocator<Address = PhysicalAddress> = GlobalPhysicalAllocator,
-> {
+pub struct UncachedCache {
     device: Arc<dyn BlockDevice>,
     block_size: usize,
-    _pd: PhantomData<A>,
 }
 
-pub enum DeviceMapper<
-    A: PhysicalMemoryAllocator<Address = PhysicalAddress> = GlobalPhysicalAllocator,
-> {
-    Uncached(UncachedCache<A>),
-    Cached(BlockCache<A>),
+pub enum DeviceMapper {
+    Uncached(UncachedCache),
+    Cached(BlockCache),
 }
 
-pub struct BlockCache<
-    A: PhysicalMemoryAllocator<Address = PhysicalAddress> = GlobalPhysicalAllocator,
-> {
+pub struct BlockCache {
     device: Arc<dyn BlockDevice>,
     block_size: usize,
     segment_size: usize,
-    cache: AsyncMutex<LruCache<u64, Arc<IrqSafeRwLock<CachedSegment<A>>>>>,
+    cache: AsyncMutex<LruCache<u64, Arc<IrqSafeRwLock<CachedSegment>>>>,
 }
 
 impl DeviceMapper {
@@ -54,7 +43,7 @@ impl DeviceMapper {
         bucket_count: usize,
         filesystem: &str,
     ) -> Result<DeviceMapper, Error> {
-        DeviceMapper::cached_with_capacity_in(
+        BlockCache::with_capacity(
             device,
             block_size,
             segment_size,
@@ -62,6 +51,7 @@ impl DeviceMapper {
             bucket_count,
             filesystem,
         )
+        .map(DeviceMapper::Cached)
     }
 
     pub fn uncached(
@@ -69,35 +59,6 @@ impl DeviceMapper {
         block_size: usize,
         filesystem: &str,
     ) -> Result<DeviceMapper, Error> {
-        DeviceMapper::uncached_in(device, block_size, filesystem)
-    }
-}
-
-impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> DeviceMapper<A> {
-    pub fn cached_with_capacity_in(
-        device: Arc<dyn BlockDevice>,
-        block_size: usize,
-        segment_size: usize,
-        bucket_capacity: usize,
-        bucket_count: usize,
-        filesystem: &str,
-    ) -> Result<DeviceMapper<A>, Error> {
-        BlockCache::<A>::with_capacity_in(
-            device,
-            block_size,
-            segment_size,
-            bucket_capacity,
-            bucket_count,
-            filesystem,
-        )
-        .map(DeviceMapper::<A>::Cached)
-    }
-
-    pub fn uncached_in(
-        device: Arc<dyn BlockDevice>,
-        block_size: usize,
-        filesystem: &str,
-    ) -> Result<DeviceMapper<A>, Error> {
         if block_size % device.block_size() != 0 {
             log::error!(
                 "Couldn't create block mapper for {filesystem}: \
@@ -105,11 +66,7 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> DeviceMapper<A> {
             );
             return Err(Error::InvalidArgument);
         }
-        let uncache = UncachedCache::<A> {
-            device,
-            block_size,
-            _pd: PhantomData,
-        };
+        let uncache = UncachedCache { device, block_size };
         Ok(Self::Uncached(uncache))
     }
 
@@ -154,7 +111,7 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> DeviceMapper<A> {
     }
 }
 
-impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> UncachedCache<A> {
+impl UncachedCache {
     pub fn device(&self) -> &Arc<dyn BlockDevice> {
         &self.device
     }
@@ -171,9 +128,11 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> UncachedCache<A> {
             );
             return Err(Error::InvalidArgument);
         }
-        let mut data = PageBox::<_, A>::new_uninit_slice_in(self.block_size)?;
-        self.device.read_aligned(pos, data.as_slice_mut()).await?;
-        let result = mapper(unsafe { data.assume_init_slice_ref() })?;
+        let mut buffer = self.device.allocate_buffer(self.block_size)?;
+        self.device
+            .read_aligned(pos, buffer.slice_mut(0..self.block_size))
+            .await?;
+        let result = mapper(unsafe { MaybeUninit::slice_assume_init_ref(&buffer[..]) })?;
         Ok(result)
     }
 
@@ -190,27 +149,31 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> UncachedCache<A> {
             );
             return Err(Error::InvalidArgument);
         }
-        let mut data = PageBox::<_, A>::new_uninit_slice_in(self.block_size)?;
+        let mut buffer = self.device.allocate_buffer(self.block_size)?;
         // No need to read a block only to then fully rewrite it
         if size != self.block_size {
-            self.device.read_aligned(pos, data.as_slice_mut()).await?;
+            self.device
+                .read_aligned(pos, buffer.slice_mut(0..self.block_size))
+                .await?;
         }
-        let mut data = unsafe { data.assume_init_slice() };
-        let result = mapper(&mut data[..])?;
-        self.device.write_aligned(pos, data.as_slice()).await?;
+        let mut buffer = unsafe { DmaBuffer::assume_init_slice(buffer) };
+        let result = mapper(&mut buffer[..])?;
+        self.device
+            .write_aligned(pos, buffer.slice(0..self.block_size))
+            .await?;
         Ok(result)
     }
 }
 
-impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> BlockCache<A> {
-    pub fn with_capacity_in(
+impl BlockCache {
+    pub fn with_capacity(
         device: Arc<dyn BlockDevice>,
         block_size: usize,
         segment_size: usize,
         bucket_capacity: usize,
         bucket_count: usize,
         filesystem: &str,
-    ) -> Result<BlockCache<A>, Error> {
+    ) -> Result<Self, Error> {
         if block_size % device.block_size() != 0 {
             log::error!(
                 "Couldn't create block cache for {filesystem}: \
@@ -251,17 +214,13 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> BlockCache<A> {
         &self.device
     }
 
-    async fn evict_block(
-        &self,
-        segment_position: u64,
-        block: Arc<IrqSafeRwLock<CachedSegment<A>>>,
-    ) {
+    async fn evict_block(&self, segment_position: u64, block: Arc<IrqSafeRwLock<CachedSegment>>) {
         let read = block.read();
         if read.dirty {
             assert_eq!(segment_position % self.segment_size as u64, 0);
             if let Err(err) = self
                 .device
-                .write_aligned(segment_position, read.data.as_slice())
+                .write_aligned(segment_position, read.data.slice(0..self.segment_size))
                 .await
             {
                 log::error!("Disk error: flushing block {}: {:?}", segment_position, err);
@@ -272,12 +231,12 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> BlockCache<A> {
     async fn fetch_block(
         &self,
         segment_position: u64,
-    ) -> Result<Arc<IrqSafeRwLock<CachedSegment<A>>>, Error> {
-        let mut data = PageBox::new_uninit_slice_in(self.segment_size)?;
+    ) -> Result<Arc<IrqSafeRwLock<CachedSegment>>, Error> {
+        let mut buffer = self.device.allocate_buffer(self.segment_size)?;
         self.device
-            .read_aligned(segment_position, data.as_slice_mut())
+            .read_aligned(segment_position, buffer.slice_mut(0..self.segment_size))
             .await?;
-        let data = unsafe { data.assume_init_slice() };
+        let data = unsafe { DmaBuffer::assume_init_slice(buffer) };
         Ok(Arc::new(IrqSafeRwLock::new(CachedSegment {
             data,
             dirty: false,
@@ -287,7 +246,7 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> BlockCache<A> {
     async fn entry(
         &self,
         segment_position: u64,
-    ) -> Result<Arc<IrqSafeRwLock<CachedSegment<A>>>, Error> {
+    ) -> Result<Arc<IrqSafeRwLock<CachedSegment>>, Error> {
         assert_eq!(segment_position % self.segment_size as u64, 0);
         let mut lock = self.cache.lock().await;
         let (value, evicted) = lock
@@ -349,354 +308,23 @@ impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> BlockCache<A> {
     }
 }
 
-impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> CachedSegment<A> {
+impl CachedSegment {
     pub fn set_dirty(&mut self) {
         self.dirty = true;
     }
 }
 
-impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> Deref for CachedSegment<A> {
-    type Target = PageBox<[u8], A>;
+impl Deref for CachedSegment {
+    type Target = DmaBuffer<[u8]>;
 
     fn deref(&self) -> &Self::Target {
         &self.data
     }
 }
 
-impl<A: PhysicalMemoryAllocator<Address = PhysicalAddress>> DerefMut for CachedSegment<A> {
+impl DerefMut for CachedSegment {
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.dirty = true;
         &mut self.data
     }
 }
-
-// #[cfg(test)]
-// mod tests {
-//     use core::{
-//         ffi::c_void,
-//         mem::MaybeUninit,
-//         ptr::null_mut,
-//         sync::atomic::{AtomicBool, Ordering},
-//     };
-//     use std::{io, sync::Mutex};
-//
-//     use async_trait::async_trait;
-//     use kernel_arch::mem::PhysicalMemoryAllocator;
-//     use libk_mm::{address::PhysicalAddress, PageBox, PageSlice};
-//     use yggdrasil_abi::error::Error;
-//
-//     use crate::vfs::block::{BlockDevice, NgBlockDevice, NgBlockDeviceWrapper};
-//
-//     use super::BlockCache;
-//
-//     struct DummyBlock {
-//         block_size: usize,
-//         block_count: usize,
-//         deny_writes: AtomicBool,
-//         data: Mutex<Vec<u8>>,
-//     }
-//
-//     struct PA;
-//
-//     impl DummyBlock {
-//         pub fn new(block_size: usize, block_count: usize) -> Self {
-//             let mut data = vec![0; block_size * block_count];
-//             for i in 0..block_count {
-//                 let block = &mut data[block_size * i..block_size * (i + 1)];
-//                 block.fill(i as u8);
-//             }
-//             Self {
-//                 data: Mutex::new(data),
-//                 deny_writes: AtomicBool::new(false),
-//                 block_size,
-//                 block_count,
-//             }
-//         }
-//     }
-//
-//     #[async_trait::async_trait]
-//     impl NgBlockDevice for DummyBlock {
-//         type Error = Error;
-//
-//         async fn read(
-//             &self,
-//             lba: u64,
-//             buffer: &mut PageSlice<MaybeUninit<u8>>,
-//         ) -> Result<(), Error> {
-//             let start = lba as usize * self.block_size;
-//             let end = start + buffer.len();
-//
-//             if end > self.block_count * self.block_size {
-//                 return Err(Error::InvalidArgument);
-//             }
-//
-//             let data = self.data.lock().unwrap();
-//             let buffer = unsafe { MaybeUninit::slice_assume_init_mut(&mut buffer[..]) };
-//             buffer.copy_from_slice(&data[start..end]);
-//
-//             Ok(())
-//         }
-//
-//         async fn write(&self, lba: u64, buffer: &PageSlice<u8>) -> Result<(), Error> {
-//             if self.deny_writes.load(Ordering::Acquire) {
-//                 panic!("write() with deny_writes = true");
-//             }
-//
-//             let start = lba as usize * self.block_size;
-//             let end = start + buffer.len();
-//
-//             if end > self.block_count * self.block_size {
-//                 return Err(Error::InvalidArgument);
-//             }
-//
-//             let mut data = self.data.lock().unwrap();
-//             data[start..end].copy_from_slice(&buffer[..]);
-//
-//             Ok(())
-//         }
-//
-//         fn block_size(&self) -> usize {
-//             self.block_size
-//         }
-//
-//         fn block_count(&self) -> usize {
-//             self.block_count
-//         }
-//     }
-//
-//     impl PhysicalMemoryAllocator for PA {
-//         type Address = PhysicalAddress;
-//
-//         unsafe fn free_page(page: Self::Address) {
-//             let base = page.try_into_usize().unwrap();
-//             let base = core::ptr::with_exposed_provenance_mut::<c_void>(base);
-//             if unsafe { libc::munmap(base, 0x1000) } != 0 {
-//                 let err = io::Error::last_os_error();
-//                 panic!("free_page: munmap returned {err}");
-//             }
-//         }
-//
-//         fn allocate_page() -> Result<Self::Address, Error> {
-//             Self::allocate_contiguous_pages(1)
-//         }
-//
-//         fn allocate_contiguous_pages(count: usize) -> Result<Self::Address, Error> {
-//             let base = unsafe {
-//                 libc::mmap(
-//                     null_mut(),
-//                     count * 0x1000,
-//                     libc::PROT_READ | libc::PROT_WRITE,
-//                     libc::MAP_ANON | libc::MAP_PRIVATE,
-//                     -1,
-//                     0,
-//                 )
-//             };
-//             if base != libc::MAP_FAILED {
-//                 let base = base.addr();
-//                 Ok(PhysicalAddress::from_usize(base))
-//             } else {
-//                 Err(Error::OutOfMemory)
-//             }
-//         }
-//     }
-//
-//     const BS: usize = 1024;
-//
-//     // The test must not crash with denied writes
-//     #[tokio::test]
-//     async fn test_no_modification() {
-//         let device = Box::leak(Box::new(DummyBlock::new(BS, 1024)));
-//         let wrapper = NgBlockDeviceWrapper::new(device);
-//         let cache = BlockCache::<PA>::with_capacity_in(wrapper, BS, BS, 64, 8);
-//
-//         device.deny_writes.store(true, Ordering::Release);
-//         cache
-//             .try_with(1 * BS as u64, |block| {
-//                 assert!(block.iter().all(|x| *x == 1));
-//                 Ok(())
-//             })
-//             .await
-//             .unwrap();
-//         cache
-//             .try_with(2 * BS as u64, |block| {
-//                 assert!(block.iter().all(|x| *x == 2));
-//                 Ok(())
-//             })
-//             .await
-//             .unwrap();
-//
-//         cache.flush().await;
-//     }
-//
-//     #[tokio::test]
-//     async fn test_partial_modification() {
-//         let device = Box::leak(Box::new(DummyBlock::new(BS, 1024)));
-//         let wrapper = NgBlockDeviceWrapper::new(device);
-//         // 8 * 8
-//         let cache = BlockCache::<PA>::with_capacity_in(wrapper, BS, BS, 8, 8);
-//
-//         const LBA: u64 = 1;
-//         cache
-//             .try_with_mut(LBA * BS as u64, 16, |block| {
-//                 block[0..16].fill(0x12);
-//                 Ok(())
-//             })
-//             .await
-//             .unwrap();
-//         cache.flush().await;
-//
-//         {
-//             let mut buffer = PageBox::<_, PA>::new_uninit_slice_in(BS).unwrap();
-//             device.read(LBA, buffer.as_slice_mut()).await.unwrap();
-//             let buffer = unsafe { buffer.assume_init_slice() };
-//             buffer[0..16].iter().for_each(|&x| {
-//                 assert_eq!(x, 0x12);
-//             });
-//             buffer[16..].iter().for_each(|&x| {
-//                 assert_eq!(x, LBA as u8);
-//             });
-//         }
-//
-//         cache
-//             .try_with_mut(LBA * BS as u64, 16, |block| {
-//                 block[16..32].fill(0x23);
-//                 Ok(())
-//             })
-//             .await
-//             .unwrap();
-//         cache
-//             .try_with_mut(LBA * BS as u64, 16, |block| {
-//                 block[48..64].fill(0x34);
-//                 Ok(())
-//             })
-//             .await
-//             .unwrap();
-//         cache
-//             .try_with_mut(LBA * BS as u64, 128, |block| {
-//                 block[128..256].fill(0xF1);
-//                 Ok(())
-//             })
-//             .await
-//             .unwrap();
-//         cache.flush().await;
-//
-//         {
-//             let mut buffer = PageBox::<_, PA>::new_uninit_slice_in(BS).unwrap();
-//             device.read(LBA, buffer.as_slice_mut()).await.unwrap();
-//             let buffer = unsafe { buffer.assume_init_slice() };
-//             buffer[0..16].iter().for_each(|&x| {
-//                 assert_eq!(x, 0x12);
-//             });
-//             buffer[16..32].iter().for_each(|&x| {
-//                 assert_eq!(x, 0x23);
-//             });
-//             buffer[48..64].iter().for_each(|&x| {
-//                 assert_eq!(x, 0x34);
-//             });
-//             buffer[128..256].iter().for_each(|&x| {
-//                 assert_eq!(x, 0xF1);
-//             });
-//             buffer[32..48].iter().for_each(|&x| {
-//                 assert_eq!(x, LBA as u8);
-//             });
-//             buffer[64..128].iter().for_each(|&x| {
-//                 assert_eq!(x, LBA as u8);
-//             });
-//             buffer[256..].iter().for_each(|&x| {
-//                 assert_eq!(x, LBA as u8);
-//             });
-//         }
-//     }
-//
-//     #[tokio::test]
-//     async fn test_implicit_eviction() {
-//         let device = Box::leak(Box::new(DummyBlock::new(BS, 1024)));
-//         let wrapper = NgBlockDeviceWrapper::new(device);
-//         // 8 * 8
-//         let cache = BlockCache::<PA>::with_capacity_in(wrapper, BS, BS, 8, 8);
-//
-//         fn mapper(x: u64) -> u8 {
-//             (x + 3) as u8
-//         }
-//
-//         // Go through all blocks, fill those with some values
-//         for i in 0..1024 {
-//             cache
-//                 .try_with_mut(i * BS as u64, BS, |block| {
-//                     block.fill(mapper(i));
-//                     Ok(())
-//                 })
-//                 .await
-//                 .unwrap();
-//         }
-//         cache.flush().await;
-//
-//         for i in 0..1024 {
-//             let mut buffer = PageBox::<_, PA>::new_uninit_slice_in(BS).unwrap();
-//             device.read(i, buffer.as_slice_mut()).await.unwrap();
-//             let buffer = unsafe { buffer.assume_init_slice() };
-//             assert!(buffer.iter().all(|x| *x == mapper(i)));
-//         }
-//
-//         for i in 0..1023 {
-//             cache
-//                 .try_with_mut(i * BS as u64, BS, |block| {
-//                     block.fill(0x12);
-//                     Ok(())
-//                 })
-//                 .await
-//                 .unwrap();
-//             cache
-//                 .try_with_mut((i + 1) * BS as u64, BS, |block| {
-//                     block.fill(0x23);
-//                     Ok(())
-//                 })
-//                 .await
-//                 .unwrap();
-//         }
-//
-//         for i in 0..1023 {
-//             cache
-//                 .try_with_mut(i * BS as u64, BS, |block| {
-//                     block.iter_mut().for_each(|x| *x += 1);
-//                     Ok(())
-//                 })
-//                 .await
-//                 .unwrap();
-//             cache
-//                 .try_with_mut((i + 1) * BS as u64, BS, |block| {
-//                     block.iter_mut().for_each(|x| *x += 2);
-//                     Ok(())
-//                 })
-//                 .await
-//                 .unwrap();
-//         }
-//
-//         cache.flush().await;
-//
-//         {
-//             let mut buffer = PageBox::<_, PA>::new_uninit_slice_in(BS).unwrap();
-//             device.read(0, buffer.as_slice_mut()).await.unwrap();
-//             let buffer = unsafe { buffer.assume_init_slice() };
-//             buffer.iter().for_each(|&x| {
-//                 assert_eq!(x, 0x13, "block 0 mismatch");
-//             });
-//         }
-//         for i in 1..1023 {
-//             let mut buffer = PageBox::<_, PA>::new_uninit_slice_in(BS).unwrap();
-//             device.read(i, buffer.as_slice_mut()).await.unwrap();
-//             let buffer = unsafe { buffer.assume_init_slice() };
-//             buffer.iter().for_each(|&x| {
-//                 assert_eq!(x, 0x15, "block {i} mismatch");
-//             });
-//         }
-//         {
-//             let mut buffer = PageBox::<_, PA>::new_uninit_slice_in(BS).unwrap();
-//             device.read(1023, buffer.as_slice_mut()).await.unwrap();
-//             let buffer = unsafe { buffer.assume_init_slice() };
-//             buffer.iter().for_each(|&x| {
-//                 assert_eq!(x, 0x25, "block 1023 mismatch");
-//             });
-//         }
-//     }
-// }

kernel/libk/src/device/block/mod.rs

@@ -3,10 +3,13 @@ use core::{any::Any, mem::MaybeUninit, ops::Deref};
 use alloc::{boxed::Box, sync::Arc};
 use async_trait::async_trait;
 use device_api::device::Device;
-use libk_mm::{PageBox, PageProvider, PageSlice};
+use libk_mm::PageProvider;
 use yggdrasil_abi::error::Error;
 
-use crate::vfs::{CommonImpl, NodeRef};
+use crate::{
+    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
+    vfs::{CommonImpl, NodeRef},
+};
 
 pub mod cache;
 pub mod partition;
@@ -90,15 +93,17 @@ impl Iterator for Chunked {
 
 #[async_trait]
 pub trait BlockDevice: Device + PageProvider {
+    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error>;
+
     async fn read_aligned(
         &self,
         position: u64,
-        buffer: &mut PageSlice<MaybeUninit<u8>>,
+        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
     ) -> Result<(), Error> {
         let _ = (position, buffer);
         Err(Error::NotImplemented)
     }
-    async fn write_aligned(&self, position: u64, buffer: &PageSlice<u8>) -> Result<(), Error> {
+    async fn write_aligned(&self, position: u64, buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
         let _ = (position, buffer);
         Err(Error::NotImplemented)
     }
@@ -108,13 +113,13 @@ pub trait BlockDevice: Device + PageProvider {
         let Some((iter, max_lba_count)) = Chunked::begin(self, position, buffer.len()) else {
             return Ok(0);
         };
-        let mut read_buffer = PageBox::new_uninit_slice(max_lba_count * bs)?;
+        let mut read_buffer = self.allocate_buffer(max_lba_count * bs)?;
+        // let mut read_buffer = PageBox::new_uninit_slice(max_lba_count * bs)?;
         let mut total = 0;
 
         for (lba, block_count, block_offset, offset, amount) in iter {
-            let read_buffer_slice = read_buffer.as_slice_mut().subslice_mut(..block_count * bs);
-            self.read_aligned(lba * bs as u64, read_buffer_slice)
-                .await?;
+            let slice = read_buffer.slice_mut(0..block_count * bs);
+            self.read_aligned(lba * bs as u64, slice).await?;
 
             let src = unsafe {
                 MaybeUninit::slice_assume_init_ref(
@@ -143,19 +148,19 @@ pub trait BlockDevice: Device + PageProvider {
 
             let amount = core::cmp::min(bs - block_offset, buf.len());
 
-            let mut block = PageBox::new_uninit_slice(bs)?;
+            let mut dma_buffer = self.allocate_buffer(bs)?;
             if amount != bs {
                 // Need to read the block first -- it's modified partially
-                self.read_aligned(lba * bs as u64, block.as_slice_mut())
+                self.read_aligned(lba * bs as u64, dma_buffer.slice_mut(0..bs))
                     .await?;
             }
 
-            let mut block = unsafe { block.assume_init_slice() };
+            let mut block = unsafe { DmaBuffer::assume_init_slice(dma_buffer) };
 
             block[block_offset..block_offset + amount].copy_from_slice(&buf[..amount]);
 
             // Write the block back
-            self.write_aligned(lba * bs as u64, block.as_slice())
+            self.write_aligned(lba * bs as u64, block.slice(0..bs))
                 .await?;
 
             buf = &buf[amount..];

kernel/libk/src/device/block/partition/gpt.rs

@@ -1,7 +1,6 @@
 use alloc::{sync::Arc, vec::Vec};
 use bytemuck::{Pod, Zeroable};
 use libk_mm::PageBox;
-use static_assertions::const_assert_eq;
 use uuid::Uuid;
 use yggdrasil_abi::error::Error;
 
@@ -9,7 +8,7 @@ use crate::device::block::BlockDevice;
 
 use super::Partition;
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Pod, Zeroable)]
 #[repr(C)]
 struct GptHeader {
     signature: [u8; 8],
@@ -26,7 +25,7 @@ struct GptHeader {
     partition_table_len: u32,
     partition_table_entry_size: u32,
     partition_table_crc32: u32,
-    _1: [u8; 420],
+    _1: u32,
 }
 
 #[derive(Clone, Copy, Zeroable, Pod)]
@@ -39,18 +38,18 @@ struct GptEntry {
     attrs: u64,
 }
 
-const_assert_eq!(size_of::<GptHeader>(), 512);
-
-async unsafe fn read_struct_lba<T>(dev: &dyn BlockDevice, lba: u64) -> Result<PageBox<T>, Error> {
-    assert_eq!(size_of::<T>(), 512);
-    let mut data = PageBox::new_uninit()?;
-    dev.read_aligned(lba * 512, PageBox::as_bytes_mut(&mut data))
-        .await?;
-    Ok(data.assume_init())
+async unsafe fn read_struct_lba<'b, T: Pod + 'b>(
+    dev: &dyn BlockDevice,
+    lba: u64,
+    buffer: &'b mut [u8],
+) -> Result<&'b mut T, Error> {
+    dev.read_exact(lba * 512, buffer).await?;
+    Ok(bytemuck::from_bytes_mut(buffer))
 }
 
 pub(crate) async fn probe_gpt(dev: &Arc<dyn BlockDevice>) -> Result<Option<Vec<Partition>>, Error> {
-    let header = unsafe { read_struct_lba::<GptHeader>(dev.as_ref(), 1) }.await?;
+    let mut header = [0; size_of::<GptHeader>()];
+    let header = unsafe { read_struct_lba::<GptHeader>(dev.as_ref(), 1, &mut header) }.await?;
 
     if &header.signature != b"EFI PART" {
         // Not a GPT partition table

kernel/libk/src/device/block/partition/mod.rs

@@ -3,9 +3,11 @@ use core::mem::MaybeUninit;
 use alloc::{boxed::Box, sync::Arc};
 use async_trait::async_trait;
 use device_api::device::Device;
-use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageProvider, PageSlice};
+use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageProvider};
 use yggdrasil_abi::error::Error;
 
+use crate::dma::{DmaBuffer, DmaSlice, DmaSliceMut};
+
 use super::BlockDevice;
 
 pub mod gpt;
@@ -49,18 +51,6 @@ impl PageProvider for Partition {
 }
 
 impl Device for Partition {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
-        todo!()
-    }
-
-    unsafe fn deinit(&self) -> Result<(), Error> {
-        todo!()
-    }
-
-    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
-        todo!()
-    }
-
     fn display_name(&self) -> &str {
         "Partition"
     }
@@ -68,28 +58,24 @@ impl Device for Partition {
 
 #[async_trait]
 impl BlockDevice for Partition {
+    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        self.device.allocate_buffer(size)
+    }
+
     async fn read_aligned(
         &self,
         position: u64,
-        buffer: &mut PageSlice<MaybeUninit<u8>>,
+        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
     ) -> Result<(), Error> {
-        // TODO check against partition range
-        debug_assert_eq!(position % self.device.block_size() as u64, 0);
-        debug_assert_eq!(buffer.len() % self.device.block_size(), 0);
-
-        self.device
-            .read_aligned(self.lba_start * self.block_size() as u64 + position, buffer)
-            .await
+        // TODO check against partition bounds
+        let lba = self.lba_start * self.block_size() as u64 + position;
+        self.device.read_aligned(lba, buffer).await
     }
 
-    async fn write_aligned(&self, position: u64, buffer: &PageSlice<u8>) -> Result<(), Error> {
-        // TODO check against partition range
-        debug_assert_eq!(position % self.device.block_size() as u64, 0);
-        debug_assert_eq!(buffer.len() % self.device.block_size(), 0);
-
-        self.device
-            .write_aligned(self.lba_start * self.block_size() as u64 + position, buffer)
-            .await
+    async fn write_aligned(&self, position: u64, buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
+        // TODO check against partition bounds
+        let lba = self.lba_start * self.block_size() as u64 + position;
+        self.device.write_aligned(lba, buffer).await
     }
 
     fn block_size(&self) -> usize {

kernel/libk/src/device/display/mod.rs

@@ -20,7 +20,7 @@ pub mod font;
 
 pub use color::Color;
 
-use crate::task::thread::Thread;
+use crate::{dma::DmaBuffer, task::thread::Thread};
 
 use super::block::BlockDevice;
 
@@ -151,6 +151,10 @@ impl DisplayWrapper {
 
 #[async_trait]
 impl BlockDevice for DisplayWrapper {
+    fn allocate_buffer(&self, _size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
+        Err(Error::NotImplemented)
+    }
+
     async fn read(&self, _pos: u64, _buf: &mut [u8]) -> Result<usize, Error> {
         Err(Error::InvalidOperation)
     }

kernel/libk/src/device/manager.rs

@@ -4,7 +4,6 @@ use core::{
 };
 
 use alloc::{collections::BTreeMap, format, sync::Arc, vec::Vec};
-use device_api::device::Device;
 use libk_util::{sync::spin_rwlock::IrqSafeRwLock, OneTimeInit};
 use yggdrasil_abi::{error::Error, io::FileMode};
 
@@ -47,26 +46,16 @@ pub struct SerialTerminalRegistry {
     registry: GenericRegistry<Arc<dyn CharDevice>>,
 }
 
-struct PendingDevice {
-    device: Arc<dyn Device>,
-    irq_only: bool,
-    failed: bool,
-}
-
 pub struct DeviceRegistry {
     pub display: DisplayDeviceRegistry,
     pub terminal: TerminalRegistry,
     pub serial_terminal: SerialTerminalRegistry,
-
-    pending_initialization: IrqSafeRwLock<Vec<PendingDevice>>,
 }
 
 pub static DEVICE_REGISTRY: DeviceRegistry = DeviceRegistry {
     display: DisplayDeviceRegistry::new(),
     terminal: TerminalRegistry::new(),
     serial_terminal: SerialTerminalRegistry::new(),
-
-    pending_initialization: IrqSafeRwLock::new(Vec::new()),
 };
 
 impl TerminalRegistry {
@@ -154,44 +143,44 @@ impl Deref for DisplayWrapper {
 }
 
 impl DeviceRegistry {
-    pub fn add_pending_initialization(&self, device: Arc<dyn Device>, irq_only: bool) {
-        self.pending_initialization.write().push(PendingDevice {
-            device,
-            irq_only,
-            failed: false,
-        });
-    }
+    // pub fn add_pending_initialization(&self, device: Arc<dyn Device>, irq_only: bool) {
+    //     self.pending_initialization.write().push(PendingDevice {
+    //         device,
+    //         irq_only,
+    //         failed: false,
+    //     });
+    // }
 
-    pub fn run_initialization(&self) {
-        let mut devices = self.pending_initialization.write();
+    // pub fn run_initialization(&self) {
+    //     let mut devices = self.pending_initialization.write();
 
-        for pending in devices.iter_mut() {
-            if pending.irq_only {
-                continue;
-            }
+    //     for pending in devices.iter_mut() {
+    //         if pending.irq_only {
+    //             continue;
+    //         }
 
-            log::debug!("Init device: {:?}", pending.device.display_name());
-            if let Err(error) = unsafe { pending.device.clone().init() } {
-                log::error!("{:?} init error: {error:?}", pending.device.display_name());
-                pending.failed = true;
-                continue;
-            }
-        }
+    //         log::debug!("Init device: {:?}", pending.device.display_name());
+    //         if let Err(error) = unsafe { pending.device.clone().init() } {
+    //             log::error!("{:?} init error: {error:?}", pending.device.display_name());
+    //             pending.failed = true;
+    //             continue;
+    //         }
+    //     }
 
-        for pending in devices.drain(..) {
-            if pending.failed {
-                continue;
-            }
+    //     for pending in devices.drain(..) {
+    //         if pending.failed {
+    //             continue;
+    //         }
 
-            log::debug!("Init IRQ: {:?}", pending.device.display_name());
-            if let Err(error) = unsafe { pending.device.clone().init_irq() } {
-                log::error!(
-                    "{:?} IRQ init error: {error:?}",
-                    pending.device.display_name()
-                );
-            }
-        }
-    }
+    //         log::debug!("Init IRQ: {:?}", pending.device.display_name());
+    //         if let Err(error) = unsafe { pending.device.clone().init_irq() } {
+    //             log::error!(
+    //                 "{:?} IRQ init error: {error:?}",
+    //                 pending.device.display_name()
+    //             );
+    //         }
+    //     }
+    // }
 }
 
 async fn probe_partition_table(

kernel/libk/src/dma.rs

@@ -0,0 +1,356 @@
+use core::{
+    alloc::Layout,
+    fmt,
+    mem::{self, MaybeUninit},
+    ops::{Deref, DerefMut, Range, Sub},
+    ptr::{self, NonNull},
+};
+
+use bytemuck::{Pod, Zeroable};
+use device_api::dma::{DmaAllocation, DmaAllocator};
+use libk_mm::{
+    address::{AsPhysicalAddress, PhysicalAddress, Virtualize},
+    phys, L3_PAGE_SIZE,
+};
+use yggdrasil_abi::error::Error;
+
+pub struct DummyDmaAllocator;
+
+impl DmaAllocator for DummyDmaAllocator {
+    fn allocate(&self, layout: Layout) -> Result<DmaAllocation, Error> {
+        if layout.align() > L3_PAGE_SIZE {
+            return Err(Error::InvalidMemoryOperation);
+        }
+        let page_count = layout.size().div_ceil(L3_PAGE_SIZE);
+        let host_physical = phys::alloc_pages_contiguous(page_count)?;
+        let host_virtual = unsafe {
+            NonNull::new_unchecked(ptr::with_exposed_provenance_mut(host_physical.virtualize()))
+        };
+        let bus_address = host_physical.into_u64();
+
+        Ok(DmaAllocation {
+            host_physical: host_physical.into_u64(),
+            host_virtual,
+            bus_address,
+            page_count,
+        })
+    }
+}
+
+pub struct DmaBuffer<T: ?Sized> {
+    host_pointer: NonNull<T>,
+    host_physical: PhysicalAddress,
+    bus_address: u64,
+    page_count: usize,
+}
+
+pub struct DmaSlice<'a, T> {
+    buffer: &'a DmaBuffer<[T]>,
+    range: Range<usize>,
+}
+
+pub struct DmaSliceMut<'a, T> {
+    buffer: &'a mut DmaBuffer<[T]>,
+    range: Range<usize>,
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, Debug, PartialOrd, Ord, Pod, Zeroable)]
+#[repr(transparent)]
+pub struct BusAddress(u64);
+
+impl<T> DmaBuffer<T> {
+    pub fn new(allocator: &dyn DmaAllocator, value: T) -> Result<DmaBuffer<T>, Error> {
+        let mut uninit = DmaBuffer::new_uninit(allocator)?;
+        uninit.write(value);
+        Ok(unsafe { DmaBuffer::assume_init(uninit) })
+    }
+
+    pub fn new_slice(
+        allocator: &dyn DmaAllocator,
+        value: T,
+        size: usize,
+    ) -> Result<DmaBuffer<[T]>, Error>
+    where
+        T: Copy,
+    {
+        let mut uninit = DmaBuffer::new_uninit_slice(allocator, size)?;
+        for i in 0..size {
+            uninit[i].write(value);
+        }
+        Ok(unsafe { DmaBuffer::assume_init_slice(uninit) })
+    }
+
+    pub fn from_slice(allocator: &dyn DmaAllocator, source: &[T]) -> Result<DmaBuffer<[T]>, Error>
+    where
+        T: Copy,
+    {
+        let mut uninit = DmaBuffer::new_uninit_slice(allocator, source.len())?;
+        MaybeUninit::copy_from_slice(&mut uninit[..], source);
+        Ok(unsafe { DmaBuffer::assume_init_slice(uninit) })
+    }
+
+    pub fn new_slice_with<F: Fn(usize) -> T>(
+        allocator: &dyn DmaAllocator,
+        init: F,
+        size: usize,
+    ) -> Result<DmaBuffer<[T]>, Error> {
+        let mut uninit = DmaBuffer::new_uninit_slice(allocator, size)?;
+        for i in 0..size {
+            uninit[i].write(init(i));
+        }
+        Ok(unsafe { DmaBuffer::assume_init_slice(uninit) })
+    }
+
+    pub fn new_uninit(allocator: &dyn DmaAllocator) -> Result<DmaBuffer<MaybeUninit<T>>, Error> {
+        let layout = Layout::new::<T>();
+        let allocation = allocator.allocate(layout)?;
+        let host_pointer = allocation.host_virtual.cast();
+        Ok(DmaBuffer {
+            host_pointer,
+            host_physical: PhysicalAddress::from_u64(allocation.host_physical),
+            bus_address: allocation.bus_address,
+            page_count: allocation.page_count,
+        })
+    }
+
+    pub fn new_uninit_slice(
+        allocator: &dyn DmaAllocator,
+        size: usize,
+    ) -> Result<DmaBuffer<[MaybeUninit<T>]>, Error> {
+        let layout = Layout::array::<T>(size).map_err(|_| Error::InvalidMemoryOperation)?;
+        let allocation = allocator.allocate(layout)?;
+        let host_pointer = NonNull::slice_from_raw_parts(allocation.host_virtual.cast(), size);
+        Ok(DmaBuffer {
+            host_pointer,
+            host_physical: PhysicalAddress::from_u64(allocation.host_physical),
+            bus_address: allocation.bus_address,
+            page_count: allocation.page_count,
+        })
+    }
+
+    pub fn new_zeroed_slice(
+        allocator: &dyn DmaAllocator,
+        size: usize,
+    ) -> Result<DmaBuffer<[MaybeUninit<T>]>, Error> {
+        let layout = Layout::array::<T>(size).map_err(|_| Error::InvalidMemoryOperation)?;
+        let allocation = allocator.allocate(layout)?;
+        unsafe {
+            let mut slice = NonNull::<[u8]>::slice_from_raw_parts(
+                allocation.host_virtual.cast(),
+                layout.size(),
+            );
+
+            slice.as_mut().fill(0);
+        }
+        let host_pointer = NonNull::slice_from_raw_parts(allocation.host_virtual.cast(), size);
+        Ok(DmaBuffer {
+            host_pointer,
+            host_physical: PhysicalAddress::from_u64(allocation.host_physical),
+            bus_address: allocation.bus_address,
+            page_count: allocation.page_count,
+        })
+    }
+
+    pub unsafe fn assume_init(buffer: DmaBuffer<MaybeUninit<T>>) -> DmaBuffer<T> {
+        let host_pointer = buffer.host_pointer;
+        let host_physical = buffer.host_physical;
+        let bus_address = buffer.bus_address;
+        let page_count = buffer.page_count;
+
+        mem::forget(buffer);
+
+        let host_pointer = host_pointer.cast();
+
+        DmaBuffer {
+            host_pointer,
+            host_physical,
+            bus_address,
+            page_count,
+        }
+    }
+
+    pub unsafe fn assume_init_slice(buffer: DmaBuffer<[MaybeUninit<T>]>) -> DmaBuffer<[T]> {
+        let host_pointer = buffer.host_pointer;
+        let host_physical = buffer.host_physical;
+        let bus_address = buffer.bus_address;
+        let page_count = buffer.page_count;
+
+        mem::forget(buffer);
+
+        let len = host_pointer.len();
+        let host_pointer = NonNull::slice_from_raw_parts(host_pointer.cast::<T>(), len);
+
+        DmaBuffer {
+            host_pointer,
+            host_physical,
+            bus_address,
+            page_count,
+        }
+    }
+}
+
+impl<T> DmaBuffer<[T]> {
+    fn slice_range_check(&self, range: &Range<usize>) {
+        assert!(
+            range.end <= self.len() && range.start <= self.len(),
+            "DMA buffer slice range out of bounds"
+        );
+        assert!(range.start <= range.end, "Invalid DMA slice range");
+    }
+
+    pub fn slice(&self, range: Range<usize>) -> DmaSlice<T> {
+        self.slice_range_check(&range);
+        DmaSlice {
+            buffer: self,
+            range,
+        }
+    }
+
+    pub fn slice_mut(&mut self, range: Range<usize>) -> DmaSliceMut<T> {
+        self.slice_range_check(&range);
+        DmaSliceMut {
+            buffer: self,
+            range,
+        }
+    }
+}
+
+impl<T: ?Sized> DmaBuffer<T> {
+    #[inline]
+    pub fn page_count(&self) -> usize {
+        self.page_count
+    }
+
+    #[inline]
+    pub fn bus_address(&self) -> BusAddress {
+        BusAddress(self.bus_address)
+    }
+}
+
+unsafe impl<T: ?Sized + Send> Send for DmaBuffer<T> {}
+unsafe impl<T: ?Sized + Sync> Sync for DmaBuffer<T> {}
+
+impl<T: ?Sized> Drop for DmaBuffer<T> {
+    fn drop(&mut self) {
+        log::trace!("Drop DmaBuffer @ {:#x}", self.host_physical);
+        unsafe {
+            ptr::drop_in_place(self.host_pointer.as_ptr());
+            for i in 0..self.page_count {
+                phys::free_page(self.host_physical.add(i * L3_PAGE_SIZE));
+            }
+        }
+    }
+}
+
+impl<T: ?Sized> AsPhysicalAddress for DmaBuffer<T> {
+    #[inline]
+    unsafe fn as_physical_address(&self) -> PhysicalAddress {
+        self.host_physical
+    }
+}
+
+impl<T: ?Sized> Deref for DmaBuffer<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { self.host_pointer.as_ref() }
+    }
+}
+
+impl<T: ?Sized> DerefMut for DmaBuffer<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { self.host_pointer.as_mut() }
+    }
+}
+
+impl<T: ?Sized> fmt::Pointer for DmaBuffer<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "<dma@{:#x}>", self.bus_address)
+    }
+}
+
+impl BusAddress {
+    pub const ZERO: Self = Self(0);
+
+    pub const fn into_u64(self) -> u64 {
+        self.0
+    }
+
+    pub fn try_into_u32(self) -> Result<u32, Error> {
+        self.0.try_into().map_err(|_| Error::InvalidMemoryOperation)
+    }
+
+    pub const fn add(self, offset: usize) -> Self {
+        Self(self.0 + offset as u64)
+    }
+}
+
+impl Sub<BusAddress> for BusAddress {
+    type Output = usize;
+
+    fn sub(self, rhs: BusAddress) -> Self::Output {
+        (self.0 - rhs.0).try_into().unwrap()
+    }
+}
+
+impl fmt::LowerHex for BusAddress {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::LowerHex::fmt(&self.0, f)
+    }
+}
+
+impl<'a, T> DmaSlice<'a, T> {
+    pub fn bus_address(&self) -> BusAddress {
+        self.buffer.bus_address().add(self.range.start)
+    }
+
+    // TODO subslicing
+    pub fn into_parts(self) -> (&'a DmaBuffer<[T]>, Range<usize>) {
+        (self.buffer, self.range)
+    }
+}
+
+impl<T> Deref for DmaSlice<'_, T> {
+    type Target = [T];
+
+    fn deref(&self) -> &Self::Target {
+        &self.buffer[self.range.clone()]
+    }
+}
+
+impl<T> fmt::Pointer for DmaSlice<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:p}[{:?}]", *self.buffer, self.range)
+    }
+}
+
+impl<'a, T> DmaSliceMut<'a, T> {
+    pub fn bus_address(&self) -> BusAddress {
+        self.buffer.bus_address().add(self.range.start)
+    }
+
+    // TODO subslicing
+    pub fn into_parts(self) -> (&'a mut DmaBuffer<[T]>, Range<usize>) {
+        (self.buffer, self.range)
+    }
+}
+
+impl<T> Deref for DmaSliceMut<'_, T> {
+    type Target = [T];
+
+    fn deref(&self) -> &Self::Target {
+        &self.buffer[self.range.clone()]
+    }
+}
+
+impl<T> DerefMut for DmaSliceMut<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.buffer[self.range.clone()]
+    }
+}
+
+impl<T> fmt::Pointer for DmaSliceMut<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:p}[{:?}]", *self.buffer, self.range)
+    }
+}

kernel/libk/src/lib.rs

@@ -6,6 +6,7 @@
     new_range_api,
     associated_type_defaults,
     maybe_uninit_slice,
+    maybe_uninit_write_slice,
     step_trait,
     const_trait_impl,
     slice_ptr_get,
@@ -44,6 +45,8 @@ pub mod random;
 pub mod time;
 pub mod vfs;
 
+pub mod dma;
+
 #[cfg(any(target_os = "none", rust_analyzer))]
 pub mod panic;
 

kernel/src/arch/aarch64/gic/gicv2m.rs

@@ -4,7 +4,7 @@ use core::{mem::offset_of, ops::Range};
 use abi::error::Error;
 use alloc::{sync::Arc, vec::Vec};
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{
         ExternalInterruptController, InterruptAffinity, InterruptHandler, Irq, IrqLevel,
         IrqOptions, IrqTrigger, IrqVector, MessageInterruptController, MsiInfo,
@@ -83,7 +83,7 @@ pub struct Gicv2m {
 }
 
 impl Device for Gicv2m {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let regs = self.regs.lock();
         log::info!("gicv2m: init @ {:#x}", self.base);
 

kernel/src/arch/aarch64/gic/mod.rs

@@ -6,7 +6,7 @@ use aarch64_cpu::asm::barrier;
 use abi::error::Error;
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{
         ExternalInterruptController, FixedInterruptTable, FullIrq, InterruptHandler,
         InterruptTable, IpiDeliveryTarget, IpiMessage, Irq, IrqLevel, IrqOptions, IrqTrigger,
@@ -58,7 +58,7 @@ impl Device for Gic {
         "ARM Generic Interrupt Controller v2"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         log::debug!(
             "Init GIC: gicd={:#x}, gicc={:#x}",
             self.gicd_base,

kernel/src/arch/aarch64/timer.rs

@@ -6,7 +6,7 @@ use aarch64_cpu::registers::{CNTFRQ_EL0, CNTPCT_EL0, CNTP_CTL_EL0, CNTP_TVAL_EL0
 use abi::{error::Error, time::NANOSECONDS_IN_SECOND};
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{FullIrq, InterruptHandler, IrqVector},
 };
 use device_tree::driver::{device_tree_driver, Node, ProbeContext};
@@ -55,7 +55,7 @@ impl Device for ArmTimer {
         "ARM Generic Timer"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         CNTP_CTL_EL0.write(CNTP_CTL_EL0::ENABLE::SET + CNTP_CTL_EL0::IMASK::SET);
         Ok(())
     }

kernel/src/arch/x86/mod.rs

@@ -4,10 +4,14 @@
 
 use abi::error::Error;
 use alloc::{sync::Arc, vec::Vec};
-use device_api::{device::Device, interrupt::Irq};
+use device_api::{
+    device::{Device, DeviceInitContext},
+    interrupt::Irq,
+};
 use kernel_arch_x86::ISA_IRQ_OFFSET;
 use libk::{
     config, debug,
+    dma::DummyDmaAllocator,
     fs::{devfs, sysfs},
     task::runtime,
 };
@@ -50,6 +54,12 @@ pub enum SelectedClockSource {
     Fallback(Arc<I8253>),
 }
 
+pub fn dummy_init_context() -> DeviceInitContext {
+    DeviceInitContext {
+        dma_allocator: Arc::new(DummyDmaAllocator),
+    }
+}
+
 // Initialize the bare minimum required to:
 // * Allocate/manage interrupts
 // * Print debug output
@@ -63,7 +73,12 @@ pub fn init_platform_early(cmdline: &str) -> Result<EarlyPlatformDevices, Error>
     // Initialize async executor queue
     runtime::init_task_queue();
 
-    let com1_3 = ComPort::setup(0x3F8, 0x3E8, Irq::External(ISA_IRQ_OFFSET + 4))?;
+    let com1_3 = ComPort::setup(
+        dummy_init_context(),
+        0x3F8,
+        0x3E8,
+        Irq::External(ISA_IRQ_OFFSET + 4),
+    )?;
     let i8259 = I8259::setup().expect("Could not initialize i8259 PIC");
 
     #[cfg(any(target_arch = "x86", rust_analyzer))]
@@ -71,7 +86,7 @@ pub fn init_platform_early(cmdline: &str) -> Result<EarlyPlatformDevices, Error>
         use libk::device::register_external_interrupt_controller;
 
         // No other interrupt handling options
-        unsafe { i8259.clone().init() }?;
+        unsafe { i8259.clone().init(dummy_init_context()) }?;
         register_external_interrupt_controller(i8259.clone());
     }
 

kernel/src/arch/x86/peripherals/hpet.rs

@@ -4,7 +4,7 @@ use abi::error::Error;
 use acpi::HpetInfo;
 use alloc::{collections::btree_map::BTreeMap, format, string::String, sync::Arc};
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{InterruptHandler, Irq, IrqOptions, IrqTrigger, IrqVector},
 };
 use libk::{device::external_interrupt_controller, task::runtime, time};
@@ -16,6 +16,8 @@ use tock_registers::{
     registers::{ReadOnly, ReadWrite},
 };
 
+use crate::arch::x86::dummy_init_context;
+
 register_bitfields! {
     u64,
     GENERAL_CAPABILITIES [
@@ -122,7 +124,7 @@ impl InterruptHandler for HpetTimer {
 }
 
 impl Device for HpetTimer {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         if self.period > u32::MAX as u64 && !self.bits64 {
             log::error!(
                 "HPET period is >32bit and the HPET itself does not support 64bit counters"
@@ -182,7 +184,7 @@ impl Device for HpetTimer {
 }
 
 impl Device for Hpet {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let regs = self.regs.write();
 
         // Reset the device until at least one timer is created
@@ -244,9 +246,10 @@ impl Hpet {
     }
 
     pub fn setup_from_acpi(acpi: &HpetInfo) -> Result<Arc<Self>, Error> {
+        let cx = dummy_init_context();
         let base = PhysicalAddress::from_usize(acpi.base_address);
         let hpet = Arc::new(Self::new(base)?);
-        unsafe { hpet.clone().init() }?;
+        unsafe { hpet.clone().init(cx) }?;
         Ok(hpet)
     }
 
@@ -279,7 +282,8 @@ impl Hpet {
         let timer = if let Ok(timer) = timer {
             timers.insert(index, timer.clone());
 
-            match unsafe { timer.clone().init() } {
+            let cx = dummy_init_context();
+            match unsafe { timer.clone().init(cx) } {
                 Ok(()) => Ok(timer),
                 Err(error) => Err(error),
             }

kernel/src/arch/x86/peripherals/i8259.rs

@@ -3,7 +3,7 @@
 use abi::error::Error;
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{
         ExternalInterruptController, FixedInterruptTable, InterruptHandler, InterruptTable, Irq,
         IrqOptions, IrqVector,
@@ -38,7 +38,7 @@ pub struct I8259 {
 }
 
 impl Device for I8259 {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         self.enable();
         Ok(())
     }

kernel/src/arch/x86/peripherals/ps2/mod.rs

@@ -5,7 +5,7 @@ use abi::{
 };
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{InterruptHandler, Irq, IrqVector},
 };
 use kernel_arch_x86::{
@@ -113,7 +113,7 @@ impl Device for PS2Controller {
         "PS/2 Controller"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         Ok(())
     }
 

kernel/src/arch/x86/peripherals/serial.rs

@@ -2,7 +2,7 @@
 use abi::{error::Error, io::TerminalOptions};
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{InterruptHandler, Irq, IrqVector},
 };
 use kernel_arch_x86::intrinsics::{IoPort, IoPortAccess};
@@ -113,7 +113,7 @@ impl Device for Port {
         "COM port"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         DEVICE_REGISTRY
             .serial_terminal
             .register(self.terminal.clone(), Some(self.clone()))
@@ -168,9 +168,14 @@ impl ComPort {
         })
     }
 
-    pub fn setup(port_a: u16, port_b: u16, irq: Irq) -> Result<Arc<Self>, Error> {
+    pub fn setup(
+        cx: DeviceInitContext,
+        port_a: u16,
+        port_b: u16,
+        irq: Irq,
+    ) -> Result<Arc<Self>, Error> {
         let this = Arc::new(Self::new(port_a, port_b, irq)?);
-        unsafe { this.port_a().clone().init() }?;
+        unsafe { this.port_a().clone().init(cx) }?;
         Ok(this)
     }
 

kernel/src/arch/x86_64/apic/ioapic.rs

@@ -3,7 +3,7 @@ use ::acpi::platform::interrupt::{Apic as AcpiApic, Polarity, TriggerMode};
 use abi::error::Error;
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{
         ExternalInterruptController, FixedInterruptTable, InterruptHandler, InterruptTable, Irq,
         IrqLevel, IrqOptions, IrqTrigger, IrqVector,
@@ -157,7 +157,7 @@ impl Device for IoApic {
         "I/O APIC"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         unreachable!()
     }
 }

kernel/src/device/bus/simple_bus.rs

@@ -7,7 +7,10 @@ use core::ops::Range;
 
 use abi::error::Error;
 use alloc::{sync::Arc, vec::Vec};
-use device_api::{bus::Bus, device::Device};
+use device_api::{
+    bus::Bus,
+    device::{Device, DeviceInitContext},
+};
 use device_tree::{
     driver::{device_tree_driver, Node, ProbeContext},
     DeviceTreePropertyRead,
@@ -18,7 +21,7 @@ struct SimpleBus {
 }
 
 impl Device for SimpleBus {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         Ok(())
     }
 

kernel/src/device/clock/bcm2835_aux.rs

@@ -2,7 +2,10 @@
 use aarch64_cpu::registers::ReadWriteable;
 use abi::error::Error;
 use alloc::sync::Arc;
-use device_api::{clock::ClockController, device::Device};
+use device_api::{
+    clock::ClockController,
+    device::{Device, DeviceInitContext},
+};
 use device_tree::driver::{device_tree_driver, Node, ProbeContext};
 use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
 use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
@@ -52,7 +55,7 @@ impl ClockController for Bcm2835Aux {
 }
 
 impl Device for Bcm2835Aux {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let regs = DeviceMemoryIo::map(self.base, Default::default())?;
         self.regs.init(IrqSafeSpinlock::new(regs));
         Ok(())

kernel/src/device/interrupt/riscv_plic.rs

@@ -4,7 +4,7 @@ use core::sync::atomic::Ordering;
 use abi::{error::Error, primitive_enum};
 use alloc::{sync::Arc, vec::Vec};
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{
         ExternalInterruptController, FixedInterruptTable, FullIrq, InterruptHandler,
         InterruptTable, Irq, IrqOptions, IrqVector,
@@ -209,7 +209,7 @@ impl ExternalInterruptController for Plic {
 }
 
 impl Device for Plic {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         log::info!("Initialize RISC-V PLIC");
 
         let common = DeviceMemoryIo::<CommonRegs>::map(self.base, Default::default())?;

kernel/src/device/power/arm_psci.rs

@@ -2,7 +2,10 @@
 
 use abi::error::Error;
 use alloc::sync::Arc;
-use device_api::{device::Device, CpuBringupDevice, ResetDevice};
+use device_api::{
+    device::{Device, DeviceInitContext},
+    CpuBringupDevice, ResetDevice,
+};
 use device_tree::{
     driver::{device_tree_driver, Node, ProbeContext},
     DeviceTreePropertyRead,
@@ -31,7 +34,7 @@ impl Device for Psci {
         "ARM PSCI"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         PLATFORM.psci.init(self.clone());
         Ok(())
     }

kernel/src/device/serial/bcm2835_aux_uart.rs

@@ -8,7 +8,7 @@ use abi::{
 };
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{FullIrq, InterruptHandler, IrqVector},
 };
 use device_tree::{
@@ -148,7 +148,7 @@ impl InterruptHandler for Bcm2835AuxUart {
 }
 
 impl Device for Bcm2835AuxUart {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         // TODO initialize pinctrl
 
         // NOTE: might as well make it an error if clock cannot be initialized

kernel/src/device/serial/ns16550a.rs

@@ -2,7 +2,7 @@
 use abi::{error::Error, io::TerminalOptions};
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{FullIrq, InterruptHandler, IrqVector},
 };
 use device_tree::driver::{device_tree_driver, Node, ProbeContext};
@@ -108,7 +108,7 @@ impl Io {
 }
 
 impl Device for Ns16550a {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         log::info!("Init ns16550a @ {:#x}", self.base);
         let mut io = Io {
             regs: DeviceMemoryIo::map(self.base, Default::default())?,

kernel/src/device/serial/pl011.rs

@@ -2,7 +2,7 @@
 use abi::{error::Error, io::TerminalOptions};
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{FullIrq, InterruptHandler, IrqVector},
 };
 use device_tree::driver::{device_tree_driver, Node, ProbeContext};
@@ -139,7 +139,7 @@ impl Device for Pl011 {
         "Primecell PL011 UART"
     }
 
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let mut io = Io {
             regs: DeviceMemoryIo::map(self.base, Default::default())?,
         };

kernel/src/device/serial/snps_dw_apb_uart.rs

@@ -2,7 +2,7 @@
 use abi::{error::Error, io::TerminalOptions};
 use alloc::sync::Arc;
 use device_api::{
-    device::Device,
+    device::{Device, DeviceInitContext},
     interrupt::{FullIrq, InterruptHandler, IrqVector},
 };
 use device_tree::driver::{device_tree_driver, Node, ProbeContext};
@@ -148,7 +148,7 @@ impl InterruptHandler for DwUart {
 }
 
 impl Device for DwUart {
-    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
+    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
         let regs = DeviceMemoryIo::map(self.base, Default::default())?;
         let mut io = Io { regs };
         io.init();