Switches to kernel stack on syscall (#1074)

kernel/src/aarch64.rs

@@ -1,3 +1,5 @@
-pub unsafe fn setup_main_cpu() {
+use crate::context::ContextArgs;
+
+pub unsafe fn setup_main_cpu() -> ContextArgs {
     todo!()
 }

kernel/src/context/aarch64.rs

@@ -1,6 +1,9 @@
 use super::Base;
 use crate::proc::Thread;
 
+/// Contains data passed from CPU setup function for context activation.
+pub struct ContextArgs {}
+
 /// Extended [Base] for AArch64.
 #[repr(C)]
 pub(super) struct Context {
@@ -8,7 +11,7 @@ pub(super) struct Context {
 }
 
 impl Context {
-    pub fn new(base: Base) -> Self {
+    pub fn new(base: Base, args: ContextArgs) -> Self {
         Self { base }
     }
 

kernel/src/context/mod.rs

@@ -25,14 +25,23 @@ mod local;
 /// - `cpu` must be unique and valid.
 /// - `pmgr` must be the same for all context.
 /// - `f` must not get inlined.
-pub unsafe fn run_with_context(cpu: usize, td: Arc<Thread>, pmgr: Arc<ProcMgr>, f: fn() -> !) -> ! {
+pub unsafe fn run_with_context(
+    cpu: usize,
+    td: Arc<Thread>,
+    pmgr: Arc<ProcMgr>,
+    args: ContextArgs,
+    f: fn() -> !,
+) -> ! {
     // We use a different mechanism here. The PS4 put all of pcpu at a global level but we put it on
     // each CPU stack instead.
-    let mut cx = Context::new(Base {
-        cpu,
-        thread: Arc::into_raw(td),
-        pmgr: Arc::into_raw(pmgr),
-    });
+    let mut cx = Context::new(
+        Base {
+            cpu,
+            thread: Arc::into_raw(td),
+            pmgr: Arc::into_raw(pmgr),
+        },
+        args,
+    );
 
     cx.activate();
     f();

kernel/src/context/x86_64.rs

@@ -3,20 +3,34 @@ use crate::arch::wrmsr;
 use core::arch::asm;
 use core::mem::offset_of;
 
+pub const fn current_trap_rsp_offset() -> usize {
+    offset_of!(Context, trap_rsp)
+}
+
 pub const fn current_user_rsp_offset() -> usize {
     offset_of!(Context, user_rsp)
 }
 
+/// Contains data passed from CPU setup function for context activation.
+pub struct ContextArgs {
+    pub trap_rsp: *mut u8,
+}
+
 /// Extended [Base] for x86-64.
 #[repr(C)]
 pub(super) struct Context {
-    base: Base,      // Must be first field.
-    user_rsp: usize, // pc_scratch_rsp
+    base: Base,        // Must be first field.
+    trap_rsp: *mut u8, // pc_rsp0
+    user_rsp: usize,   // pc_scratch_rsp
 }
 
 impl Context {
-    pub fn new(base: Base) -> Self {
-        Self { base, user_rsp: 0 }
+    pub fn new(base: Base, args: ContextArgs) -> Self {
+        Self {
+            base,
+            trap_rsp: args.trap_rsp,
+            user_rsp: 0,
+        }
     }
 
     /// Set kernel `GS` segment register to `cx`.

kernel/src/main.rs

@@ -47,7 +47,9 @@ unsafe extern "C" fn _start(env: &'static BootEnv, conf: &'static Config) -> ! {
 
     info!("Starting Obliteration Kernel.");
 
-    self::arch::setup_main_cpu();
+    // Setup the CPU after the first print to let the bootloader developer know (some of) their code
+    // are working.
+    let cx = self::arch::setup_main_cpu();
 
     // Setup thread0 to represent this thread.
     let thread0 = Thread::new_bare();
@@ -58,7 +60,7 @@ unsafe extern "C" fn _start(env: &'static BootEnv, conf: &'static Config) -> ! {
     // Activate CPU context.
     let thread0 = Arc::new(thread0);
 
-    self::context::run_with_context(0, thread0, pmgr, main);
+    self::context::run_with_context(0, thread0, pmgr, cx, main);
 }
 
 #[inline(never)] // See self::context::run_with_context docs.
@@ -88,13 +90,14 @@ fn main() -> ! {
     }
 }
 
-/// # Interupt safety
-/// This function is interupt safe.
+/// # Context safety
+/// This function does not require a CPU context.
+///
+/// # Interrupt safety
+/// This function is interrupt safe.
 #[allow(dead_code)]
 #[cfg_attr(target_os = "none", panic_handler)]
 fn panic(i: &PanicInfo) -> ! {
-    // This function is not allowed to access the CPU context due to it can be called before the
-    // context has been activated.
     let (file, line) = match i.location() {
         Some(v) => (v.file(), v.line()),
         None => ("unknown", 0),

kernel/src/x86_64.rs

@@ -1,4 +1,4 @@
-use crate::context::current_user_rsp_offset;
+use crate::context::{current_trap_rsp_offset, current_user_rsp_offset, ContextArgs};
 use crate::trap::interrupt_handler;
 use bitfield_struct::bitfield;
 use core::arch::{asm, global_asm};
@@ -12,7 +12,7 @@ pub const GDT_USER_CS32: SegmentSelector = SegmentSelector::new().with_si(5).wit
 
 /// # Safety
 /// This function can be called only once and must be called by main CPU entry point.
-pub unsafe fn setup_main_cpu() {
+pub unsafe fn setup_main_cpu() -> ContextArgs {
     // Setup GDT.
     static mut GDT: [SegmentDescriptor; 10] = [
         // Null descriptor.
@@ -148,6 +148,10 @@ pub unsafe fn setup_main_cpu() {
         .unwrap();
 
     wrmsr(0xC0000080, efer);
+
+    ContextArgs {
+        trap_rsp: TSS.rsp0 as _,
+    }
 }
 
 pub unsafe fn wrmsr(reg: u32, val: usize) {
@@ -196,9 +200,11 @@ global_asm!(
 global_asm!(
     "syscall_entry64:",
     "swapgs",
-    "mov gs:[{user_rsp}], rsp",
+    "mov gs:[{user_rsp}], rsp", // Save user RSP.
+    "mov rsp, gs:[{trap_rsp}]",
     "ud2",
-    user_rsp = const current_user_rsp_offset()
+    user_rsp = const current_user_rsp_offset(),
+    trap_rsp = const current_trap_rsp_offset()
 );
 
 // See Xfast_syscall32 on the PS4 for a reference.