Use cfg_if to group tdx cfg block

kernel/aster-nix/src/device/mod.rs

@@ -1,20 +1,26 @@
 // SPDX-License-Identifier: MPL-2.0
 
+use cfg_if::cfg_if;
+
 mod null;
 mod pty;
 mod random;
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-mod tdxguest;
 pub mod tty;
 mod urandom;
 mod zero;
 
+cfg_if! {
+    if #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))] {
+        mod tdxguest;
+
+        use tdx_guest::tdx_is_enabled;
+
+        pub use tdxguest::TdxGuest;
+    }
+}
+
 pub use pty::{new_pty_pair, PtyMaster, PtySlave};
 pub use random::Random;
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use tdx_guest::tdx_is_enabled;
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-pub use tdxguest::TdxGuest;
 pub use urandom::Urandom;
 
 use self::tty::get_n_tty;
@@ -34,11 +40,14 @@ pub fn init() -> Result<()> {
     add_node(console, "console")?;
     let tty = Arc::new(tty::TtyDevice);
     add_node(tty, "tty")?;
-    #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-    let tdx_guest = Arc::new(tdxguest::TdxGuest);
-    #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-    if tdx_is_enabled() {
-        add_node(tdx_guest, "tdx_guest")?;
+    cfg_if! {
+        if #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))] {
+            let tdx_guest = Arc::new(tdxguest::TdxGuest);
+
+            if tdx_is_enabled() {
+                add_node(tdx_guest, "tdx_guest")?;
+            }
+        }
     }
     let random = Arc::new(random::Random);
     add_node(random, "random")?;

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

@@ -10,20 +10,24 @@ use core::{
 };
 
 use bitflags::bitflags;
+use cfg_if::cfg_if;
 use log::debug;
-#[cfg(feature = "intel_tdx")]
-use tdx_guest::tdcall;
 pub use trapframe::GeneralRegs as RawGeneralRegs;
 use trapframe::UserContext as RawUserContext;
 use x86_64::registers::rflags::RFlags;
 
-#[cfg(feature = "intel_tdx")]
-use crate::arch::tdx_guest::{handle_virtual_exception, TdxTrapFrame};
 use crate::{
     trap::call_irq_callback_functions,
     user::{ReturnReason, UserContextApi, UserContextApiInternal},
 };
 
+cfg_if! {
+    if #[cfg(feature = "intel_tdx")] {
+        use tdx_guest::tdcall;
+        use crate::arch::tdx_guest::{handle_virtual_exception, TdxTrapFrame};
+    }
+}
+
 /// Cpu context, including both general-purpose registers and floating-point registers.
 #[derive(Clone, Default, Copy, Debug)]
 #[repr(C)]

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

@@ -4,20 +4,24 @@
 
 use alloc::{vec, vec::Vec};
 
-#[cfg(feature = "intel_tdx")]
-use ::tdx_guest::tdx_is_enabled;
 use acpi::PlatformInfo;
 use bit_field::BitField;
+use cfg_if::cfg_if;
 use log::info;
 use spin::Once;
 
-#[cfg(feature = "intel_tdx")]
-use crate::arch::tdx_guest;
 use crate::{
     arch::x86::kernel::acpi::ACPI_TABLES, mm::paddr_to_vaddr, sync::SpinLock, trap::IrqLine, Error,
     Result,
 };
 
+cfg_if! {
+    if #[cfg(feature = "intel_tdx")] {
+        use ::tdx_guest::tdx_is_enabled;
+        use crate::arch::tdx_guest;
+    }
+}
+
 /// I/O Advanced Programmable Interrupt Controller. It is used to distribute external interrupts
 /// in a more advanced manner than that of the standard 8259 PIC.
 ///

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

@@ -5,6 +5,7 @@
 use alloc::fmt;
 use core::ops::Range;
 
+use cfg_if::cfg_if;
 pub(crate) use util::__memcpy_fallible;
 use x86_64::{instructions::tlb, structures::paging::PhysFrame, VirtAddr};
 
@@ -136,11 +137,13 @@ pub fn current_page_table_paddr() -> Paddr {
 }
 
 impl PageTableEntry {
-    /// 51:12
-    #[cfg(not(feature = "intel_tdx"))]
-    const PHYS_ADDR_MASK: usize = 0xF_FFFF_FFFF_F000;
-    #[cfg(feature = "intel_tdx")]
-    const PHYS_ADDR_MASK: usize = 0x7_FFFF_FFFF_F000;
+    cfg_if! {
+        if #[cfg(feature = "intel_tdx")] {
+            const PHYS_ADDR_MASK: usize = 0x7_FFFF_FFFF_F000;
+        } else {
+            const PHYS_ADDR_MASK: usize = 0xF_FFFF_FFFF_F000;
+        }
+    }
     const PROP_MASK: usize = !Self::PHYS_ADDR_MASK & !PageTableFlags::HUGE.bits();
 }
 

ostd/src/arch/x86/mod.rs

@@ -14,11 +14,22 @@ pub(crate) mod pci;
 pub mod qemu;
 pub mod serial;
 pub mod task;
-#[cfg(feature = "intel_tdx")]
-pub(crate) mod tdx_guest;
 pub mod timer;
 pub mod trap;
 
+use cfg_if::cfg_if;
+
+cfg_if! {
+    if #[cfg(feature = "intel_tdx")] {
+        pub(crate) mod tdx_guest;
+
+        use {
+            crate::early_println,
+            ::tdx_guest::{init_tdx, tdcall::InitError, tdx_is_enabled},
+        };
+    }
+}
+
 use core::{
     arch::x86_64::{_rdrand64_step, _rdtsc},
     sync::atomic::Ordering,
@@ -26,11 +37,6 @@ use core::{
 
 use kernel::apic::ioapic;
 use log::{info, warn};
-#[cfg(feature = "intel_tdx")]
-use {
-    crate::early_println,
-    ::tdx_guest::{init_tdx, tdcall::InitError, tdx_is_enabled},
-};
 
 #[cfg(feature = "intel_tdx")]
 pub(crate) fn check_tdx_init() {
@@ -79,18 +85,22 @@ pub(crate) fn init_on_bsp() {
 
     timer::init();
 
-    #[cfg(feature = "intel_tdx")]
-    if !tdx_is_enabled() {
-        match iommu::init() {
-            Ok(_) => {}
-            Err(err) => warn!("IOMMU initialization error:{:?}", err),
+    cfg_if! {
+        if #[cfg(feature = "intel_tdx")] {
+            if !tdx_is_enabled() {
+                match iommu::init() {
+                    Ok(_) => {}
+                    Err(err) => warn!("IOMMU initialization error:{:?}", err),
+                }
+            }
+        } else {
+            match iommu::init() {
+                Ok(_) => {}
+                Err(err) => warn!("IOMMU initialization error:{:?}", err),
+            }
         }
     }
-    #[cfg(not(feature = "intel_tdx"))]
-    match iommu::init() {
-        Ok(_) => {}
-        Err(err) => warn!("IOMMU initialization error:{:?}", err),
-    }
+
     // Some driver like serial may use PIC
     kernel::pic::init();
 }

ostd/src/arch/x86/trap.rs

@@ -3,14 +3,11 @@
 //! Handles trap.
 
 use align_ext::AlignExt;
+use cfg_if::cfg_if;
 use log::debug;
-#[cfg(feature = "intel_tdx")]
-use tdx_guest::{tdcall, tdx_is_enabled};
 use trapframe::TrapFrame;
 
 use super::ex_table::ExTable;
-#[cfg(feature = "intel_tdx")]
-use crate::arch::{cpu::VIRTUALIZATION_EXCEPTION, tdx_guest::handle_virtual_exception};
 use crate::{
     cpu::{CpuException, CpuExceptionInfo, PageFaultErrorCode, PAGE_FAULT},
     cpu_local_cell,
@@ -23,6 +20,13 @@ use crate::{
     trap::call_irq_callback_functions,
 };
 
+cfg_if! {
+    if #[cfg(feature = "intel_tdx")] {
+        use tdx_guest::{tdcall, tdx_is_enabled};
+        use crate::arch::{cpu::VIRTUALIZATION_EXCEPTION, tdx_guest::handle_virtual_exception};
+    }
+}
+
 cpu_local_cell! {
     static IS_KERNEL_INTERRUPTED: bool = false;
 }
@@ -134,14 +138,18 @@ fn handle_kernel_page_fault(f: &TrapFrame, page_fault_vaddr: u64) {
     let vaddr = (page_fault_vaddr as usize).align_down(PAGE_SIZE);
     let paddr = vaddr - LINEAR_MAPPING_BASE_VADDR;
 
-    #[cfg(not(feature = "intel_tdx"))]
-    let priv_flags = PrivFlags::GLOBAL;
-    #[cfg(feature = "intel_tdx")]
-    let priv_flags = if tdx_is_enabled() {
-        PrivFlags::SHARED | PrivFlags::GLOBAL
-    } else {
-        PrivFlags::GLOBAL
-    };
+    cfg_if! {
+        if #[cfg(feature = "intel_tdx")] {
+            let priv_flags = if tdx_is_enabled() {
+                PrivFlags::SHARED | PrivFlags::GLOBAL
+            } else {
+                PrivFlags::GLOBAL
+            };
+        } else {
+            let priv_flags = PrivFlags::GLOBAL;
+        }
+    }
+
     // SAFETY:
     // 1. We have checked that the page fault address falls within the address range of the direct
     //    mapping of physical memory.

ostd/src/bus/mmio/mod.rs

@@ -10,18 +10,22 @@ pub mod common_device;
 use alloc::vec::Vec;
 use core::ops::Range;
 
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use ::tdx_guest::tdx_is_enabled;
+use cfg_if::cfg_if;
 use log::debug;
 
 use self::bus::MmioBus;
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use crate::arch::tdx_guest;
 use crate::{
     arch::kernel::IO_APIC, bus::mmio::common_device::MmioCommonDevice, mm::paddr_to_vaddr,
     sync::SpinLock, trap::IrqLine,
 };
 
+cfg_if! {
+    if #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))] {
+        use ::tdx_guest::tdx_is_enabled;
+        use crate::arch::tdx_guest;
+    }
+}
+
 const VIRTIO_MMIO_MAGIC: u32 = 0x74726976;
 
 /// MMIO bus instance

ostd/src/bus/pci/capability/msix.rs

@@ -7,11 +7,8 @@
 
 use alloc::{sync::Arc, vec::Vec};
 
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use ::tdx_guest::tdx_is_enabled;
+use cfg_if::cfg_if;
 
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use crate::arch::tdx_guest;
 use crate::{
     bus::pci::{
         cfg_space::{Bar, Command, MemoryBar},
@@ -22,6 +19,13 @@ use crate::{
     trap::IrqLine,
 };
 
+cfg_if! {
+    if #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))] {
+        use ::tdx_guest::tdx_is_enabled;
+        use crate::arch::tdx_guest;
+    }
+}
+
 /// MSI-X capability. It will set the BAR space it uses to be hidden.
 #[derive(Debug)]
 #[repr(C)]

ostd/src/mm/dma/dma_coherent.rs

@@ -3,12 +3,9 @@
 use alloc::sync::Arc;
 use core::ops::Deref;
 
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use ::tdx_guest::tdx_is_enabled;
+use cfg_if::cfg_if;
 
 use super::{check_and_insert_dma_mapping, remove_dma_mapping, DmaError, HasDaddr};
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use crate::arch::tdx_guest;
 use crate::{
     arch::{iommu, mm::tlb_flush_addr_range},
     mm::{
@@ -21,6 +18,13 @@ use crate::{
     prelude::*,
 };
 
+cfg_if! {
+    if #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))] {
+        use ::tdx_guest::tdx_is_enabled;
+        use crate::arch::tdx_guest;
+    }
+}
+
 /// A coherent (or consistent) DMA mapping,
 /// which guarantees that the device and the CPU can
 /// access the data in parallel.

ostd/src/mm/dma/dma_stream.rs

@@ -3,12 +3,9 @@
 use alloc::sync::Arc;
 use core::ops::Range;
 
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use ::tdx_guest::tdx_is_enabled;
+use cfg_if::cfg_if;
 
 use super::{check_and_insert_dma_mapping, remove_dma_mapping, DmaError, HasDaddr};
-#[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))]
-use crate::arch::tdx_guest;
 use crate::{
     arch::iommu,
     error::Error,
@@ -18,6 +15,13 @@ use crate::{
     },
 };
 
+cfg_if! {
+    if #[cfg(all(target_arch = "x86_64", feature = "intel_tdx"))] {
+        use ::tdx_guest::tdx_is_enabled;
+        use crate::arch::tdx_guest;
+    }
+}
+
 /// A streaming DMA mapping. Users must synchronize data
 /// before reading or after writing to ensure consistency.
 ///