Remove the system device's IO memory access

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

@@ -11,7 +11,7 @@ mod registers;
 pub(crate) use dma_remapping::{has_dma_remapping, map, unmap};
 pub(crate) use interrupt_remapping::{alloc_irt_entry, has_interrupt_remapping, IrtEntryHandle};
 
-use crate::mm::page_table::PageTableError;
+use crate::{io::IoMemAllocatorBuilder, mm::page_table::PageTableError};
 
 /// An enumeration representing possible errors related to IOMMU.
 #[derive(Debug)]
@@ -22,8 +22,8 @@ pub enum IommuError {
     ModificationError(PageTableError),
 }
 
-pub(crate) fn init() -> Result<(), IommuError> {
-    registers::init()?;
+pub(crate) fn init(io_mem_builder: &IoMemAllocatorBuilder) -> Result<(), IommuError> {
+    registers::init(io_mem_builder)?;
     invalidate::init();
     dma_remapping::init();
     interrupt_remapping::init();

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

@@ -39,7 +39,8 @@ use crate::{
         },
         x86::kernel::acpi::dmar::{Dmar, Remapping},
     },
-    mm::paddr_to_vaddr,
+    io::IoMemAllocatorBuilder,
+    mm::{paddr_to_vaddr, PAGE_SIZE},
     sync::{LocalIrqDisabled, SpinLock},
 };
 
@@ -251,7 +252,7 @@ impl IommuRegisters {
     }
 
     /// Creates an instance from base address
-    fn new() -> Option<Self> {
+    fn new(io_mem_builder: &IoMemAllocatorBuilder) -> Option<Self> {
         let dmar = Dmar::new()?;
         debug!("DMAR: {:#x?}", dmar);
 
@@ -265,6 +266,7 @@ impl IommuRegisters {
         assert_ne!(base_address, 0, "IOMMU address should not be zero");
         debug!("IOMMU base address: {:#x?}", base_address);
 
+        io_mem_builder.remove(base_address as usize..(base_address as usize + PAGE_SIZE));
         let base = NonNull::new(paddr_to_vaddr(base_address as usize) as *mut u8).unwrap();
 
         // SAFETY: All offsets and sizes are strictly adhered to in the manual, and the base
@@ -303,8 +305,8 @@ impl IommuRegisters {
 
 pub(super) static IOMMU_REGS: Once<SpinLock<IommuRegisters, LocalIrqDisabled>> = Once::new();
 
-pub(super) fn init() -> Result<(), IommuError> {
-    let iommu_regs = IommuRegisters::new().ok_or(IommuError::NoIommu)?;
+pub(super) fn init(io_mem_builder: &IoMemAllocatorBuilder) -> Result<(), IommuError> {
+    let iommu_regs = IommuRegisters::new(io_mem_builder).ok_or(IommuError::NoIommu)?;
     IOMMU_REGS.call_once(|| SpinLock::new(iommu_regs));
     Ok(())
 }

ostd/src/arch/x86/kernel/apic/ioapic.rs

@@ -17,6 +17,7 @@ use volatile::{
 use crate::{
     arch::{iommu::has_interrupt_remapping, x86::kernel::acpi::get_platform_info},
     if_tdx_enabled,
+    io::IoMemAllocatorBuilder,
     mm::paddr_to_vaddr,
     sync::SpinLock,
     trap::IrqLine,
@@ -142,7 +143,8 @@ impl IoApicAccess {
     /// # Safety
     ///
     /// User must ensure the base address is valid.
-    unsafe fn new(base_address: usize) -> Self {
+    unsafe fn new(base_address: usize, io_mem_builder: &IoMemAllocatorBuilder) -> Self {
+        io_mem_builder.remove(base_address..(base_address + 0x20));
         let base = NonNull::new(paddr_to_vaddr(base_address) as *mut u8).unwrap();
         let register = VolatileRef::new_restricted(WriteOnly, base.cast::<u32>());
         let data = VolatileRef::new(base.add(0x10).cast::<u32>());
@@ -178,7 +180,7 @@ impl IoApicAccess {
 
 pub static IO_APIC: Once<Vec<SpinLock<IoApic>>> = Once::new();
 
-pub fn init() {
+pub fn init(io_mem_builder: &IoMemAllocatorBuilder) {
     let Some(platform_info) = get_platform_info() else {
         IO_APIC.call_once(|| {
             // FIXME: Is it possible to have an address that is not the default 0xFEC0_0000?
@@ -193,7 +195,7 @@ pub fn init() {
                     tdx_guest::unprotect_gpa_range(IO_APIC_DEFAULT_ADDRESS, 1).unwrap();
                 }
             });
-            let mut io_apic = unsafe { IoApicAccess::new(IO_APIC_DEFAULT_ADDRESS) };
+            let mut io_apic = unsafe { IoApicAccess::new(IO_APIC_DEFAULT_ADDRESS, io_mem_builder) };
             io_apic.set_id(0);
             let id = io_apic.id();
             let version = io_apic.version();
@@ -225,7 +227,8 @@ pub fn init() {
                     }
                 });
                 let interrupt_base = io_apic.global_system_interrupt_base;
-                let mut io_apic = unsafe { IoApicAccess::new(io_apic.address as usize) };
+                let mut io_apic =
+                    unsafe { IoApicAccess::new(io_apic.address as usize, io_mem_builder) };
                 io_apic.set_id(id as u8);
                 let id = io_apic.id();
                 let version = io_apic.version();

ostd/src/arch/x86/kernel/apic/mod.rs

@@ -8,8 +8,9 @@ use core::{
 
 use bit_field::BitField;
 use spin::Once;
+use xapic::get_xapic_base_address;
 
-use crate::{cpu::PinCurrentCpu, cpu_local};
+use crate::{cpu::PinCurrentCpu, cpu_local, io::IoMemAllocatorBuilder};
 
 pub mod ioapic;
 pub mod x2apic;
@@ -346,7 +347,7 @@ pub enum DivideConfig {
     Divide128 = 0b1010,
 }
 
-pub fn init() -> Result<(), ApicInitError> {
+pub fn init(io_mem_builder: &IoMemAllocatorBuilder) -> Result<(), ApicInitError> {
     crate::arch::x86::kernel::pic::disable_temp();
     if x2apic::X2Apic::has_x2apic() {
         log::info!("x2APIC found!");
@@ -354,6 +355,8 @@ pub fn init() -> Result<(), ApicInitError> {
         Ok(())
     } else if xapic::XApic::has_xapic() {
         log::info!("xAPIC found!");
+        let base_address = get_xapic_base_address();
+        io_mem_builder.remove(base_address..(base_address + size_of::<[u32; 256]>()));
         APIC_TYPE.call_once(|| ApicType::XApic);
         Ok(())
     } else {

ostd/src/arch/x86/kernel/apic/xapic.rs

@@ -23,7 +23,7 @@ impl XApic {
         if !Self::has_xapic() {
             return None;
         }
-        let address = mm::paddr_to_vaddr(get_apic_base_address());
+        let address = mm::paddr_to_vaddr(get_xapic_base_address());
         Some(Self {
             mmio_start: address as *mut u32,
         })
@@ -47,7 +47,7 @@ impl XApic {
 
     pub fn enable(&mut self) {
         // Enable xAPIC
-        set_apic_base_address(get_apic_base_address());
+        set_apic_base_address(get_xapic_base_address());
 
         // Set SVR, Enable APIC and set Spurious Vector to 15 (Reserved irq number)
         let svr: u32 = (1 << 8) | 15;
@@ -119,8 +119,8 @@ fn set_apic_base_address(address: usize) {
     }
 }
 
-/// Gets APIC base address
-fn get_apic_base_address() -> usize {
+/// Gets xAPIC base address
+pub(super) fn get_xapic_base_address() -> usize {
     unsafe {
         (x86_64::registers::model_specific::Msr::new(IA32_APIC_BASE_MSR).read() & 0xf_ffff_f000)
             as usize

ostd/src/arch/x86/mod.rs

@@ -80,11 +80,11 @@ pub(crate) unsafe fn late_init_on_bsp() {
 
     kernel::acpi::init();
 
-    let builder = construct_io_mem_allocator_builder();
+    let io_mem_builder = construct_io_mem_allocator_builder();
 
-    match kernel::apic::init() {
+    match kernel::apic::init(&io_mem_builder) {
         Ok(_) => {
-            ioapic::init();
+            ioapic::init(&io_mem_builder);
         }
         Err(err) => {
             info!("APIC init error:{:?}", err);
@@ -99,7 +99,7 @@ pub(crate) unsafe fn late_init_on_bsp() {
 
     if_tdx_enabled!({
     } else {
-        match iommu::init() {
+        match iommu::init(&io_mem_builder) {
             Ok(_) => {}
             Err(err) => warn!("IOMMU initialization error:{:?}", err),
         }
@@ -110,7 +110,7 @@ pub(crate) unsafe fn late_init_on_bsp() {
 
     // SAFETY: All the system device memory I/Os have been removed from the builder.
     unsafe {
-        crate::io::init(builder);
+        crate::io::init(io_mem_builder);
     }
 }
 

ostd/src/io/io_mem/allocator.rs

@@ -113,8 +113,8 @@ pub static IO_MEM_ALLOCATOR: Once<IoMemAllocator> = Once::new();
 ///
 /// User must ensure all the memory I/O regions that belong to the system device have been removed by calling the
 /// `remove` function.
-pub(crate) unsafe fn init(builder: IoMemAllocatorBuilder) {
-    IO_MEM_ALLOCATOR.call_once(|| IoMemAllocator::new(builder.allocators));
+pub(crate) unsafe fn init(io_mem_builder: IoMemAllocatorBuilder) {
+    IO_MEM_ALLOCATOR.call_once(|| IoMemAllocator::new(io_mem_builder.allocators));
 }
 
 fn find_allocator<'a>(