cpu_local/register/
mod.rs

1//! Architecture register primitives and shared validation.
2
3use core::{pin::Pin, ptr::NonNull, sync::atomic::Ordering};
4
5use crate::{CpuAreaRef, CpuLocalError, CpuPin, CurrentThreadHeader, ThreadSwitchError};
6
7#[cfg(all(not(feature = "host-test"), target_arch = "aarch64"))]
8mod aarch64;
9#[cfg(feature = "host-test")]
10mod host;
11#[cfg(all(not(feature = "host-test"), target_arch = "loongarch64"))]
12mod loongarch64;
13#[cfg(all(
14    not(feature = "host-test"),
15    any(target_arch = "riscv32", target_arch = "riscv64")
16))]
17mod riscv;
18#[cfg(all(not(feature = "host-test"), target_arch = "x86_64"))]
19mod x86_64;
20
21#[cfg(all(not(feature = "host-test"), target_arch = "aarch64"))]
22use aarch64 as imp;
23#[cfg(feature = "host-test")]
24use host as imp;
25#[cfg(all(not(feature = "host-test"), target_arch = "loongarch64"))]
26use loongarch64 as imp;
27#[cfg(all(
28    not(feature = "host-test"),
29    any(target_arch = "riscv32", target_arch = "riscv64")
30))]
31use riscv as imp;
32#[cfg(all(not(feature = "host-test"), target_arch = "x86_64"))]
33use x86_64 as imp;
34
35#[cfg(all(
36    not(feature = "host-test"),
37    not(any(
38        target_arch = "x86_64",
39        target_arch = "aarch64",
40        target_arch = "riscv32",
41        target_arch = "riscv64",
42        target_arch = "loongarch64"
43    ))
44))]
45compile_error!("cpu-local supports x86_64, AArch64, RISC-V, and LoongArch64 only");
46
47/// Installs the final area of an offline CPU.
48///
49/// # Safety
50///
51/// The area must remain mapped until shutdown. The CPU must be offline with
52/// traps disabled, and no previous area may be installed on this physical CPU.
53#[doc(hidden)]
54pub unsafe fn install_cpu_area(area: CpuAreaRef) -> Result<(), CpuLocalError> {
55    imp::validate_environment()?;
56    let boot_thread = area.prefix().boot_thread().header();
57    let boot_pointer = boot_thread as *const CurrentThreadHeader as usize;
58    // SAFETY: the caller owns the offline register installation boundary.
59    unsafe { imp::install_cpu_base(area.base(), boot_pointer) };
60    if unsafe { imp::read_cpu_base()? } != area.base() {
61        fatal_register_invariant();
62    }
63    Ok(())
64}
65
66pub(crate) fn current_area() -> Result<CpuAreaRef, CpuLocalError> {
67    let area_base = unsafe { imp::read_cpu_base()? };
68    if area_base == 0 {
69        return Err(CpuLocalError::AreaNotInstalled);
70    }
71    // SAFETY: only install_cpu_area writes the architecture-owned base, and
72    // its contract requires a shutdown-lifetime initialized area.
73    unsafe { CpuAreaRef::from_initialized_base(area_base) }
74}
75
76/// Publishes the scheduler anchor before the architecture switch tail.
77///
78/// # Safety
79///
80/// The caller must own the final IRQ-disabled context-switch boundary. `value`
81/// must identify the prepared pinned header and remain alive while current.
82pub(crate) unsafe fn commit_current_thread(area: CpuAreaRef, value: usize) {
83    area.runtime_anchor()
84        .current_thread_slot()
85        .store(value, Ordering::Release);
86}
87
88/// Returns the pinned current-thread header after checking both sources.
89pub fn current_thread(pin: &CpuPin<'_>) -> Result<NonNull<CurrentThreadHeader>, CpuLocalError> {
90    let area = pin.area();
91    let slot = area.runtime_anchor().current_thread_raw();
92    let register = unsafe { imp::read_current_thread(area.base()) };
93    if slot == 0
94        || slot != register
95        || !slot.is_multiple_of(core::mem::align_of::<CurrentThreadHeader>())
96    {
97        return Err(CpuLocalError::CurrentThreadMismatch);
98    }
99    let pointer = NonNull::new(slot as *mut CurrentThreadHeader)
100        .ok_or(CpuLocalError::CurrentThreadMismatch)?;
101    // SAFETY: scheduler publication only accepts pinned headers that remain
102    // alive while current, and the caller holds the required CPU pin.
103    let thread_area = unsafe { pointer.as_ref() }
104        .cpu_area()
105        .ok_or(CpuLocalError::CurrentThreadMismatch)?;
106    if thread_area != area {
107        return Err(CpuLocalError::CurrentThreadMismatch);
108    }
109    Ok(pointer)
110}
111
112/// Reads the current header before the scheduler can construct its guard.
113///
114/// # Safety
115///
116/// The caller must keep the scheduler-owned current task alive and must not
117/// dereference the result after a context switch.
118#[doc(hidden)]
119pub unsafe fn scheduler_current_thread() -> Result<NonNull<CurrentThreadHeader>, CpuLocalError> {
120    #[cfg(not(feature = "tls"))]
121    {
122        // LinuxCurrent images keep the task pointer in one architecture-owned
123        // source. Reading a CPU area first would race migration because this
124        // function is itself used to construct the preemption guard.
125        let register = unsafe { imp::read_current_thread(0) };
126        NonNull::new(register as *mut CurrentThreadHeader)
127            .ok_or(CpuLocalError::CurrentThreadMismatch)
128    }
129
130    #[cfg(feature = "tls")]
131    loop {
132        // UnikernelTls images keep current in the CPU area's runtime anchor.
133        // Retry if migration changes the base between sampling the area and
134        // loading its slot; the caller cannot be pinned before this lookup.
135        let area = current_area()?;
136        let register = unsafe { imp::read_current_thread(area.base()) };
137        if unsafe { imp::read_cpu_base()? } != area.base() {
138            continue;
139        }
140        return NonNull::new(register as *mut CurrentThreadHeader)
141            .ok_or(CpuLocalError::CurrentThreadMismatch);
142    }
143}
144
145#[cfg(all(test, feature = "host-test"))]
146mod tests {
147    use core::mem::MaybeUninit;
148
149    use super::*;
150    use crate::{CpuAreaPrefix, CpuIndex};
151
152    fn modeled_area(cpu_index: usize) -> CpuAreaRef {
153        let storage = Box::leak(Box::new(MaybeUninit::<CpuAreaPrefix>::uninit()));
154        let base = storage.as_mut_ptr() as usize;
155        storage.write(
156            CpuAreaPrefix::initialize(CpuIndex::try_from(cpu_index).unwrap(), base).unwrap(),
157        );
158        // SAFETY: the initialized fixture is leaked for the process lifetime.
159        unsafe { CpuAreaRef::from_initialized_base(base) }.unwrap()
160    }
161
162    #[test]
163    fn scheduler_current_thread_survives_migration_during_bootstrap_read() {
164        let first = modeled_area(0);
165        let second = modeled_area(1);
166        let first_boot = first.prefix().boot_thread().header();
167        let second_boot = second.prefix().boot_thread().header();
168
169        // SAFETY: this host thread serially owns both leaked CPU fixtures.
170        unsafe { imp::install_cpu_base(first.base(), first_boot as *const _ as usize) };
171        imp::migrate_on_next_current_read(second.base());
172
173        assert_eq!(
174            // SAFETY: both boot headers have process-lifetime storage.
175            unsafe { scheduler_current_thread() },
176            Ok(NonNull::from(second_boot)),
177        );
178    }
179
180    #[test]
181    fn scheduler_current_thread_rejects_an_uninstalled_host_area() {
182        #[cfg(feature = "tls")]
183        let expected_error = CpuLocalError::AreaNotInstalled;
184        #[cfg(not(feature = "tls"))]
185        let expected_error = CpuLocalError::CurrentThreadMismatch;
186
187        let rejected = std::thread::spawn(move || {
188            // SAFETY: the fresh host thread has no installed CPU area, so no
189            // scheduler-owned pointer can be returned.
190            matches!(
191                unsafe { scheduler_current_thread() },
192                Err(error) if error == expected_error
193            )
194        })
195        .join()
196        .expect("host current-thread probe panicked");
197
198        assert!(rejected);
199    }
200}
201
202/// Binds and publishes the first scheduler task on an offline CPU.
203///
204/// # Safety
205///
206/// The CPU must remain offline and trap-free. `header` must stay pinned and
207/// alive until the scheduler replaces it through a prepared switch.
208#[doc(hidden)]
209pub unsafe fn install_bootstrap_thread(
210    pin: &CpuPin<'_>,
211    header: Pin<&CurrentThreadHeader>,
212) -> Result<(), ThreadSwitchError> {
213    let epoch = unsafe { header.bind_cpu(pin.area()) }?;
214    let pointer = header.as_non_null().as_ptr() as usize;
215    unsafe { commit_current_thread(pin.area(), pointer) };
216    // Bootstrap has no raw switch tail. Install the architecture-owned current
217    // register directly while this CPU remains offline and trap-free.
218    unsafe { imp::write_current_thread(pointer) };
219    if current_thread(pin) != Ok(header.as_non_null()) {
220        // The register is already committed, so continuing would make all
221        // later Rust execution unsound. Rollback is intentionally impossible.
222        let _ = epoch;
223        fatal_register_invariant();
224    }
225    Ok(())
226}
227
228/// Reads task-owned kernel TLS under an explicit CPU pin.
229#[cfg(feature = "tls")]
230pub fn kernel_tls(_pin: &CpuPin<'_>) -> usize {
231    unsafe { imp::read_kernel_tls() }
232}
233
234/// Installs task-owned kernel TLS at an offline bootstrap boundary.
235///
236/// # Safety
237///
238/// The caller must own the offline CPU or IRQ-disabled final task switch, and
239/// `value` must remain a valid TLS base for the installed execution context.
240#[cfg(feature = "tls")]
241#[doc(hidden)]
242pub unsafe fn install_kernel_tls(_pin: &CpuPin<'_>, value: usize) {
243    unsafe { imp::write_kernel_tls(value) };
244}
245
246#[cold]
247#[inline(never)]
248fn fatal_register_invariant() -> ! {
249    panic!("CPU-local register commit did not retain the validated state")
250}
cpu_local/register/mod.rs

cpu_local/register/
mod.rs
