fix(vt-d): VT-d interrupt remapping initialization with Invalidation Queue support

awkernel_lib/src/arch/x86_64/interrupt_remap.rs

@@ -1,6 +1,8 @@
 //! Interrupt remapping of Vt-d.
 
 use acpi::AcpiTables;
+use alloc::format;
+use alloc::string::String;
 
 use crate::{
     addr::{phy_addr::PhyAddr, virt_addr::VirtAddr, Addr},
@@ -17,6 +19,42 @@ use super::acpi::AcpiMapper;
 const TABLE_SIZE: usize = 256 * 4096; // 1MiB
 const TABLE_ENTRY_NUM: usize = TABLE_SIZE / 16;
 
+/// Error type for Vt-d interrupt remapping initialization.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum VtdError {
+    /// Failed to enable Interrupt Remapping Table Pointer.
+    IrtpEnableFailed,
+    /// Failed to disable Interrupt Remapping Table Pointer.
+    IrtpDisableFailed,
+    /// Failed to enable Interrupt Remapping.
+    IreEnableFailed,
+    /// Failed to disable Interrupt Remapping.
+    IreDisableFailed,
+    /// Failed to enable Queued Invalidation.
+    QieEnableFailed,
+    /// Failed to disable Queued Invalidation.
+    QieDisableFailed,
+    /// Failed to allocate DMA memory.
+    DmaAllocationFailed,
+    /// Timeout waiting for command completion.
+    Timeout,
+}
+
+impl core::fmt::Display for VtdError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            VtdError::IrtpEnableFailed => write!(f, "Failed to enable IRTP"),
+            VtdError::IrtpDisableFailed => write!(f, "Failed to disable IRTP"),
+            VtdError::IreEnableFailed => write!(f, "Failed to enable IRE"),
+            VtdError::IreDisableFailed => write!(f, "Failed to disable IRE"),
+            VtdError::QieEnableFailed => write!(f, "Failed to enable QIE"),
+            VtdError::QieDisableFailed => write!(f, "Failed to disable QIE"),
+            VtdError::DmaAllocationFailed => write!(f, "DMA allocation failed"),
+            VtdError::Timeout => write!(f, "Timeout waiting for command completion"),
+        }
+    }
+}
+
 #[allow(dead_code)]
 mod registers {
     use bitflags::bitflags;
@@ -60,8 +98,15 @@ mod registers {
         }
     }
 
+    pub const ICS_IWC: u32 = 1; // Invalidation Wait Descriptor Complete
+
+    mmio_r!(offset 0x10 => pub EXTENDED_CAPABILITY<u64>);
     mmio_w!(offset 0x18 => pub GLOBAL_COMMAND<GlobalCommandStatus>);
     mmio_r!(offset 0x1c => pub GLOBAL_STATUS<GlobalCommandStatus>);
+    mmio_rw!(offset 0x80 => pub IQH<u64>); // Invalidation Queue Head
+    mmio_rw!(offset 0x88 => pub IQT<u64>); // Invalidation Queue Tail
+    mmio_rw!(offset 0x90 => pub IQA<u64>); // Invalidation Queue Address
+    mmio_r!(offset 0x9c => pub ICS<u32>); // Invalidation Completion Status
     mmio_rw!(offset 0xb8 => pub IRTA<u64>);
 }
 
@@ -216,37 +261,49 @@ pub unsafe fn init_interrupt_remap(
     phy_offset: VirtAddr,
     acpi: &AcpiTables<AcpiMapper>,
     is_x2apic: bool,
-) -> Result<(), &'static str> {
+) -> Result<(), VtdError> {
     let mut remap_table = [None; 32];
 
     if let Ok(dmar) = acpi.find_table::<Dmar>() {
-        dmar.entries().for_each(|entry| {
+        for entry in dmar.entries() {
             if let DmarEntry::Drhd(drhd) = entry {
                 drhd.device_scopes()
                     .find(|(scope, _path)| scope.entry_type == 1);
 
                 if let Some((phy_addr, _virt_addr)) = &remap_table[drhd.segment_number as usize] {
                     let segment_number = drhd.segment_number as usize;
-                    // Set Interrupt Remapping Table Address Register
-                    set_irta(segment_number, phy_offset, drhd, *phy_addr, is_x2apic);
+                    // Enable Interrupt Remapping
+                    enable_interrupt_remapping(
+                        segment_number,
+                        phy_offset,
+                        drhd,
+                        *phy_addr,
+                        is_x2apic,
+                    )?;
                 } else {
                     let segment_number = drhd.segment_number as usize;
 
                     let pool =
                         DMAPool::<[u8; TABLE_SIZE]>::new(segment_number, TABLE_SIZE / PAGESIZE)
-                            .expect("DMAPool::new() failed.");
+                            .ok_or(VtdError::DmaAllocationFailed)?;
 
                     let virt_addr = pool.get_virt_addr();
                     let phy_addr = pool.get_phy_addr();
                     pool.leak();
 
-                    // Set Interrupt Remapping Table Address Register
-                    set_irta(segment_number, phy_offset, drhd, phy_addr, is_x2apic);
+                    // Enable Interrupt Remapping
+                    enable_interrupt_remapping(
+                        segment_number,
+                        phy_offset,
+                        drhd,
+                        phy_addr,
+                        is_x2apic,
+                    )?;
 
                     remap_table[segment_number] = Some((phy_addr, virt_addr));
                 }
             }
-        });
+        }
     }
 
     // Set INTERRUPT_REMAPPING
@@ -268,14 +325,134 @@ pub unsafe fn init_interrupt_remap(
     Ok(())
 }
 
+/// Update GLOBAL_COMMAND register following Intel VT-d specification.
+///
+/// Steps:
+/// 1. Read GLOBAL_STATUS register
+/// 2. Reset one-shot bits (keeping only persistent bits)
+/// 3. Set/clear the target bit
+/// 4. Write to GLOBAL_COMMAND register
+/// 5. Wait until GLOBAL_STATUS indicates command is serviced
+fn update_global_command(
+    vt_d_base: VirtAddr,
+    bit_to_set: registers::GlobalCommandStatus,
+    enable: bool,
+) -> Result<(), VtdError> {
+    // Step 1: Read GLOBAL_STATUS
+    let status = registers::GLOBAL_STATUS.read(vt_d_base.as_usize());
+
+    // Step 2: Reset one-shot bits
+    // One-shot bits mask: 0x96FFFFFF (bits that should be preserved)
+    let persistent_mask = registers::GlobalCommandStatus::from_bits_truncate(0x96FF_FFFF);
+    let status_persistent = status & persistent_mask;
+
+    // Step 3: Set or clear the target bit
+    let command = if enable {
+        status_persistent | bit_to_set
+    } else {
+        status_persistent & !bit_to_set
+    };
+
+    // Step 4: Write to GLOBAL_COMMAND
+    registers::GLOBAL_COMMAND.write(command, vt_d_base.as_usize());
+
+    // Step 5: Wait until GLOBAL_STATUS indicates command is serviced
+    if wait_toggle_then_set(vt_d_base, bit_to_set, enable).is_err() {
+        // Timeout: bit not updated as expected
+        let current = registers::GLOBAL_STATUS.read(vt_d_base.as_usize());
+        Err(
+            if bit_to_set.contains(registers::GlobalCommandStatus::IRTP)
+                && (enable && !current.contains(registers::GlobalCommandStatus::IRTP))
+            {
+                VtdError::IrtpEnableFailed
+            } else if bit_to_set.contains(registers::GlobalCommandStatus::IRTP)
+                && (!enable && current.contains(registers::GlobalCommandStatus::IRTP))
+            {
+                VtdError::IrtpDisableFailed
+            } else if bit_to_set.contains(registers::GlobalCommandStatus::IRE)
+                && (enable && !current.contains(registers::GlobalCommandStatus::IRE))
+            {
+                VtdError::IreEnableFailed
+            } else if bit_to_set.contains(registers::GlobalCommandStatus::IRE)
+                && (!enable && current.contains(registers::GlobalCommandStatus::IRE))
+            {
+                VtdError::IreDisableFailed
+            } else if bit_to_set.contains(registers::GlobalCommandStatus::QIE)
+                && (enable && !current.contains(registers::GlobalCommandStatus::QIE))
+            {
+                VtdError::QieEnableFailed
+            } else if bit_to_set.contains(registers::GlobalCommandStatus::QIE)
+                && (!enable && current.contains(registers::GlobalCommandStatus::QIE))
+            {
+                VtdError::QieDisableFailed
+            } else {
+                VtdError::Timeout // Default fallback
+            },
+        )
+    } else {
+        Ok(())
+    }
+}
+
+fn wait_toggle_then_set(
+    vt_d_base: VirtAddr,
+    status_bit: registers::GlobalCommandStatus,
+    enable: bool,
+) -> Result<(), VtdError> {
+    if enable {
+        // wait until 1
+        for _ in 0..1000 {
+            if registers::GLOBAL_STATUS
+                .read(vt_d_base.as_usize())
+                .contains(status_bit)
+            {
+                return Ok(());
+            }
+        }
+    } else {
+        // wait until 0
+        for _ in 0..1000 {
+            if !registers::GLOBAL_STATUS
+                .read(vt_d_base.as_usize())
+                .contains(status_bit)
+            {
+                return Ok(());
+            }
+        }
+    }
+
+    Err(VtdError::Timeout)
+}
+
+/// Wait for register-based invalidation to complete.
+/// This must be called before enabling Queued Invalidation.
+fn wait_register_based_invalidation_complete(vt_d_base: VirtAddr) -> Result<(), VtdError> {
+    // Wait for any pending register-based invalidation to complete
+    // by checking the ICS register's IWC bit is 0
+    for _ in 0..1000 {
+        let ics = registers::ICS.read(vt_d_base.as_usize());
+        if (ics & registers::ICS_IWC) == 0 {
+            return Ok(());
+        }
+    }
+
+    log::error!("Timeout waiting for register-based invalidation to complete");
+    Err(VtdError::Timeout)
+}
+
 /// Set Interrupt Remapping Table Address Register.
-fn set_irta(
+///
+/// 1. Set IRTA register
+/// 2. Enable IRTP (Interrupt Remapping Table Pointer)
+/// 3. Enable IRE (Interrupt Remapping Enable)
+/// 4. Initialize Invalidation Queue
+fn enable_interrupt_remapping(
     segment_number: usize,
     phy_offset: VirtAddr,
     drhd: &DmarDrhd,
     phy_addr: PhyAddr,
     is_x2apic: bool,
-) {
+) -> Result<(), VtdError> {
     let vt_d_base = phy_offset + drhd.register_base_address as usize;
 
     // Extended Interrupt Mode (x2APIC) Enable
@@ -284,18 +461,84 @@ fn set_irta(
 
     registers::IRTA.write(val, vt_d_base.as_usize());
 
-    let mut stat = registers::GLOBAL_STATUS.read(vt_d_base.as_usize());
-    stat |= registers::GlobalCommandStatus::IRTP | registers::GlobalCommandStatus::IRE;
-    registers::GLOBAL_COMMAND.write(stat, vt_d_base.as_usize());
+    // Enable IRTP (Interrupt Remapping Table Pointer)
+    update_global_command(vt_d_base, registers::GlobalCommandStatus::IRTP, true)?;
+
+    // Enable IRE (Interrupt Remapping Enable)
+    update_global_command(vt_d_base, registers::GlobalCommandStatus::IRE, true)?;
+
+    // Initialize Invalidation Queue
+    init_invalidation_queue(segment_number, vt_d_base)?;
+
+    // Re-read status after queue initialization
+    let status = registers::GLOBAL_STATUS.read(vt_d_base.as_usize());
+
+    let mut scope_str = String::new();
+
+    for scope in drhd.device_scopes() {
+        let (device_scope, path) = scope;
+        scope_str = format!(
+            "{scope_str}\r\nVt-d Device Scope: entry type = {}, length = {}, flags = 0x{:x}, enumeration ID = {}, start bus number = {}, path = {path:x?}",
+            device_scope.entry_type,
+            device_scope.length,
+            device_scope.flags,
+            device_scope.enumeration_id,
+            device_scope.start_bus_number,
+        );
+    }
 
     log::info!(
-        "Vt-d Interrupt Remapping: Segment = {}, Vt-d Base = 0x{:x}, Table PhyAddr = 0x{:x}, enabled = {}, mode = {}",
+            "Vt-d Interrupt Remapping: Segment = {segment_number}, Vt-d Base = 0x{:x}, Table PhyAddr = 0x{:x}, enabled = {}, mode = {}, extended capability = 0x{:x}, global status = 0x{:x}, IRTA = 0x{:x}, DHRD.flags = 0x{:x}{scope_str}",
+            drhd.register_base_address as usize,
+            phy_addr.as_usize(),
+            status.contains(registers::GlobalCommandStatus::IRTP | registers::GlobalCommandStatus::IRE),
+            if is_x2apic { "x2APIC" } else { "xAPIC" },
+            registers::EXTENDED_CAPABILITY.read(vt_d_base.as_usize()),
+            status.bits(),
+            registers::IRTA.read(vt_d_base.as_usize()),
+            drhd.flags
+        );
+
+    Ok(())
+}
+
+fn init_invalidation_queue(segment_number: usize, vt_d_base: VirtAddr) -> Result<(), VtdError> {
+    // Wait for any pending register-based invalidation to complete
+    // before enabling Queued Invalidation
+    wait_register_based_invalidation_complete(vt_d_base)?;
+
+    // Reset Invalidation Queue Tail
+    registers::IQT.write(0, vt_d_base.as_usize());
+
+    // Allocate a page for the Invalidation Queue
+    let pool =
+        DMAPool::<[u8; PAGESIZE]>::new(segment_number, 1).ok_or(VtdError::DmaAllocationFailed)?;
+
+    let phy_addr = pool.get_phy_addr().as_usize() as u64;
+
+    // bits[63:12] = Base (4KB aligned)
+    // bit[11]     = DW (128/256)
+    // bits[2:0]   = queue size (0 = 4KB)
+    // bits[10:3]  = Reserved
+    let iqa_value = phy_addr & !((1 << 12) - 1);
+    registers::IQA.write(iqa_value, vt_d_base.as_usize());
+
+    // Leak the pool to keep it allocated
+    pool.leak();
+
+    // Enable Queued Invalidation (QIE)
+    update_global_command(vt_d_base, registers::GlobalCommandStatus::QIE, true)?;
+
+    log::info!(
+        "Vt-d Queued Invalidation enabled: Segment = {}, Queue PhyAddr = 0x{:x}, IQA = 0x{:x}, IQH = {}, IQT = {}",
         segment_number,
-        drhd.register_base_address as usize,
-        phy_addr.as_usize(),
-        stat.contains(registers::GlobalCommandStatus::IRE | registers::GlobalCommandStatus::IRTP),
-        if is_x2apic { "x2APIC" } else { "xAPIC" },
-    )
+        phy_addr,
+        registers::IQA.read(vt_d_base.as_usize()),
+        registers::IQH.read(vt_d_base.as_usize()),
+        registers::IQT.read(vt_d_base.as_usize())
+    );
+
+    Ok(())
 }
 
 pub fn allocate_remapping_entry(