monolithic: switch user page table

modules/axhal/src/arch/aarch64/context.rs

@@ -68,7 +68,13 @@ pub struct TaskContext {
 }
 
 impl TaskContext {
-    /// Creates a new default context for a new task.
+    /// Creates a dummy context for a new task.
+    ///
+    /// Note the context is not initialized, it will be filled by [`switch_to`]
+    /// (for initial tasks) and [`init`] (for regular tasks) methods.
+    ///
+    /// [`init`]: TaskContext::init
+    /// [`switch_to`]: TaskContext::switch_to
     pub const fn new() -> Self {
         unsafe { core::mem::MaybeUninit::zeroed().assume_init() }
     }

modules/axhal/src/arch/riscv/context.rs

@@ -90,7 +90,13 @@ pub struct TaskContext {
 }
 
 impl TaskContext {
-    /// Creates a new default context for a new task.
+    /// Creates a dummy context for a new task.
+    ///
+    /// Note the context is not initialized, it will be filled by [`switch_to`]
+    /// (for initial tasks) and [`init`] (for regular tasks) methods.
+    ///
+    /// [`init`]: TaskContext::init
+    /// [`switch_to`]: TaskContext::switch_to
     pub const fn new() -> Self {
         unsafe { core::mem::MaybeUninit::zeroed().assume_init() }
     }

modules/axhal/src/arch/x86_64/context.rs

@@ -1,10 +1,9 @@
-use core::{arch::asm, fmt};
-use memory_addr::VirtAddr;
+#![allow(unused_imports)]
 
-#[cfg(all(feature = "irq", feature = "uspace"))]
+use core::{arch::asm, fmt};
+use memory_addr::{PhysAddr, VirtAddr};
 use x86_64::registers::rflags::RFlags;
 
-#[cfg(feature = "uspace")]
 use super::gdt::GdtStruct;
 
 /// Saved registers when a trap (interrupt or exception) occurs.
@@ -53,6 +52,11 @@ pub struct UspaceContext(TrapFrame);
 
 #[cfg(feature = "uspace")]
 impl UspaceContext {
+    /// Creates an empty context with all registers set to zero.
+    pub const fn empty() -> Self {
+        unsafe { core::mem::MaybeUninit::zeroed().assume_init() }
+    }
+
     /// Creates a new context with the given entry point, user stack pointer,
     /// and the argument.
     pub fn new(entry: usize, ustack_top: VirtAddr, arg0: usize) -> Self {
@@ -79,6 +83,26 @@ impl UspaceContext {
         Self(tf)
     }
 
+    /// Gets the instruction pointer.
+    pub const fn get_ip(&self) -> usize {
+        self.0.rip as _
+    }
+
+    /// Gets the stack pointer.
+    pub const fn get_sp(&self) -> usize {
+        self.0.rsp as _
+    }
+
+    /// Sets the stack pointer.
+    pub const fn set_sp(&mut self, rsp: usize) {
+        self.0.rsp = rsp as _;
+    }
+
+    /// Sets the return value register.
+    pub const fn set_ret_reg(&mut self, rax: usize) {
+        self.0.rax = rax as _;
+    }
+
     /// Enters user space.
     ///
     /// It restores the user registers and jumps to the user entry point
@@ -91,7 +115,7 @@ impl UspaceContext {
     /// This function is unsafe because it changes processor mode and the stack.
     pub unsafe fn enter_uspace(&self, kstack_top: VirtAddr) -> ! {
         super::disable_irqs();
-        super::tss_set_rsp0(kstack_top);
+        assert_eq!(super::tss_get_rsp0(), kstack_top);
         asm!("
             mov     rsp, {tf}
             pop     rax
@@ -218,15 +242,26 @@ pub struct TaskContext {
     /// Extended states, i.e., FP/SIMD states.
     #[cfg(feature = "fp_simd")]
     pub ext_state: ExtendedState,
+    /// The `CR3` register value, i.e., the page table root.
+    #[cfg(feature = "uspace")]
+    pub cr3: PhysAddr,
 }
 
 impl TaskContext {
-    /// Creates a new default context for a new task.
-    pub const fn new() -> Self {
+    /// Creates a dummy context for a new task.
+    ///
+    /// Note the context is not initialized, it will be filled by [`switch_to`]
+    /// (for initial tasks) and [`init`] (for regular tasks) methods.
+    ///
+    /// [`init`]: TaskContext::init
+    /// [`switch_to`]: TaskContext::switch_to
+    pub fn new() -> Self {
         Self {
             kstack_top: VirtAddr::from(0),
             rsp: 0,
             fs_base: 0,
+            #[cfg(feature = "uspace")]
+            cr3: crate::paging::kernel_page_table_root(),
             #[cfg(feature = "fp_simd")]
             ext_state: ExtendedState::default(),
         }
@@ -254,6 +289,17 @@ impl TaskContext {
         self.fs_base = tls_area.as_usize();
     }
 
+    /// Changes the page table root (`CR3` register for x86_64).
+    ///
+    /// If not set, the kernel page table root is used (obtained by
+    /// [`axhal::paging::kernel_page_table_root`][1]).
+    ///
+    /// [1]: crate::paging::kernel_page_table_root
+    #[cfg(feature = "uspace")]
+    pub fn set_page_table_root(&mut self, cr3: PhysAddr) {
+        self.cr3 = cr3;
+    }
+
     /// Switches to another task.
     ///
     /// It first saves the current task's context from CPU to this place, and then
@@ -265,9 +311,16 @@ impl TaskContext {
             next_ctx.ext_state.restore();
         }
         #[cfg(feature = "tls")]
-        {
+        unsafe {
             self.fs_base = super::read_thread_pointer();
-            unsafe { super::write_thread_pointer(next_ctx.fs_base) };
+            super::write_thread_pointer(next_ctx.fs_base);
+        }
+        #[cfg(feature = "uspace")]
+        unsafe {
+            super::tss_set_rsp0(next_ctx.kstack_top);
+            if next_ctx.cr3 != self.cr3 {
+                super::write_page_table_root(next_ctx.cr3);
+            }
         }
         unsafe { context_switch(&mut self.rsp, &next_ctx.rsp) }
     }

modules/axhal/src/arch/x86_64/gdt.rs

@@ -1,5 +1,5 @@
 use core::fmt;
-use spin::Once;
+use lazy_init::LazyInit;
 
 use x86_64::instructions::tables::{lgdt, load_tss};
 use x86_64::registers::segmentation::{Segment, SegmentSelector, CS};
@@ -9,10 +9,10 @@ use x86_64::{addr::VirtAddr, PrivilegeLevel};
 
 #[no_mangle]
 #[percpu::def_percpu]
-static TSS: Once<TaskStateSegment> = Once::new();
+static TSS: TaskStateSegment = TaskStateSegment::new();
 
 #[percpu::def_percpu]
-static GDT: Once<GdtStruct> = Once::new();
+static GDT: LazyInit<GdtStruct> = LazyInit::new();
 
 /// A wrapper of the Global Descriptor Table (GDT) with maximum 16 entries.
 #[repr(align(16))]
@@ -99,21 +99,25 @@ impl fmt::Debug for GdtStruct {
 /// current CPU.
 pub fn init_gdt() {
     unsafe {
-        let tss = TSS.current_ref_raw();
         let gdt = GDT.current_ref_raw();
-        let tss = tss.call_once(TaskStateSegment::new);
-        let gdt = gdt.call_once(|| GdtStruct::new(tss));
+        gdt.init_by(GdtStruct::new(TSS.current_ref_raw()));
         gdt.load();
         gdt.load_tss();
     }
 }
 
+/// Returns the stack pointer for privilege level 0 (RSP0) of the current TSS.
+pub fn tss_get_rsp0() -> memory_addr::VirtAddr {
+    let tss = unsafe { TSS.current_ref_raw() };
+    memory_addr::VirtAddr::from(tss.privilege_stack_table[0].as_u64() as usize)
+}
+
 /// Sets the stack pointer for privilege level 0 (RSP0) of the current TSS.
 ///
 /// # Safety
 ///
 /// Must be called after initialization and preemption is disabled.
 pub unsafe fn tss_set_rsp0(rsp0: memory_addr::VirtAddr) {
-    let tss = unsafe { TSS.current_ref_mut_raw().get_mut_unchecked() };
-    tss.privilege_stack_table[0] = x86_64::VirtAddr::new(rsp0.as_usize() as u64);
+    let tss = unsafe { TSS.current_ref_mut_raw() };
+    tss.privilege_stack_table[0] = VirtAddr::new_truncate(rsp0.as_usize() as u64);
 }

modules/axhal/src/arch/x86_64/mod.rs

@@ -15,7 +15,7 @@ use x86::{controlregs, msr, tlb};
 use x86_64::instructions::interrupts;
 
 pub use self::context::{ExtendedState, FxsaveArea, TaskContext, TrapFrame};
-pub use self::gdt::{init_gdt, tss_set_rsp0, GdtStruct};
+pub use self::gdt::{init_gdt, tss_get_rsp0, tss_set_rsp0, GdtStruct};
 pub use self::idt::{init_idt, IdtStruct};
 
 #[cfg(feature = "uspace")]

modules/axhal/src/lib.rs

@@ -28,6 +28,7 @@
 #![feature(asm_const)]
 #![feature(naked_functions)]
 #![feature(const_option)]
+#![feature(const_mut_refs)]
 #![feature(doc_auto_cfg)]
 
 #[allow(unused_imports)]

modules/axhal/src/paging.rs

@@ -2,6 +2,7 @@
 
 use axalloc::global_allocator;
 use page_table::PagingIf;
+use spin::Once;
 
 use crate::mem::{phys_to_virt, virt_to_phys, MemRegionFlags, PhysAddr, VirtAddr, PAGE_SIZE_4K};
 
@@ -64,3 +65,24 @@ cfg_if::cfg_if! {
         pub type PageTable = page_table::aarch64::A64PageTable<PagingIfImpl>;
     }
 }
+
+static KERNEL_PAGE_TABLE_ROOT: Once<PhysAddr> = Once::new();
+
+/// Saves the kernel page table information, which may be used on context
+/// switch.
+pub fn set_kernel_page_table(kern_pt: &PageTable) {
+    let root_paddr = kern_pt.root_paddr();
+    KERNEL_PAGE_TABLE_ROOT.call_once(|| root_paddr);
+    unsafe { crate::arch::write_page_table_root(root_paddr) };
+}
+
+/// Get the root physical address of the kernel page table.
+///
+/// # Panics
+///
+/// It must be called after [`set_kernel_page_table`], otherwise it will panic.
+pub fn kernel_page_table_root() -> PhysAddr {
+    *KERNEL_PAGE_TABLE_ROOT
+        .get()
+        .expect("kernel page table not initialized")
+}

modules/axruntime/src/lib.rs

@@ -251,8 +251,7 @@ fn remap_kernel_memory() -> Result<(), axhal::paging::PagingError> {
         }
         KERNEL_PAGE_TABLE.init_by(kernel_page_table);
     }
-
-    unsafe { axhal::arch::write_page_table_root(KERNEL_PAGE_TABLE.root_paddr()) };
+    axhal::paging::set_kernel_page_table(&KERNEL_PAGE_TABLE);
     Ok(())
 }
 

modules/axtask/src/api.rs

@@ -89,16 +89,21 @@ pub fn on_timer_tick() {
     RUN_QUEUE.lock().scheduler_timer_tick();
 }
 
+/// Add the given task to the run queue, returns the task reference.
+pub fn spawn_task(task: TaskInner) -> AxTaskRef {
+    let task_ref = task.into_arc();
+    RUN_QUEUE.lock().add_task(task_ref.clone());
+    task_ref
+}
+
 /// Spawns a new task with the given parameters.
 ///
 /// Returns the task reference.
 pub fn spawn_raw<F>(f: F, name: String, stack_size: usize) -> AxTaskRef
 where
     F: FnOnce() + Send + 'static,
 {
-    let task = TaskInner::new(f, name, stack_size);
-    RUN_QUEUE.lock().add_task(task.clone());
-    task
+    spawn_task(TaskInner::new(f, name, stack_size))
 }
 
 /// Spawns a new task with the default parameters.

modules/axtask/src/lib.rs

@@ -29,6 +29,8 @@
 #![feature(doc_cfg)]
 #![feature(doc_auto_cfg)]
 #![feature(linkage)]
+#![feature(const_mut_refs)]
+#![feature(const_unsafecell_get_mut)]
 
 #[cfg(test)]
 mod tests;

modules/axtask/src/run_queue.rs

@@ -24,7 +24,7 @@ pub(crate) struct AxRunQueue {
 
 impl AxRunQueue {
     pub fn new() -> SpinNoIrq<Self> {
-        let gc_task = TaskInner::new(gc_entry, "gc".into(), axconfig::TASK_STACK_SIZE);
+        let gc_task = TaskInner::new(gc_entry, "gc".into(), axconfig::TASK_STACK_SIZE).into_arc();
         let mut scheduler = Scheduler::new();
         scheduler.add_task(gc_task);
         SpinNoIrq::new(Self { scheduler })
@@ -214,21 +214,22 @@ fn gc_entry() {
 }
 
 pub(crate) fn init() {
-    // Put the subsequent execution into the `main` task.
+    // Create the `idle` task (not current task).
     const IDLE_TASK_STACK_SIZE: usize = 4096;
     let idle_task = TaskInner::new(|| crate::run_idle(), "idle".into(), IDLE_TASK_STACK_SIZE);
-    IDLE_TASK.with_current(|i| i.init_by(idle_task.clone()));
+    IDLE_TASK.with_current(|i| i.init_by(idle_task.into_arc()));
 
-    let main_task = TaskInner::new_init("main".into());
+    // Put the subsequent execution into the `main` task.
+    let main_task = TaskInner::new_init("main".into()).into_arc();
     main_task.set_state(TaskState::Running);
+    unsafe { CurrentTask::init_current(main_task) };
 
     RUN_QUEUE.init_by(AxRunQueue::new());
-    unsafe { CurrentTask::init_current(main_task) }
 }
 
 pub(crate) fn init_secondary() {
     // Put the subsequent execution into the `idle` task.
-    let idle_task = TaskInner::new_init("idle".into());
+    let idle_task = TaskInner::new_init("idle".into()).into_arc();
     idle_task.set_state(TaskState::Running);
     IDLE_TASK.with_current(|i| i.init_by(idle_task.clone()));
     unsafe { CurrentTask::init_current(idle_task) }