Fix safety reasoning about get_on_cpu

ostd/src/boot/smp.rs

@@ -140,15 +140,14 @@ pub fn register_ap_entry(entry: fn()) {
 }
 
 #[no_mangle]
-fn ap_early_entry(local_apic_id: u32) -> ! {
-    crate::arch::enable_cpu_features();
-
-    // SAFETY: we are on the AP and they are only called once with the correct
-    // CPU ID.
+fn ap_early_entry(cpu_id: u32) -> ! {
+    // SAFETY: `cpu_id` is the correct value of the CPU ID.
     unsafe {
-        cpu::set_this_cpu_id(local_apic_id);
+        cpu::init_on_ap(cpu_id);
     }
 
+    crate::arch::enable_cpu_features();
+
     // SAFETY: this function is only called once on this AP.
     unsafe {
         crate::arch::trap::init(false);
@@ -169,11 +168,11 @@ fn ap_early_entry(local_apic_id: u32) -> ! {
 
     // Mark the AP as started.
     let ap_boot_info = AP_BOOT_INFO.get().unwrap();
-    ap_boot_info.per_ap_info[local_apic_id as usize - 1]
+    ap_boot_info.per_ap_info[cpu_id as usize - 1]
         .is_started
         .store(true, Ordering::Release);
 
-    log::info!("Processor {} started. Spinning for tasks.", local_apic_id);
+    log::info!("Processor {} started. Spinning for tasks.", cpu_id);
 
     let ap_late_entry = AP_LATE_ENTRY.wait();
     ap_late_entry();

ostd/src/cpu/local/cpu_local.rs

@@ -84,7 +84,7 @@ impl<T: 'static> CpuLocal<T> {
         Self(val)
     }
 
-    /// Get access to the underlying value on the current CPU with a
+    /// Gets access to the underlying value on the current CPU with a
     /// provided IRQ guard.
     ///
     /// By this method, you can borrow a reference to the underlying value
@@ -100,15 +100,10 @@ impl<T: 'static> CpuLocal<T> {
         }
     }
 
-    /// Get access to the underlying value through a raw pointer.
+    /// Gets access to the underlying value through a raw pointer.
     ///
-    /// This function calculates the virtual address of the CPU-local object
-    /// based on the CPU-local base address and the offset in the BSP.
-    ///
-    /// # Safety
-    ///
-    /// The caller must ensure that the reference to `self` is static.
-    pub(crate) unsafe fn as_ptr(&'static self) -> *const T {
+    /// This method is safe, but using the returned pointer will be unsafe.
+    pub(crate) fn as_ptr(&'static self) -> *const T {
         super::is_used::debug_set_true();
 
         let offset = self.get_offset();
@@ -122,7 +117,7 @@ impl<T: 'static> CpuLocal<T> {
         local_va as *mut T
     }
 
-    /// Get the offset of the CPU-local object in the CPU-local area.
+    /// Gets the offset of the CPU-local object in the CPU-local area.
     fn get_offset(&'static self) -> usize {
         let bsp_va = self as *const _ as usize;
         let bsp_base = __cpu_local_start as usize;
@@ -134,11 +129,10 @@ impl<T: 'static> CpuLocal<T> {
 }
 
 impl<T: 'static + Sync> CpuLocal<T> {
-    /// Get access to the copy of value on a specific CPU.
+    /// Gets access to the CPU-local value on a specific CPU.
     ///
-    /// # Panics
-    ///
-    /// Panics if the CPU ID is out of range.
+    /// This allows the caller to access CPU-local data from a remote CPU,
+    /// so the data type must be `Sync`.
     pub fn get_on_cpu(&'static self, cpu_id: CpuId) -> &'static T {
         super::is_used::debug_set_true();
 
@@ -149,21 +143,22 @@ impl<T: 'static + Sync> CpuLocal<T> {
             return &self.0;
         }
 
-        // SAFETY: Here we use `Once::get_unchecked` to make getting the CPU-
-        // local base faster. The storages must be initialized here (since this
-        // is not the BSP) so it is safe to do so.
-        let base = unsafe {
-            *super::CPU_LOCAL_STORAGES
-                .get_unchecked()
-                .get_unchecked(cpu_id - 1)
-        };
-        let base = crate::mm::paddr_to_vaddr(base);
+        // SAFETY: At this time we have a non-BSP `CpuId`, which means that
+        // `init_cpu_nums` must have been called, so `copy_bsp_for_ap` must
+        // also have been called (see the implementation of `cpu::init_on_bsp`),
+        // so `CPU_LOCAL_STORAGES` must already be initialized.
+        let storages = unsafe { super::CPU_LOCAL_STORAGES.get_unchecked() };
+        // SAFETY: `cpu_id` is guaranteed to be in range because the type
+        // invariant of `CpuId`.
+        let storage = unsafe { *storages.get_unchecked(cpu_id - 1) };
+        let base = crate::mm::paddr_to_vaddr(storage);
 
         let offset = self.get_offset();
-
         let ptr = (base + offset) as *const T;
 
-        // SAFETY: The pointer is valid since the initialization is completed.
+        // SAFETY: `ptr` represents CPU-local data on a remote CPU. It
+        // contains valid data, the type is `Sync`, and no one will mutably
+        // borrow it, so creating an immutable borrow here is valid.
         unsafe { &*ptr }
     }
 }

ostd/src/cpu/local/mod.rs

@@ -65,8 +65,11 @@ static CPU_LOCAL_STORAGES: Once<&'static [Paddr]> = Once::new();
 /// function to copy it for the APs. Otherwise, the copied data will
 /// contain non-constant (also non-`Copy`) data, resulting in undefined
 /// behavior when it's loaded on the APs.
-pub(crate) unsafe fn copy_bsp_for_ap() {
-    let num_aps = super::num_cpus() - 1; // BSP does not need allocated storage.
+///
+/// The caller must ensure that the `num_cpus` matches the number of all
+/// CPUs that will access the CPU-local storage.
+pub(crate) unsafe fn copy_bsp_for_ap(num_cpus: usize) {
+    let num_aps = num_cpus - 1; // BSP does not need allocated storage.
     if num_aps == 0 {
         return;
     }

ostd/src/cpu/mod.rs

@@ -5,6 +5,8 @@
 pub mod local;
 pub mod set;
 
+pub use set::{AtomicCpuSet, CpuSet};
+
 cfg_if::cfg_if! {
     if #[cfg(target_arch = "x86_64")] {
         pub use crate::arch::x86::cpu::*;
@@ -13,10 +15,7 @@ cfg_if::cfg_if! {
     }
 }
 
-pub use set::{AtomicCpuSet, CpuSet};
-use spin::Once;
-
-use crate::{arch::boot::smp::count_processors, cpu_local_cell, task::atomic_mode::InAtomicMode};
+use crate::{cpu_local_cell, task::atomic_mode::InAtomicMode};
 
 /// The ID of a CPU in the system.
 ///
@@ -50,39 +49,38 @@ impl TryFrom<usize> for CpuId {
 }
 
 /// The number of CPUs.
-static NUM_CPUS: Once<u32> = Once::new();
+static mut NUM_CPUS: u32 = 1;
 
 /// Initializes the number of CPUs.
 ///
 /// # Safety
 ///
-/// The caller must ensure that this function is called only once on the BSP
-/// at the correct time when the number of CPUs is available from the platform.
-pub(crate) unsafe fn init_num_cpus() {
-    let num_processors = count_processors().unwrap_or(1);
-    NUM_CPUS.call_once(|| num_processors);
-}
+/// The caller must ensure that
+/// 1. We're in the boot context of the BSP and APs have not yet booted.
+/// 2. The argument is the correct value of the number of CPUs (which
+///    is a constant, since we don't support CPU hot-plugging anyway).
+unsafe fn init_num_cpus(num_cpus: u32) {
+    assert!(num_cpus >= 1);
 
-/// Initializes the number of the current CPU.
-///
-/// # Safety
-///
-/// The caller must ensure that this function is called only once on the
-/// correct CPU with the correct CPU ID.
-pub(crate) unsafe fn set_this_cpu_id(id: u32) {
-    CURRENT_CPU.store(id);
+    // SAFETY: It is safe to mutate this global variable because we
+    // are in the boot context.
+    unsafe { NUM_CPUS = num_cpus };
+
+    // Note that decreasing the number of CPUs may break existing
+    // `CpuId`s (which have a type invariant to say that the ID is
+    // less than the number of CPUs).
+    //
+    // However, this never happens: due to the safety conditions
+    // it's only legal to call this function to increase the number
+    // of CPUs from one (the initial value) to the actual number of
+    // CPUs.
 }
 
 /// Returns the number of CPUs.
 pub fn num_cpus() -> usize {
-    debug_assert!(
-        NUM_CPUS.get().is_some(),
-        "The number of CPUs is not initialized"
-    );
-    // SAFETY: The number of CPUs is initialized. The unsafe version is used
-    // to avoid the overhead of the check.
-    let num = unsafe { *NUM_CPUS.get_unchecked() };
-    num as usize
+    // SAFETY: As far as the safe APIs are concerned, `NUM_CPUS` is
+    // read-only, so it is always valid to read.
+    (unsafe { NUM_CPUS }) as usize
 }
 
 /// Returns an iterator over all CPUs.
@@ -90,6 +88,33 @@ pub fn all_cpus() -> impl Iterator<Item = CpuId> {
     (0..num_cpus()).map(|id| CpuId(id as u32))
 }
 
+cpu_local_cell! {
+    /// The current CPU ID.
+    static CURRENT_CPU: u32 = 0;
+    /// The initialization state of the current CPU ID.
+    #[cfg(debug_assertions)]
+    static IS_CURRENT_CPU_INITED: bool = false;
+}
+
+/// Initializes the current CPU ID.
+///
+/// # Safety
+///
+/// This method must be called on each processor during the early
+/// boot phase of the processor.
+///
+/// The caller must ensure that this function is called with
+/// the correct value of the CPU ID.
+unsafe fn set_this_cpu_id(id: u32) {
+    // FIXME: If there are safe APIs that rely on the correctness of
+    // the CPU ID for soundness, we'd better make the CPU ID a global
+    // invariant and initialize it before entering `ap_early_entry`.
+    CURRENT_CPU.store(id);
+
+    #[cfg(debug_assertions)]
+    IS_CURRENT_CPU_INITED.store(true);
+}
+
 /// A marker trait for guard types that can "pin" the current task to the
 /// current CPU.
 ///
@@ -117,9 +142,10 @@ pub unsafe trait PinCurrentCpu {
 /// To ensure that the CPU ID is up-to-date, do it under any guards that
 /// implements the [`PinCurrentCpu`] trait.
 pub fn current_cpu_racy() -> CpuId {
-    let id = CURRENT_CPU.load();
-    debug_assert_ne!(id, u32::MAX, "This CPU is not initialized");
-    CpuId(id)
+    #[cfg(debug_assertions)]
+    assert!(IS_CURRENT_CPU_INITED.load());
+
+    CpuId(CURRENT_CPU.load())
 }
 
 // SAFETY: A guard that enforces the atomic mode requires disabling any
@@ -127,7 +153,35 @@ pub fn current_cpu_racy() -> CpuId {
 unsafe impl<T: InAtomicMode> PinCurrentCpu for T {}
 unsafe impl PinCurrentCpu for dyn InAtomicMode + '_ {}
 
-cpu_local_cell! {
-    /// The number of the current CPU.
-    static CURRENT_CPU: u32 = u32::MAX;
+/// # Safety
+///
+/// The caller must ensure that
+/// 1. We're in the boot context of the BSP and APs have not yet booted.
+/// 2. The number of available processors is available.
+/// 3. No CPU-local objects have been accessed.
+pub(crate) unsafe fn init_on_bsp() {
+    let num_cpus = crate::arch::boot::smp::count_processors().unwrap_or(1);
+
+    // SAFETY: The safety is upheld by the caller and
+    // the correctness of the `get_num_processors` method.
+    unsafe {
+        local::copy_bsp_for_ap(num_cpus as usize);
+
+        set_this_cpu_id(0);
+
+        // Note that `init_num_cpus` should be called after `copy_bsp_for_ap`.
+        // This helps to build the safety reasoning in `CpuLocal::get_on_cpu`.
+        // See its implementation for details.
+        init_num_cpus(num_cpus);
+    }
+}
+
+/// # Safety
+///
+/// The caller must ensure that:
+/// 1. We're in the boot context of an AP.
+/// 2. The CPU ID of the AP is `cpu_id`.
+pub(crate) unsafe fn init_on_ap(cpu_id: u32) {
+    // SAFETY: The safety is upheld by the caller.
+    unsafe { set_this_cpu_id(cpu_id) };
 }

ostd/src/lib.rs

@@ -88,11 +88,7 @@ unsafe fn init() {
     // 1. They are only called once in the boot context of the BSP.
     // 2. The number of CPUs are available because ACPI has been initialized.
     // 3. No CPU-local objects have been accessed yet.
-    unsafe {
-        cpu::init_num_cpus();
-        cpu::local::copy_bsp_for_ap();
-        cpu::set_this_cpu_id(0);
-    }
+    unsafe { cpu::init_on_bsp() };
 
     // SAFETY: We are on the BSP and APs are not yet started.
     let meta_pages = unsafe { mm::frame::meta::init() };