kernel: fix and make robust the ReentrantLock kernel ID fetch function

Author: Qix-

oro-kernel/src/lib.rs

@@ -21,6 +21,7 @@ pub mod module;
 pub mod port;
 pub mod ring;
 pub mod scheduler;
+pub mod sync;
 pub mod thread;
 
 use core::{
@@ -74,21 +75,6 @@ pub struct Kernel<A: Arch> {
 	handle:    A::CoreHandle,
 }
 
-// NOTE(qix-): This is an ergonomics hack to avoid `A: Arch` in a lot of places.
-#[doc(hidden)]
-static mut KERNEL_ID_FN: MaybeUninit<fn() -> u32> = MaybeUninit::uninit();
-
-#[doc(hidden)]
-#[no_mangle]
-unsafe extern "C" fn oro_sync_current_core_id() -> u32 {
-	KERNEL_ID_FN.assume_init()()
-}
-
-#[doc(hidden)]
-fn get_arch_kernel_id<A: Arch>() -> u32 {
-	Kernel::<A>::get().id()
-}
-
 impl<A: Arch> Kernel<A> {
 	/// Initializes a new core-local instance of the Oro kernel.
 	///
@@ -116,11 +102,6 @@ impl<A: Arch> Kernel<A> {
 	) -> Result<&'static Self, MapError> {
 		assert::fits::<Self, 4096>();
 
-		#[expect(static_mut_refs)]
-		{
-			KERNEL_ID_FN.write(get_arch_kernel_id::<A>);
-		}
-
 		let mapper = AddressSpace::<A>::current_supervisor_space();
 		let core_local_segment = AddressSpace::<A>::kernel_core_local();
 
@@ -141,6 +122,13 @@ impl<A: Arch> Kernel<A> {
 			mapper,
 		});
 
+		// SAFETY(qix-): Now that the kernel has been mapped in, we can initialize the _real_
+		// SAFETY(qix-): core-local ID function. This is effectively a no-op for secondary
+		// SAFETY(qix-): cores, but doing it here ensures that any stray usage of `ReentrantLock`s
+		// SAFETY(qix-): at least see _some_ core ID. This isn't ideal, it's a bit of a hack, but
+		// SAFETY(qix-): it's a one-off situation that isn't trivial to avoid.
+		sync::initialize_kernel_id_fn::<A>();
+
 		(*kernel_ptr)
 			.scheduler
 			.write(TicketMutex::new(Scheduler::new(&*kernel_ptr)));
@@ -242,6 +230,9 @@ impl<A: Arch> KernelState<A> {
 	/// or else registry accesses will page fault.
 	#[allow(clippy::missing_panics_doc)]
 	pub unsafe fn init(this: &'static mut MaybeUninit<Self>) -> Result<(), MapError> {
+		// SAFETY(qix-): Must be first, before anything else happens in the kernel.
+		self::sync::install_dummy_kernel_id_fn();
+
 		let root_ring = ring::Ring::<A>::new_root()?;
 		let root_ring_weak = Arc::downgrade(&root_ring);
 

oro-kernel/src/sync.rs

@@ -0,0 +1,83 @@
+//! `oro-sync` backing functionality.
+//!
+//! These are mostly just implementations to allow certain advanced features of
+//! the `oro-sync` crate work (e.g. `ReentrantMutex`).
+
+use core::mem::MaybeUninit;
+
+use crate::arch::Arch;
+
+/// Holds the "global" kernel ID function pointer, which retrieves the **core local**
+/// kernel ID for the calling context's current core.
+///
+/// This is an ergonomics hack to avoid `A: Arch` in a lot of places.
+static mut KERNEL_ID_FN: MaybeUninit<fn() -> u32> = MaybeUninit::uninit();
+
+/// Debug field for ensuring the kernel ID function is set.
+#[cfg(debug_assertions)]
+static HAS_SET_KERNEL_ID_FN: core::sync::atomic::AtomicBool =
+	core::sync::atomic::AtomicBool::new(false);
+
+/// Retrieves the current core's kernel ID. This is linked to by `oro-sync` for the
+/// [`oro_sync::ReentrantLock`] implementation.
+#[doc(hidden)]
+#[no_mangle]
+unsafe extern "C" fn oro_sync_current_core_id() -> u32 {
+	#[cfg(debug_assertions)]
+	{
+		assert!(
+			HAS_SET_KERNEL_ID_FN.load(core::sync::atomic::Ordering::Relaxed),
+			"kernel ID function not set"
+		);
+	}
+
+	KERNEL_ID_FN.assume_init()()
+}
+
+/// The generic kernel ID fetcher, based on the [`Arch`] type.
+#[doc(hidden)]
+fn get_arch_kernel_id<A: Arch>() -> u32 {
+	crate::Kernel::<A>::get().id()
+}
+
+/// Initializes the kernel ID function pointer.
+///
+/// # Safety
+/// This must be called **exactly once** prior to any [`oro_sync::ReentrantMutex`] locking.
+///
+/// Further, `A` **must** be the same [`Arch`] type as the kernel being initialized,
+/// and **must** be the same across **all cores**.
+pub unsafe fn initialize_kernel_id_fn<A: Arch>() {
+	#[cfg(debug_assertions)]
+	{
+		HAS_SET_KERNEL_ID_FN.store(true, core::sync::atomic::Ordering::Relaxed);
+	}
+
+	// SAFETY(qix-): We have offloaded safety considerations to the caller here.
+	#[expect(static_mut_refs)]
+	{
+		KERNEL_ID_FN.write(get_arch_kernel_id::<A>);
+	}
+}
+
+/// Installs a dummy kernel ID for use during early boot.
+///
+/// # Safety
+/// This function *may* be called prior to [`initialize_kernel_id_fn`], but
+/// **must** be called prior to any locking with [`oro_sync::ReentrantMutex`]
+/// if it is.
+///
+/// **The handler installed by this function must not be used once multiple
+/// cores are active.**
+pub unsafe fn install_dummy_kernel_id_fn() {
+	#[cfg(debug_assertions)]
+	{
+		HAS_SET_KERNEL_ID_FN.store(true, core::sync::atomic::Ordering::Relaxed);
+	}
+
+	// SAFETY(qix-): We have offloaded safety considerations to the caller here.
+	#[expect(static_mut_refs)]
+	{
+		KERNEL_ID_FN.write(|| 0);
+	}
+}