From ffb7e194f95c652c5f2e2038c8f0a5820db140dc Mon Sep 17 00:00:00 2001
From: WangRunji <wangrunji0408@163.com>
Date: Sun, 20 May 2018 22:12:18 +0800
Subject: [PATCH] Fix set TSS in SMP

---
 src/arch/x86_64/driver/apic/lapic.rs |  3 ++
 src/arch/x86_64/gdt.rs               | 46 +++++++++++++++++++---------
 src/arch/x86_64/interrupt/handler.rs |  4 +--
 src/lib.rs                           |  3 +-
 4 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/src/arch/x86_64/driver/apic/lapic.rs b/src/arch/x86_64/driver/apic/lapic.rs
index 3692bf6..80e2b54 100644
--- a/src/arch/x86_64/driver/apic/lapic.rs
+++ b/src/arch/x86_64/driver/apic/lapic.rs
@@ -35,6 +35,9 @@ pub fn start_ap(apicid: u8, addr: u32) {
 
 pub fn lapic_id() -> u8 {
 	unsafe{
+        if lapic.is_null() {
+            return 0;
+        }
 		(*(lapic as *const u32).offset(0x0020/4) >> 24) as u8
 	}
 }
\ No newline at end of file
diff --git a/src/arch/x86_64/gdt.rs b/src/arch/x86_64/gdt.rs
index b93f969..daca7ad 100644
--- a/src/arch/x86_64/gdt.rs
+++ b/src/arch/x86_64/gdt.rs
@@ -3,9 +3,10 @@ use core::fmt::Debug;
 use x86_64::structures::tss::TaskStateSegment;
 use x86_64::structures::gdt::SegmentSelector;
 use x86_64::{PrivilegeLevel, VirtualAddress};
-use spin::Once;
+use spin::{Once, Mutex, MutexGuard};
 use alloc::boxed::Box;
-use core::ptr::Unique;
+use arch::driver::apic::lapic_id;
+use consts::MAX_CPU_NUM;
 
 /// Alloc TSS & GDT at kernel heap, then init and load it.
 /// The double fault stack will be allocated at kernel heap too.
@@ -26,8 +27,7 @@ pub fn init() {
 
         tss
     });
-    unsafe{ TSS_PTR = Unique::new_unchecked(Box::into_raw(tss)); }
-    let tss = unsafe{ TSS_PTR.as_ref() };
+    let tss = Box::into_raw(tss);
 
     let gdt = Box::new({
         let mut gdt = Gdt::new();
@@ -37,7 +37,7 @@ pub fn init() {
         gdt.add_entry(UDATA);
         gdt.add_entry(UCODE32);
         gdt.add_entry(UDATA32);
-        gdt.add_entry(Descriptor::tss_segment(&tss));
+        gdt.add_entry(Descriptor::tss_segment(unsafe { &*tss }));
         gdt
     });
     let gdt = unsafe{ &*Box::into_raw(gdt) };
@@ -49,19 +49,35 @@ pub fn init() {
         // load TSS
         load_tss(TSS_SELECTOR);
     }
+
+    CPUS[lapic_id() as usize].call_once(||
+        Mutex::new(Cpu { gdt, tss: unsafe { &mut *tss } }));
 }
 
-// TODO: more elegant?
-static mut TSS_PTR: Unique<TaskStateSegment> = unsafe{ Unique::new_unchecked(0 as *mut _) };
+static CPUS: [Once<Mutex<Cpu>>; MAX_CPU_NUM] = [
+    // TODO: More elegant ?
+    Once::new(), Once::new(), Once::new(), Once::new(),
+    Once::new(), Once::new(), Once::new(), Once::new(),
+];
 
-/// 设置从Ring3跳到Ring0时，自动切换栈的地址
-///
-/// 每次进入用户态前，都要调用此函数，才能保证正确返回内核态
-pub fn set_ring0_rsp(rsp: usize) {
-    trace!("gdt.set_ring0_rsp: {:#x}", rsp);
-    unsafe {
-        TSS_PTR.as_mut().privilege_stack_table[0] = VirtualAddress(rsp);
-        trace!("TSS:\n{:?}", TSS_PTR.as_ref());
+pub struct Cpu {
+    gdt: &'static Gdt,
+    tss: &'static mut TaskStateSegment,
+}
+
+impl Cpu {
+    pub fn current() -> MutexGuard<'static, Cpu> {
+        CPUS[lapic_id() as usize].try().unwrap().lock()
+    }
+
+    /// 设置从Ring3跳到Ring0时，自动切换栈的地址
+    ///
+    /// 每次进入用户态前，都要调用此函数，才能保证正确返回内核态
+    pub fn set_ring0_rsp(&mut self, rsp: usize) {
+        trace!("gdt.set_ring0_rsp: {:#x}", rsp);
+        unsafe {
+            self.tss.privilege_stack_table[0] = VirtualAddress(rsp);
+        }
     }
 }
 
diff --git a/src/arch/x86_64/interrupt/handler.rs b/src/arch/x86_64/interrupt/handler.rs
index fc76ab2..cd83f8e 100644
--- a/src/arch/x86_64/interrupt/handler.rs
+++ b/src/arch/x86_64/interrupt/handler.rs
@@ -126,10 +126,10 @@ pub extern fn rust_trap(tf: &mut TrapFrame) -> usize {
 }
 
 fn set_return_rsp(tf: &TrapFrame) {
-    use arch::gdt;
+    use arch::gdt::Cpu;
     use core::mem::size_of;
     if tf.cs & 0x3 == 3 {
-        gdt::set_ring0_rsp(tf as *const _ as usize + size_of::<TrapFrame>());
+        Cpu::current().set_ring0_rsp(tf as *const _ as usize + size_of::<TrapFrame>());
     }
 }
 
diff --git a/src/lib.rs b/src/lib.rs
index 1b3d5f9..b190b9f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -81,8 +81,7 @@ pub extern "C" fn rust_main(multiboot_information_address: usize) -> ! {
         |addr: usize| memory_controller.map_page_identity(addr));
     // memory_controller.print_page_table();
 
-    // FIXME: 开启SMP后，导致switch_to_user中设置rsp无效
-//    arch::smp::start_other_cores(&acpi, &mut memory_controller);
+    arch::smp::start_other_cores(&acpi, &mut memory_controller);
     process::init(memory_controller);
 
     fs::load_sfs();