Make some boot entry functions unsafe and document safety

The following functions are unsafe now:
- riscv_boot
- loongarch_boot
- ap_early_entry
- kernel_task_entry

Remove extern declaration of ap_early_entry from riscv_ap_early_entry
And
- Directly call ap_early_entry from module path.
- Clarify what safety requirements are on riscv_ap_early_entry
- Clarify how safety requirements are met on calling ap_early_entry

ostd/src/arch/loongarch/boot/mod.rs

@@ -99,8 +99,17 @@ fn check_cpu_config() {
 /// The entry point of the Rust code portion of Asterinas.
 ///
 /// Reference: <https://docs.kernel.org/arch/loongarch/booting.html#information-passed-from-bootloader-to-kernel>
+///
+/// # Safety
+///
+/// - This function must be called only once through assembly code.
+/// - The caller must follow C calling conventions and put the right arguments in registers.
 #[no_mangle]
-pub extern "C" fn loongarch_boot(_efi_boot: usize, cmdline_paddr: usize, systab_paddr: usize) -> ! {
+unsafe extern "C" fn loongarch_boot(
+    _efi_boot: usize,
+    cmdline_paddr: usize,
+    systab_paddr: usize,
+) -> ! {
     check_cpu_config();
 
     let systab_ptr = paddr_to_vaddr(systab_paddr) as *const EfiSystemTable;

ostd/src/arch/riscv/boot/bsp_boot.S

@@ -25,7 +25,7 @@ _start:
     # Arguments passed from SBI:
     #   a0 = hart id
     #   a1 = device tree paddr
-    # We do not touch them here. They are passed to the Rust entrypoint.
+    # We do not touch them here. They are passed to the Rust entrypoint `riscv_boot`.
 
     # Set up the Sv48 page table.
     #   sv48_boot_l4pt[511] = (PPN(sv48_boot_l3pt) << PTE_PPN_SHIFT) | PTE_V

ostd/src/arch/riscv/boot/mod.rs

@@ -107,8 +107,15 @@ fn parse_initramfs_range() -> Option<(usize, usize)> {
 static mut BOOTSTRAP_HART_ID: u32 = u32::MAX;
 
 /// The entry point of the Rust code portion of Asterinas.
+///
+/// `BOOTSTRAP_HART_ID` is initialized to be `hart_id` and accessible after calling this.
+///
+/// # Safety
+///
+/// - This function must be called only once through assembly code.
+/// - The caller must follow C calling conventions and put the right arguments in registers.
 #[no_mangle]
-pub extern "C" fn riscv_boot(hart_id: usize, device_tree_paddr: usize) -> ! {
+unsafe extern "C" fn riscv_boot(hart_id: usize, device_tree_paddr: usize) -> ! {
     early_println!("Enter riscv_boot");
 
     // SAFETY: We only write it once this time. Other processors will only read

ostd/src/arch/riscv/boot/smp.rs

@@ -5,7 +5,7 @@
 use core::arch::global_asm;
 
 use crate::{
-    boot::smp::PerApRawInfo,
+    boot::smp::{ap_early_entry, PerApRawInfo},
     cpu_local_cell,
     mm::{Paddr, Vaddr},
 };
@@ -187,16 +187,19 @@ cpu_local_cell! {
 // Since in RISC-V we cannot read the hart ID in S mode, the hart ID is
 // delivered from the bootloader. We need to record the hart ID with another
 // layer of entry point.
+//
+/// # Safety
+///
+/// - This function must be called only once on each AP through assembly code.
+/// - The caller must follow C calling conventions and put the right arguments in registers.
 #[no_mangle]
 unsafe extern "C" fn riscv_ap_early_entry(cpu_id: u32, hart_id: u32) -> ! {
-    unsafe extern "C" {
-        fn ap_early_entry(cpu_id: u32) -> !;
-    }
-
     // CPU local memory could be accessed here since we are the AP and the BSP
     // must have initialized it.
     AP_CURRENT_HART_ID.store(hart_id);
 
-    // SAFETY: This is valid to call and only called once.
-    unsafe { ap_early_entry(cpu_id) };
+    // SAFETY: This is valid to call, because
+    // - the caller guarantees being called only once on each AP.
+    // - the compiler guarantees the C calling conventions on this call.
+    unsafe { ap_early_entry(cpu_id) }
 }

ostd/src/arch/x86/boot/ap_boot.S

@@ -152,6 +152,7 @@ __ap_boot_cpu_id_tail:
 .code64
 
 ap_long_mode:
+    // Argument passed to ap_early_entry: rdi = cpu_id
     mov rdi, 1
     lock xadd [__ap_boot_cpu_id_tail], rdi
 

ostd/src/boot/smp.rs

@@ -124,8 +124,14 @@ pub fn register_ap_entry(entry: fn()) {
     AP_LATE_ENTRY.call_once(|| entry);
 }
 
+/// Early boot entry of an AP.
+///
+/// # Safety
+///
+/// - This function must be called only once on each AP.
+/// - The caller must follow C calling conventions and put the right arguments in registers.
 #[no_mangle]
-fn ap_early_entry(cpu_id: u32) -> ! {
+pub(crate) unsafe extern "C" fn ap_early_entry(cpu_id: u32) -> ! {
     // SAFETY:
     // 1. We're in the boot context of an AP.
     // 2. The CPU ID of the AP is correct.

ostd/src/task/mod.rs

@@ -165,8 +165,12 @@ impl TaskOptions {
         //
         // We provide an assembly wrapper for this function as the end of call stack so we
         // have to disable name mangling for it.
+        //
+        // # Safety
+        //
+        // This function must be called from `switch.S` when the context is prepared correctly.
         #[no_mangle]
-        extern "C" fn kernel_task_entry() -> ! {
+        unsafe extern "C" fn kernel_task_entry() -> ! {
             // SAFETY: The new task is switched on a CPU for the first time, `after_switching_to`
             // hasn't been called yet.
             unsafe { processor::after_switching_to() };