Treat GS as always pointing to TSS in kernel space.

src/arch/x86_64/gdt.rs

@@ -85,15 +85,39 @@ pub static mut GDT: [GdtEntry; 10] = [
     GdtEntry::new(0, 0, 0, 0),
 ];
 
+#[repr(packed)]
+pub struct TssWrapper {
+    base: TaskStateSegment,
+    _pad: u64,
+    _user_stack: u64,
+}
+impl core::ops::Deref for TssWrapper {
+    type Target = TaskStateSegment;
+
+    fn deref(&self) -> &Self::Target {
+        &self.base
+    }
+}
+impl core::ops::DerefMut for TssWrapper {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.base
+    }
+}
+
 #[thread_local]
-pub static mut TSS: TaskStateSegment = TaskStateSegment {
-    reserved: 0,
-    rsp: [0; 3],
-    reserved2: 0,
-    ist: [0; 7],
-    reserved3: 0,
-    reserved4: 0,
-    iomap_base: 0xFFFF
+pub static mut TSS: TssWrapper = TssWrapper {
+    base: TaskStateSegment {
+        reserved: 0,
+        rsp: [0; 3],
+        reserved2: 0,
+        ist: [0; 7],
+        reserved3: 0,
+        reserved4: 0,
+        iomap_base: 0xFFFF
+    },
+    _pad: 0_u64,
+    // Accessed only from assembly, at `gs:[0x70]`
+    _user_stack: 0_u64,
 };
 
 pub unsafe fn set_tcb(pid: usize) {
@@ -167,7 +191,11 @@ pub unsafe fn init_paging(tcb_offset: usize, stack_offset: usize) {
     segmentation::load_ds(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
     segmentation::load_es(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
     segmentation::load_fs(SegmentSelector::new(GDT_KERNEL_TLS as u16, Ring::Ring0));
+
     segmentation::load_gs(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
+    // Ensure that GS always points to the TSS segment in kernel space.
+    x86::msr::wrmsr(x86::msr::IA32_GS_BASE, &TSS as *const _ as usize as u64);
+
     segmentation::load_ss(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
 
     // Load the task register

src/arch/x86_64/interrupt/exception.rs

@@ -41,6 +41,8 @@ interrupt_stack!(debug, |stack| {
     }
 });
 
+// TODO: Give NMI and double faults a separate stack, as they can trigger a triple fault due to the
+// way swapgs is used.
 interrupt_stack!(non_maskable, |stack| {
     println!("Non-maskable interrupt");
     stack.dump();

src/arch/x86_64/interrupt/handler.rs

@@ -309,6 +309,17 @@ macro_rules! pop_fs {
     " };
 }
 
+macro_rules! swapgs_if_ring3 {
+    () => { "
+        // Check whether the last two bits RSP+8 (code segment) are equal to zero.
+        test BYTE PTR [rsp + 8], 0x03
+        // Skip the SWAPGS instruction if CS & 0b11 == 0b00.
+        jz 1f
+        swapgs
+        1:
+    " }
+}
+
 #[macro_export]
 macro_rules! interrupt_stack {
     ($name:ident, |$stack:ident| $code:block) => {
@@ -327,6 +338,7 @@ macro_rules! interrupt_stack {
 
             function!($name => {
                 // Backup all userspace registers to stack
+                swapgs_if_ring3!(),
                 "push rax\n",
                 push_scratch!(),
                 push_preserved!(),
@@ -347,6 +359,7 @@ macro_rules! interrupt_stack {
                 pop_preserved!(),
                 pop_scratch!(),
 
+                swapgs_if_ring3!(),
                 "iretq\n",
             });
         }
@@ -364,6 +377,7 @@ macro_rules! interrupt {
 
             function!($name => {
                 // Backup all userspace registers to stack
+                swapgs_if_ring3!(),
                 "push rax\n",
                 push_scratch!(),
                 push_fs!(),
@@ -381,6 +395,7 @@ macro_rules! interrupt {
                 pop_fs!(),
                 pop_scratch!(),
 
+                swapgs_if_ring3!(),
                 "iretq\n",
             });
         }
@@ -404,6 +419,7 @@ macro_rules! interrupt_error {
             }
 
             function!($name => {
+                swapgs_if_ring3!(),
                 // Move rax into code's place, put code in last instead (to be
                 // compatible with InterruptStack)
                 "xchg [rsp], rax\n",
@@ -434,6 +450,7 @@ macro_rules! interrupt_error {
                 pop_preserved!(),
                 pop_scratch!(),
 
+                swapgs_if_ring3!(),
                 "iretq\n",
             });
         }

src/arch/x86_64/interrupt/syscall.rs

@@ -16,13 +16,16 @@ pub unsafe fn init() {
     // The base selector of the three consecutive segments (of which two are used) for user code
     // and user data. It points to a 32-bit code segment, which must be followed by a data segment
     // (stack), and a 64-bit code segment.
-    let sysret_cs_ss_base = ((gdt::GDT_USER_CODE32_UNUSED as u16) << 3) | u16::from(gdt::GDT_A_RING_3);
+    let sysret_cs_ss_base = ((gdt::GDT_USER_CODE32_UNUSED as u16) << 3) | 3;
     let star_high = u32::from(syscall_cs_ss_base) | (u32::from(sysret_cs_ss_base) << 16);
 
     msr::wrmsr(msr::IA32_STAR, u64::from(star_high) << 32);
     msr::wrmsr(msr::IA32_LSTAR, syscall_instruction as u64);
     msr::wrmsr(msr::IA32_FMASK, 0x0300); // Clear trap flag and interrupt enable
-    msr::wrmsr(msr::IA32_KERNEL_GSBASE, &gdt::TSS as *const _ as u64);
+
+    // Inside kernel space, GS should _always_ point to the TSS. When leaving userspace, `swapgs`
+    // is called again, making the userspace GS always point to user data.
+    msr::wrmsr(msr::IA32_KERNEL_GSBASE, 0);
 
     let efer = msr::rdmsr(msr::IA32_EFER);
     msr::wrmsr(msr::IA32_EFER, efer | 1);
@@ -63,10 +66,10 @@ function!(syscall_instruction => {
     // Yes, this is magic. No, you don't need to understand
     "
         swapgs                    // Set gs segment to TSS
-        mov gs:[28], rsp          // Save userspace stack pointer
+        mov gs:[0x70], rsp        // Save userspace stack pointer
         mov rsp, gs:[4]           // Load kernel stack pointer
         push WORD PTR 5 * 8 + 3   // Push fake userspace SS (resembling iret frame)
-        push QWORD PTR gs:[28]    // Push userspace rsp
+        push QWORD PTR gs:[0x70]  // Push userspace rsp
         push r11                  // Push rflags
         push WORD PTR 6 * 8 + 3   // Push fake userspace CS (resembling iret frame)
         push rcx                  // Push userspace return pointer
@@ -96,11 +99,12 @@ function!(syscall_instruction => {
     // Return
     "
         pop rcx                 // Pop userspace return pointer
-        add rsp, 2
+        add rsp, 2              // Pop CS
         pop r11                 // Pop rflags
-        add rsp, 10             // Pop SS and rsp
-        mov rsp, gs:[28]        // Restore userspace stack pointer
-        swapgs                  // Restore gs from TSS to kernel data
+        pop QWORD PTR gs:[0x70] // Pop userspace stack pointer
+        add rsp, 2              // Pop SS
+        mov rsp, gs:[0x70]      // Restore userspace stack pointer
+        swapgs                  // Restore gs from TSS to user data
         sysretq                 // Return into userspace; RCX=>RIP,R11=>RFLAGS
     ",
 });

src/arch/x86_64/start.rs

@@ -243,19 +243,6 @@ pub unsafe extern fn kstart_ap(args_ptr: *const KernelArgsAp) -> ! {
 #[inline(never)]
 // TODO: AbiCompatBool
 pub unsafe extern "C" fn usermode(_ip: usize, _sp: usize, _arg: usize, _singlestep: u32) -> ! {
-    /*asm!("push r10
-          push r11
-          push r12
-          push r13
-          push r14
-          push r15",
-         in("r10") (gdt::GDT_USER_DATA << 3 | 3), // Data segment
-         in("r11") sp, // Stack pointer
-         in("r12") flags, // Flags
-         in("r13") (gdt::GDT_USER_CODE << 3 | 3), // Code segment
-         in("r14") ip, // IP
-         in("r15") arg, // Argument
-    );*/
     // rdi, rsi, rdx, rcx
     asm!(
         "
@@ -265,25 +252,34 @@ pub unsafe extern "C" fn usermode(_ip: usize, _sp: usize, _arg: usize, _singlest
             or rbx, {flag_singlestep}
 
             .after_singlestep_branch:
-            mov r12, rdi
-            mov r13, rsi
-            mov rdi, rdx
+
+            // save `ip` (rdi), `sp` (rsi), and `arg` (rdx) in callee-preserved registers, so that
+            // they are not modified by `pti_unmap`
+
+            mov r13, rdi
+            mov r14, rsi
+            mov r15, rdx
             call {pti_unmap}
 
             // Go to usermode
-            mov r14, {user_data_seg_selector}
-            mov r15, {user_tls_seg_selector}
-            mov ds, r14d
-            mov es, r14d
-            mov fs, r15d
-            mov gs, r14d
+            mov r8, {user_data_seg_selector}
+            mov r9, {user_tls_seg_selector}
+            mov ds, r8d
+            mov es, r8d
+            mov fs, r9d
+            // Exchange the old kernel GS (pointing to TSS) and kernel data
+            swapgs
+            // Replace kernel data segment with user data segment
+            mov gs, r8d
 
             // Target RFLAGS
             mov r11, rbx
             // Target instruction pointer
-            mov rcx, r12
+            mov rcx, r13
             // Target stack pointer
-            mov rsp, r13
+            mov rsp, r14
+            // Target argument
+            mov rdi, r13
 
             xor rax, rax
             xor rbx, rbx