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: super::read_page_table_root(),
             #[cfg(feature = "fp_simd")]
             ext_state: ExtendedState::default(),
         }
@@ -254,6 +289,12 @@ impl TaskContext {
         self.fs_base = tls_area.as_usize();
     }
 
+    /// Sets the page table root (`CR3` register for x86_64).
+    #[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 +306,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

@@ -108,12 +108,18 @@ pub fn init_gdt() {
     }
 }
 
+/// 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_mut_raw().get_mut_unchecked() };
+    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);
+    tss.privilege_stack_table[0] = VirtAddr::new(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/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) }

modules/axtask/src/task.rs

@@ -154,7 +154,7 @@ impl TaskInner {
     }
 
     /// Create a new task with the given entry function and stack size.
-    pub(crate) fn new<F>(entry: F, name: String, stack_size: usize) -> AxTaskRef
+    pub fn new<F>(entry: F, name: String, stack_size: usize) -> Self
     where
         F: FnOnce() + Send + 'static,
     {
@@ -168,12 +168,12 @@ impl TaskInner {
         let tls = VirtAddr::from(0);
 
         t.entry = Some(Box::into_raw(Box::new(entry)));
-        t.ctx.get_mut().init(task_entry as usize, kstack.top(), tls);
+        t.ctx_mut().init(task_entry as usize, kstack.top(), tls);
         t.kstack = Some(kstack);
         if t.name == "idle" {
             t.is_idle = true;
         }
-        Arc::new(AxTask::new(t))
+        t
     }
 
     /// Creates an "init task" using the current CPU states, to use as the
@@ -184,13 +184,17 @@ impl TaskInner {
     ///
     /// And there is no need to set the `entry`, `kstack` or `tls` fields, as
     /// they will be filled automatically when the task is switches out.
-    pub(crate) fn new_init(name: String) -> AxTaskRef {
+    pub(crate) fn new_init(name: String) -> Self {
         let mut t = Self::new_common(TaskId::new(), name);
         t.is_init = true;
         if t.name == "idle" {
             t.is_idle = true;
         }
-        Arc::new(AxTask::new(t))
+        t
+    }
+
+    pub(crate) fn into_arc(self) -> AxTaskRef {
+        Arc::new(AxTask::new(self))
     }
 
     #[inline]
@@ -297,6 +301,21 @@ impl TaskInner {
     pub(crate) const unsafe fn ctx_mut_ptr(&self) -> *mut TaskContext {
         self.ctx.get()
     }
+
+    /// Returns a mutable reference to the task context.
+    #[inline]
+    pub const fn ctx_mut(&mut self) -> &mut TaskContext {
+        self.ctx.get_mut()
+    }
+
+    /// Returns the top address of the kernel stack.
+    #[inline]
+    pub const fn kernel_stack_top(&self) -> Option<VirtAddr> {
+        match &self.kstack {
+            Some(s) => Some(s.top()),
+            None => None,
+        }
+    }
 }
 
 impl fmt::Debug for TaskInner {
@@ -343,6 +362,8 @@ impl Drop for TaskStack {
 use core::mem::ManuallyDrop;
 
 /// A wrapper of [`AxTaskRef`] as the current task.
+///
+/// It won't change the reference count of the task when created or dropped.
 pub struct CurrentTask(ManuallyDrop<AxTaskRef>);
 
 impl CurrentTask {