set kernel stack in tss

Author: 0x6D70

kernel/src/arch/amd64/gdt.rs

@@ -122,3 +122,23 @@ pub fn init_ap(ap_id: usize) {
         load_tss(tss_selector);
     }
 }
+
+pub fn set_tss_kernel_stack(stack_top: u64) {
+    unsafe {
+        // For BSP
+        if TSS.is_completed() {
+            let tss = &*TSS as *const TaskStateSegment as *mut TaskStateSegment;
+            (*tss).privilege_stack_table[0] = VirtAddr::new(stack_top);
+            return;
+        }
+
+        // For APs
+        let gdt_ap = GDT_AP.lock();
+        let ap_id = crate::cpu::current_cpu().cpu_id;
+        
+        if let Some(gdt) = gdt_ap.iter().find(|g| g.ap_id == ap_id) {
+            let tss = gdt.tss as *const TaskStateSegment as *mut TaskStateSegment;
+            (*tss).privilege_stack_table[0] = VirtAddr::new(stack_top);
+        }
+    }
+}

kernel/src/cpu.rs

@@ -122,7 +122,7 @@ pub struct Cpu {
     user_space_stack: usize,
     pub kernel_stack: Cell<usize>,
 
-    cpu_id: usize,
+    pub cpu_id: usize,
     pub lapic_id: u32,
     pub run_queue: Spinlock<VecDeque<Arc<Thread>>>,
     pub wait_queue: Spinlock<VecDeque<Arc<Thread>>>,

kernel/src/sched/process.rs

@@ -8,7 +8,7 @@ use crate::fs::file::File;
 use crate::fs::vnode::VNode;
 use crate::mem::frame::FRAME_ALLOCATOR;
 use crate::mem::vm::{protflags_from_ptflags, Vm};
-use crate::mem::{HIGHER_HALF_OFFSET, KERNEL_THREAD_STACK_TOP, PROCESS_START, STACK_SIZE};
+use crate::mem::{HIGHER_HALF_OFFSET, KERNEL_THREAD_STACK_TOP, STACK_SIZE};
 use crate::VFS;
 use alloc::collections::BTreeMap;
 use alloc::sync::Arc;
@@ -122,7 +122,7 @@ impl Process {
 
         for ph in elf.segments().expect("Failed to get program headers") {
             if ph.p_type == elf::abi::PT_LOAD {
-                let start = ph.p_vaddr + PROCESS_START;
+                let start = ph.p_vaddr;
                 let end = start + ph.p_memsz;
 
                 let page_start = align_down(start, Size4KiB::SIZE);
@@ -172,7 +172,7 @@ impl Process {
             }
         }
 
-        VirtAddr::new(elf.ehdr.e_entry + PROCESS_START)
+        VirtAddr::new(elf.ehdr.e_entry)
     }
 
     pub fn next_tid(&mut self) -> usize {

kernel/src/sched/scheduler.rs

@@ -1,4 +1,4 @@
-use crate::arch::amd64::gdt::GDT_BSP;
+use crate::arch::amd64::gdt::{set_tss_kernel_stack, GDT_BSP};
 use crate::arch::amd64::switch_context;
 use crate::cpu::current_cpu;
 use crate::timer::timer_event::TimerEvent;
@@ -83,9 +83,13 @@ pub fn switch_threads(old: Arc<Thread>, new: Arc<Thread>) {
 
             DS::set_reg(GDT_BSP.1.user_data_selector);
 
+            let kernel_stack_top = new.kernel_stack.as_ref().unwrap().kernel_stack_top;
+            
             current_cpu()
                 .kernel_stack
-                .set(new.kernel_stack.as_ref().unwrap().kernel_stack_top);
+                .set(kernel_stack_top);
+            
+            set_tss_kernel_stack(kernel_stack_top as u64);
         }
     }