Allow arbitrary closures in start_core

src/entry.rs

@@ -4,7 +4,9 @@
 
 //! Entrypoint code
 
-use core::arch::naked_asm;
+use core::{arch::naked_asm, mem::offset_of};
+
+use crate::StartCoreStack;
 
 /// This is a generic entry point for an image. It carries out the operations required to prepare the
 /// loaded image to be run. Specifically, it zeroes the bss section using registers x25 and above,
@@ -51,8 +53,8 @@ unsafe extern "C" fn entry() -> ! {
 /// An assembly entry point for secondary cores.
 ///
 /// It will enable the MMU, disable trapping of floating point instructions, initialise the
-/// stack pointer to `stack_end` and then jump to the function pointer at the bottom of the
-/// stack with the u64 value second on the stack as a parameter.
+/// stack pointer to `stack_end` and then jump to the trampoline function pointer at the bottom
+/// of the stack with the closure pointer second on the stack as a parameter.
 ///
 /// # Safety
 ///
@@ -69,15 +71,20 @@ pub unsafe extern "C" fn secondary_entry(stack_end: *mut u64) -> ! {
         "isb",
         // Set the stack pointer which was passed.
         "mov sp, x0",
-        // Load Rust entry point address and argument from the bottom of the stack into
-        // callee-saved registers.
-        "ldp x19, x20, [sp, #-16]",
+        // Load the closure address into x19 and the trampoline address into x20.
+        // This is loaded from StartCoreStack.
+        "ldr x19, [sp, #{entry_ptr_offset}]",
+        "ldr x20, [sp, #{trampoline_ptr_offset}]",
         // Set the exception vector.
         "bl {set_exception_vector}",
-        // Pass argument to Rust entry point.
+        // Pass the entry point (closure) address to the trampoline function.
         "mov x0, x19",
-        // Call into Rust code.
+        // Call into Rust trampoline.
         "br x20",
+        entry_ptr_offset = const offset_of!(StartCoreStack<()>, entry_ptr) as isize
+            - size_of::<StartCoreStack<()>>() as isize,
+        trampoline_ptr_offset = const offset_of!(StartCoreStack<()>, trampoline_ptr) as isize
+            - size_of::<StartCoreStack<()>>() as isize,
         set_exception_vector = sym crate::set_exception_vector,
     )
 }

src/lib.rs

@@ -46,6 +46,8 @@ pub use pagetable::{
     InitialPagetable,
 };
 
+use core::mem::ManuallyDrop;
+
 #[cfg(not(feature = "initial-pagetable"))]
 #[unsafe(naked)]
 #[unsafe(link_section = ".init")]
@@ -209,37 +211,71 @@ impl StackPage {
     }
 }
 
+#[repr(C)]
+pub(crate) struct StartCoreStack<F> {
+    entry_ptr: *mut ManuallyDrop<F>,
+    trampoline_ptr: unsafe extern "C" fn(&mut ManuallyDrop<F>) -> !,
+}
+
 #[cfg(feature = "psci")]
 /// Issues a PSCI CPU_ON call to start the CPU core with the given MPIDR.
 ///
 /// This starts the core with an assembly entry point which will enable the MMU, disable trapping of
 /// floating point instructions, initialise the stack pointer to the given value, and then jump to
 /// the given Rust entry point function, passing it the given argument value.
 ///
+/// The closure passed as `rust_entry` **should never return**. Because the
+/// [never type has not been stabilized](https://github.com/rust-lang/rust/issues/35121)), this
+/// cannot be enforced by the type system yet.
+///
 /// # Safety
 ///
 /// `stack` must point to a region of memory which is reserved for this core's stack. It must remain
 /// valid as long as the core is running, and there must not be any other access to it during that
 /// time. It must be mapped both for the current core to write to it (to pass initial parameters)
 /// and in the initial page table which the core being started will used, with the same memory
 /// attributes for both.
-pub unsafe fn start_core<C: smccc::Call, const N: usize>(
+// TODO: change `F` generic bounds to `FnOnce() -> !` when the never type is stabilized:
+// https://github.com/rust-lang/rust/issues/35121
+pub unsafe fn start_core<C: smccc::Call, F: FnOnce() + Send + 'static, const N: usize>(
     mpidr: u64,
     stack: *mut Stack<N>,
-    rust_entry: extern "C" fn(arg: u64) -> !,
-    arg: u64,
+    rust_entry: F,
 ) -> Result<(), smccc::psci::Error> {
+    const {
+        assert!(
+            core::mem::size_of::<StartCoreStack<F>>()
+                + 2 * core::mem::size_of::<F>()
+                + 2 * core::mem::align_of::<F>()
+                + 1024 // trampoline stack frame overhead
+                <= core::mem::size_of::<Stack<N>>(),
+            "the `rust_entry` closure is too big to fit in the core stack"
+        );
+    }
+
+    let rust_entry = ManuallyDrop::new(rust_entry);
+
+    let stack_start = stack.cast::<u8>();
+    let align_offfset = stack_start.align_offset(core::mem::align_of::<F>());
+    let entry_ptr = stack_start
+        .wrapping_add(align_offfset)
+        .cast::<ManuallyDrop<F>>();
+
     assert!(stack.is_aligned());
     // The stack grows downwards on aarch64, so get a pointer to the end of the stack.
     let stack_end = stack.wrapping_add(1);
+    let params = stack_end.cast::<StartCoreStack<F>>().wrapping_sub(1);
 
-    // Write Rust entry point to the stack, so the assembly entry point can jump to it.
-    let params = stack_end as *mut u64;
+    // Write the trampoline and entry closure, so the assembly entry point can jump to it.
     // SAFETY: Our caller promised that the stack is valid and nothing else will access it.
     unsafe {
-        *params.wrapping_sub(1) = rust_entry as usize as _;
-        *params.wrapping_sub(2) = arg;
-    }
+        entry_ptr.write(rust_entry);
+        *params = StartCoreStack {
+            entry_ptr,
+            trampoline_ptr: trampoline::<F>,
+        };
+    };
+
     // Wait for the stores above to complete before starting the secondary CPU core.
     dsb_st();
 
@@ -250,6 +286,24 @@ pub unsafe fn start_core<C: smccc::Call, const N: usize>(
     )
 }
 
+#[cfg(feature = "psci")]
+/// Used by [`start_core`] as an entry point for the secondary CPU core.
+///
+/// # Safety
+///
+/// This calls [`ManuallyDrop::take`] on the provided argument, so this function must be
+/// called at most once for a given instance of `F`.
+// TODO: change `F` generic bounds to `FnOnce() -> !` when the never type is stabilized:
+// https://github.com/rust-lang/rust/issues/35121
+unsafe extern "C" fn trampoline<F: FnOnce() + Send + 'static>(entry: &mut ManuallyDrop<F>) -> ! {
+    // SAFETY: the trampoline function is only ever called once after creating ManuallyDrop
+    // instance, so we won't call ManuallyDrop::take more than once.
+    let entry = unsafe { ManuallyDrop::take(entry) };
+    entry();
+
+    panic!("rust_entry function passed to start_core should never return");
+}
+
 /// Data synchronisation barrier that waits for stores to complete, for the full system.
 #[cfg(feature = "psci")]
 fn dsb_st() {