Extract DMA remapping feature

ostd/src/arch/x86/iommu/dma_remapping/context_table.rs

@@ -6,6 +6,7 @@ use alloc::collections::BTreeMap;
 use core::mem::size_of;
 
 use log::warn;
+use ostd_pod::Pod;
 
 use super::second_stage::{DeviceMode, PageTableEntry, PagingConsts};
 use crate::{
@@ -16,7 +17,6 @@ use crate::{
         page_table::{PageTableError, PageTableItem},
         Frame, FrameAllocOptions, Paddr, PageFlags, PageTable, VmIo, PAGE_SIZE,
     },
-    Pod,
 };
 
 /// Bit 0 is `Present` bit, indicating whether this entry is present.
@@ -51,18 +51,23 @@ pub enum ContextTableError {
 }
 
 impl RootTable {
-    pub fn new() -> Self {
+    pub fn root_paddr(&self) -> Paddr {
+        self.root_frame.start_paddr()
+    }
+
+    pub(super) fn new() -> Self {
         Self {
             root_frame: FrameAllocOptions::new(1).alloc_single().unwrap(),
             context_tables: BTreeMap::new(),
         }
     }
 
+    /// Mapping device address to physical address.
     ///
     /// # Safety
     ///
     /// User must ensure the given paddr is a valid one.
-    pub unsafe fn map(
+    pub(super) unsafe fn map(
         &mut self,
         device: PciDeviceLocation,
         daddr: Daddr,
@@ -78,7 +83,7 @@ impl RootTable {
         Ok(())
     }
 
-    pub fn unmap(
+    pub(super) fn unmap(
         &mut self,
         device: PciDeviceLocation,
         daddr: Daddr,
@@ -93,37 +98,11 @@ impl RootTable {
         Ok(())
     }
 
-    pub fn paddr(&self) -> Paddr {
-        self.root_frame.start_paddr()
-    }
-
-    fn get_or_create_context_table(&mut self, device_id: PciDeviceLocation) -> &mut ContextTable {
-        let bus_entry = self
-            .root_frame
-            .read_val::<RootEntry>(device_id.bus as usize * size_of::<RootEntry>())
-            .unwrap();
-
-        if !bus_entry.is_present() {
-            let table = ContextTable::new();
-            let address = table.paddr();
-            self.context_tables.insert(address, table);
-            let entry = RootEntry(address as u128 | 1);
-            self.root_frame
-                .write_val::<RootEntry>(device_id.bus as usize * size_of::<RootEntry>(), &entry)
-                .unwrap();
-            self.context_tables.get_mut(&address).unwrap()
-        } else {
-            self.context_tables
-                .get_mut(&(bus_entry.addr() as usize))
-                .unwrap()
-        }
-    }
-
     /// Specify the device page table instead of creating a page table if not exists.
     ///
     /// This will be useful if we want all the devices to use the same page table.
     /// The original page table will be overwritten.
-    pub fn specify_device_page_table(
+    pub(super) fn specify_device_page_table(
         &mut self,
         device_id: PciDeviceLocation,
         page_table: PageTable<DeviceMode, PageTableEntry, PagingConsts>,
@@ -153,6 +132,28 @@ impl RootTable {
             .unwrap();
         context_table.page_tables.get_mut(&address).unwrap();
     }
+
+    fn get_or_create_context_table(&mut self, device_id: PciDeviceLocation) -> &mut ContextTable {
+        let bus_entry = self
+            .root_frame
+            .read_val::<RootEntry>(device_id.bus as usize * size_of::<RootEntry>())
+            .unwrap();
+
+        if !bus_entry.is_present() {
+            let table = ContextTable::new();
+            let address = table.paddr();
+            self.context_tables.insert(address, table);
+            let entry = RootEntry(address as u128 | 1);
+            self.root_frame
+                .write_val::<RootEntry>(device_id.bus as usize * size_of::<RootEntry>(), &entry)
+                .unwrap();
+            self.context_tables.get_mut(&address).unwrap()
+        } else {
+            self.context_tables
+                .get_mut(&(bus_entry.addr() as usize))
+                .unwrap()
+        }
+    }
 }
 
 /// Context Entry in the Context Table, used in Intel iommu.

ostd/src/arch/x86/iommu/dma_remapping/mod.rs

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: MPL-2.0
+
+pub use context_table::RootTable;
+use log::info;
+use second_stage::{DeviceMode, PageTableEntry, PagingConsts};
+use spin::Once;
+
+use super::IommuError;
+use crate::{
+    arch::iommu::registers::IOMMU_REGS,
+    bus::pci::PciDeviceLocation,
+    mm::{Daddr, PageTable},
+    prelude::Paddr,
+    sync::SpinLock,
+};
+
+mod context_table;
+mod second_stage;
+
+pub fn has_dma_remapping() -> bool {
+    PAGE_TABLE.get().is_some()
+}
+
+/// Mapping device address to physical address.
+///
+/// # Safety
+///
+/// Mapping an incorrect address may lead to a kernel data leak.
+pub unsafe fn map(daddr: Daddr, paddr: Paddr) -> Result<(), IommuError> {
+    let Some(table) = PAGE_TABLE.get() else {
+        return Err(IommuError::NoIommu);
+    };
+    // The page table of all devices is the same. So we can use any device ID.
+    table
+        .lock()
+        .map(PciDeviceLocation::zero(), daddr, paddr)
+        .map_err(|err| match err {
+            context_table::ContextTableError::InvalidDeviceId => unreachable!(),
+            context_table::ContextTableError::ModificationError(err) => {
+                IommuError::ModificationError(err)
+            }
+        })
+}
+
+pub fn unmap(daddr: Daddr) -> Result<(), IommuError> {
+    let Some(table) = PAGE_TABLE.get() else {
+        return Err(IommuError::NoIommu);
+    };
+    // The page table of all devices is the same. So we can use any device ID.
+    table
+        .lock()
+        .unmap(PciDeviceLocation::zero(), daddr)
+        .map_err(|err| match err {
+            context_table::ContextTableError::InvalidDeviceId => unreachable!(),
+            context_table::ContextTableError::ModificationError(err) => {
+                IommuError::ModificationError(err)
+            }
+        })
+}
+
+pub fn init() {
+    // Create Root Table instance
+    let mut root_table = RootTable::new();
+    // For all PCI Device, use the same page table.
+    let page_table = PageTable::<DeviceMode, PageTableEntry, PagingConsts>::empty();
+    for table in PciDeviceLocation::all() {
+        root_table.specify_device_page_table(table, unsafe { page_table.shallow_copy() })
+    }
+    PAGE_TABLE.call_once(|| SpinLock::new(root_table));
+
+    // Enable DMA remapping
+    let mut iommu_regs = IOMMU_REGS.get().unwrap().lock_irq_disabled();
+    iommu_regs.enable_dma_remapping(PAGE_TABLE.get().unwrap());
+    info!("[IOMMU] DMA remapping enabled");
+}
+
+static PAGE_TABLE: Once<SpinLock<RootTable>> = Once::new();

ostd/src/arch/x86/iommu/mod.rs

@@ -2,22 +2,13 @@
 
 //! The IOMMU support.
 
-mod context_table;
+mod dma_remapping;
 mod fault;
 mod registers;
-mod second_stage;
 
-use log::info;
-pub use second_stage::DeviceMode;
-use second_stage::{PageTableEntry, PagingConsts};
-use spin::Once;
+pub(crate) use dma_remapping::{has_dma_remapping, map, unmap};
 
-use crate::{
-    arch::iommu::context_table::RootTable,
-    bus::pci::PciDeviceLocation,
-    mm::{dma::Daddr, page_table::PageTableError, Paddr, PageTable},
-    sync::Mutex,
-};
+use crate::mm::page_table::PageTableError;
 
 /// An enumeration representing possible errors related to IOMMU.
 #[derive(Debug)]
@@ -28,59 +19,9 @@ pub enum IommuError {
     ModificationError(PageTableError),
 }
 
-///
-/// # Safety
-///
-/// Mapping an incorrect address may lead to a kernel data leak.
-pub(crate) unsafe fn map(daddr: Daddr, paddr: Paddr) -> Result<(), IommuError> {
-    let Some(table) = PAGE_TABLE.get() else {
-        return Err(IommuError::NoIommu);
-    };
-    // The page table of all devices is the same. So we can use any device ID.
-    table
-        .lock()
-        .map(PciDeviceLocation::zero(), daddr, paddr)
-        .map_err(|err| match err {
-            context_table::ContextTableError::InvalidDeviceId => unreachable!(),
-            context_table::ContextTableError::ModificationError(err) => {
-                IommuError::ModificationError(err)
-            }
-        })
-}
-
-pub(crate) fn unmap(daddr: Daddr) -> Result<(), IommuError> {
-    let Some(table) = PAGE_TABLE.get() else {
-        return Err(IommuError::NoIommu);
-    };
-    // The page table of all devices is the same. So we can use any device ID.
-    table
-        .lock()
-        .unmap(PciDeviceLocation::zero(), daddr)
-        .map_err(|err| match err {
-            context_table::ContextTableError::InvalidDeviceId => unreachable!(),
-            context_table::ContextTableError::ModificationError(err) => {
-                IommuError::ModificationError(err)
-            }
-        })
-}
-
 pub(crate) fn init() -> Result<(), IommuError> {
     registers::init()?;
+    dma_remapping::init();
 
-    let mut root_table = RootTable::new();
-    // For all PCI Device, use the same page table.
-    let page_table = PageTable::<DeviceMode, PageTableEntry, PagingConsts>::empty();
-    for table in PciDeviceLocation::all() {
-        root_table.specify_device_page_table(table, unsafe { page_table.shallow_copy() })
-    }
-    remapping::init(&root_table)?;
-    PAGE_TABLE.call_once(|| Mutex::new(root_table));
-    info!("IOMMU enabled");
     Ok(())
 }
-
-pub(crate) fn has_iommu() -> bool {
-    PAGE_TABLE.get().is_some()
-}
-
-static PAGE_TABLE: Once<Mutex<RootTable>> = Once::new();

ostd/src/arch/x86/iommu/registers/mod.rs

@@ -19,7 +19,7 @@ use volatile::{
     Volatile,
 };
 
-use super::{dma_remapping::context_table::RootTable, IommuError};
+use super::{dma_remapping::RootTable, IommuError};
 use crate::{
     arch::{
         iommu::fault,
@@ -89,7 +89,7 @@ impl IommuRegisters {
     pub(super) fn enable_dma_remapping(&mut self, root_table: &'static SpinLock<RootTable>) {
         // Set root table address
         self.root_table_address
-            .write(root_table.lock().paddr() as u64);
+            .write(root_table.lock().root_paddr() as u64);
         self.write_global_command(GlobalCommand::SRTP, true);
         while !self.global_status().contains(GlobalStatus::RTPS) {}
 

ostd/src/mm/dma/mod.rs

@@ -11,7 +11,7 @@ use inherit_methods_macro::inherit_methods;
 use spin::Once;
 
 use super::Paddr;
-use crate::{arch::iommu::has_iommu, mm::PAGE_SIZE, sync::SpinLock};
+use crate::{arch::iommu::has_dma_remapping, mm::PAGE_SIZE, sync::SpinLock};
 
 /// The devide address.
 ///
@@ -49,7 +49,7 @@ impl<T: HasDaddr> HasDaddr for &T {
 static DMA_MAPPING_SET: Once<SpinLock<BTreeSet<Paddr>>> = Once::new();
 
 pub fn dma_type() -> DmaType {
-    if has_iommu() {
+    if has_dma_remapping() {
         DmaType::Iommu
     } else {
         DmaType::Direct