From 4ced24e8d1e5347aa42b945a3e49ab669f07b752 Mon Sep 17 00:00:00 2001
From: Jeremy Soller <jackpot51@gmail.com>
Date: Mon, 22 Aug 2016 14:53:03 -0600
Subject: [PATCH] Setup TSS, using thread local storage

---
 arch/x86_64/src/gdt.rs   | 10 +++++++---
 arch/x86_64/src/start.rs |  4 ++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/arch/x86_64/src/gdt.rs b/arch/x86_64/src/gdt.rs
index 1dc94e8..7ef3450 100644
--- a/arch/x86_64/src/gdt.rs
+++ b/arch/x86_64/src/gdt.rs
@@ -89,7 +89,7 @@ pub static mut TSS: TaskStateSegment = TaskStateSegment {
 };
 
 /// Initialize GDT on the BSP
-pub unsafe fn init(tcb_offset: usize) {
+pub unsafe fn init(tcb_offset: usize, stack_offset: usize) {
     // Setup the initial GDT with TLS, so we can setup the TLS GDT (a little confusing)
     // This means that each CPU will have its own GDT, but we only need to define it once as a thread local
     INIT_GDTR.limit = (INIT_GDT.len() * mem::size_of::<GdtEntry>() - 1) as u16;
@@ -99,11 +99,11 @@ pub unsafe fn init(tcb_offset: usize) {
     INIT_GDT[GDT_KERNEL_TLS].set_offset(tcb_offset as u32);
 
     // Run the AP GDT initialization, which does the rest
-    init_ap(tcb_offset);
+    init_ap(tcb_offset, stack_offset);
 }
 
 /// Initialize GDT for an AP
-pub unsafe fn init_ap(tcb_offset: usize) {
+pub unsafe fn init_ap(tcb_offset: usize, stack_offset: usize) {
     // Load the initial GDT, before we have access to thread locals
     dtables::lgdt(&INIT_GDTR);
 
@@ -118,12 +118,16 @@ pub unsafe fn init_ap(tcb_offset: usize) {
     GDTR.limit = (GDT.len() * mem::size_of::<GdtEntry>() - 1) as u16;
     GDTR.base = GDT.as_ptr() as u64;
 
+    // Set the TLS segment to the offset of the Thread Control Block
     GDT[GDT_KERNEL_TLS].set_offset(tcb_offset as u32);
 
     // We can now access our TSS, which is a thread local
     GDT[GDT_TSS].set_offset(&TSS as *const _ as u32);
     GDT[GDT_TSS].set_limit(mem::size_of::<TaskStateSegment>() as u32);
 
+    // Set the stack pointer when coming back from userspace
+    TSS.rsp[0] = stack_offset as u64;
+
     // Load the new GDT, which is correctly located in thread local storage
     dtables::lgdt(&GDTR);
 
diff --git a/arch/x86_64/src/start.rs b/arch/x86_64/src/start.rs
index 34481ee..08b2113 100644
--- a/arch/x86_64/src/start.rs
+++ b/arch/x86_64/src/start.rs
@@ -75,7 +75,7 @@ pub unsafe extern fn kstart() -> ! {
         let mut active_table = paging::init(stack_start, stack_end);
 
         // Set up GDT
-        gdt::init((&__tbss_end as *const u8 as *const usize).offset(-1) as usize);
+        gdt::init((&__tbss_end as *const u8 as *const usize).offset(-1) as usize, stack_end);
 
         // Set up IDT
         idt::init();
@@ -146,7 +146,7 @@ pub unsafe extern fn kstart_ap(stack_start: usize, stack_end: usize) -> ! {
         let mut active_table = paging::init(stack_start, stack_end);
 
         // Set up GDT for AP
-        gdt::init_ap((&__tbss_end as *const u8 as *const usize).offset(-1) as usize);
+        gdt::init_ap((&__tbss_end as *const u8 as *const usize).offset(-1) as usize, stack_end);
 
         // Set up IDT for AP
         idt::init();