From 68bdda4c4c2241991662e1bea4bac99bf22992bd Mon Sep 17 00:00:00 2001
From: Zhang Junyang <junyang@stu.pku.edu.cn>
Date: Fri, 13 Dec 2024 21:23:52 +0800
Subject: [PATCH] Set the page table for APs before kicking

---
 ostd/src/arch/x86/boot/ap_boot.S  | 12 +++++++++---
 ostd/src/arch/x86/boot/smp.rs     | 25 ++++++++++++++++++++++---
 ostd/src/mm/page_table/boot_pt.rs | 15 ++++++++++-----
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/ostd/src/arch/x86/boot/ap_boot.S b/ostd/src/arch/x86/boot/ap_boot.S
index ba8273446..76ab6c4df 100644
--- a/ostd/src/arch/x86/boot/ap_boot.S
+++ b/ostd/src/arch/x86/boot/ap_boot.S
@@ -83,7 +83,13 @@ x2apic_mode:
     rdmsr
     jmp ap_protect
 
-.code32
+// This is a pointer to the page table used by the APs.
+// The BSP will fill this pointer before kicking the APs.
+.global __boot_page_table_pointer
+.align 4
+__boot_page_table_pointer:
+    .skip 4
+
 ap_protect:
     // Save the local APIC ID in an unused register.
     // We will calculate the stack pointer of this core 
@@ -103,7 +109,7 @@ ap_protect:
     // Set the page table. The application processors use
     // the same page table as the bootstrap processor's
     // boot phase page table.
-    lea eax, [boot_page_table_start]
+    mov eax, __boot_page_table_pointer
     mov cr3, eax
 
     // Enable long mode.
@@ -132,9 +138,9 @@ ap_long_mode_in_low_address:
     mov rax, offset ap_long_mode
     jmp rax
 
+.data
 // This is a pointer to be filled by the BSP when boot stacks
 // of all APs are allocated and initialized.
-.data
 .global __ap_boot_stack_array_pointer
 .align 8
 __ap_boot_stack_array_pointer:
diff --git a/ostd/src/arch/x86/boot/smp.rs b/ostd/src/arch/x86/boot/smp.rs
index 55df78520..3b5ac609d 100644
--- a/ostd/src/arch/x86/boot/smp.rs
+++ b/ostd/src/arch/x86/boot/smp.rs
@@ -57,7 +57,8 @@ pub(crate) fn get_num_processors() -> Option<u32> {
 /// Brings up all application processors.
 pub(crate) fn bringup_all_aps() {
     copy_ap_boot_code();
-    init_boot_stack_array();
+    fill_boot_stack_array_ptr();
+    fill_boot_pt_ptr();
     send_boot_ipis();
 }
 
@@ -85,7 +86,7 @@ fn copy_ap_boot_code() {
 }
 
 /// Initializes the boot stack array in the AP boot code with the given pages.
-fn init_boot_stack_array() {
+fn fill_boot_stack_array_ptr() {
     let pages = &crate::boot::smp::AP_BOOT_INFO
         .get()
         .unwrap()
@@ -97,7 +98,7 @@ fn init_boot_stack_array() {
     }
     let ap_boot_stack_arr_ptr: *mut u64 = __ap_boot_stack_array_pointer as usize as *mut u64;
     log::debug!(
-        "__ap_boot_stack_array_pointer: {:#x?}",
+        "Setting __ap_boot_stack_array_pointer={:#x?} for AP boot stacks",
         ap_boot_stack_arr_ptr
     );
 
@@ -107,6 +108,24 @@ fn init_boot_stack_array() {
     }
 }
 
+fn fill_boot_pt_ptr() {
+    // This is defined in the boot assembly code.
+    extern "C" {
+        fn __boot_page_table_pointer();
+    }
+    let boot_pt_ptr: *mut u32 = __boot_page_table_pointer as usize as *mut u32;
+    let boot_pt = crate::mm::page_table::boot_pt::with_borrow(|pt| pt.root_address()).unwrap();
+    log::debug!(
+        "Setting __boot_page_table_pointer={:#x?} for AP boot page tables",
+        boot_pt
+    );
+
+    // SAFETY: this pointer points to a static variable defined in the `ap_boot.S`.
+    unsafe {
+        boot_pt_ptr.write_volatile(boot_pt as u32);
+    }
+}
+
 // The symbols are defined in the linker script.
 extern "C" {
     fn __ap_boot_start();
diff --git a/ostd/src/mm/page_table/boot_pt.rs b/ostd/src/mm/page_table/boot_pt.rs
index 39f01001a..539e55fe1 100644
--- a/ostd/src/mm/page_table/boot_pt.rs
+++ b/ostd/src/mm/page_table/boot_pt.rs
@@ -16,7 +16,7 @@ use crate::{
     cpu::num_cpus,
     cpu_local_cell,
     mm::{
-        nr_subpage_per_huge, paddr_to_vaddr, page::allocator::PAGE_ALLOCATOR, PageProperty,
+        nr_subpage_per_huge, paddr_to_vaddr, page::allocator::PAGE_ALLOCATOR, Paddr, PageProperty,
         PagingConstsTrait, Vaddr, PAGE_SIZE,
     },
     sync::SpinLock,
@@ -32,9 +32,9 @@ type FrameNumber = usize;
 ///
 /// The boot page table will be dropped when there's no CPU activating it.
 /// This function will return an [`Err`] if the boot page table is dropped.
-pub(crate) fn with_borrow<F>(f: F) -> Result<(), ()>
+pub(crate) fn with_borrow<F, R>(f: F) -> Result<R, ()>
 where
-    F: FnOnce(&mut BootPageTable),
+    F: FnOnce(&mut BootPageTable) -> R,
 {
     let mut boot_pt = BOOT_PAGE_TABLE.lock();
 
@@ -48,9 +48,9 @@ where
         *boot_pt = Some(unsafe { BootPageTable::from_current_pt() });
     }
 
-    f(boot_pt.as_mut().unwrap());
+    let r = f(boot_pt.as_mut().unwrap());
 
-    Ok(())
+    Ok(r)
 }
 
 /// Dismiss the boot page table.
@@ -115,6 +115,11 @@ impl<E: PageTableEntryTrait, C: PagingConstsTrait> BootPageTable<E, C> {
         }
     }
 
+    /// Returns the root physical address of the boot page table.
+    pub(crate) fn root_address(&self) -> Paddr {
+        self.root_pt * C::BASE_PAGE_SIZE
+    }
+
     /// Maps a base page to a frame.
     ///
     /// # Panics