SMP fixes

deny.toml

@@ -1,6 +1,6 @@
 [licenses]
-allow = ["Apache-2.0", "BSD-3-Clause", "ISC", "MIT", "MPL-2.0", "Unlicense"]
+allow = ["Apache-2.0", "BSD-3-Clause", "ISC", "MIT", "MPL-2.0", "Unlicense", "Zlib"]
 exceptions = [
   {allow = ["BSD-3-Clause-acpica"], crate = "acpica-rs"},
-  {allow = ["Unicode-DFS-2016"], crate = "unicode-ident"}
+  {allow = ["Unicode-3.0"], crate = "unicode-ident"}
 ]

kernel/src/cpu/idt.rs

@@ -561,7 +561,7 @@ pub fn setup_idt() {
     set_address(&mut idt.entries[249], isr249);
     set_address(&mut idt.entries[250], isr250);
     set_address(&mut idt.entries[251], isr251);
-    set_address(&mut idt.entries[252], isr252);
+    set_address(&mut idt.entries[252], isr252); // Invalidate TLB
     set_address(&mut idt.entries[253], isr253); // SCI
     set_address(&mut idt.entries[254], isr254); // Scheduler
     set_address(&mut idt.entries[255], isr255); // Spurious interrupt

kernel/src/cpu/isr.rs

@@ -1,9 +1,16 @@
-use core::{arch::global_asm, sync::atomic::Ordering};
+use core::{
+    arch::{asm, global_asm},
+    sync::atomic::Ordering,
+};
 
 use log::warn;
 use spin::Mutex;
 
-use crate::sys::{lapic::send_eoi, scheduler::scheduler, sync::IN_ISR_HANDLER};
+use crate::sys::{
+    lapic::{get_current_lapic_id, send_eoi},
+    scheduler::scheduler,
+    sync::{IN_ISR_HANDLER, LOCKS_HELD},
+};
 
 global_asm!(include_str!("isr.s"), options(att_syntax));
 
@@ -69,6 +76,8 @@ const EXCEPTIONS: [&str; 32] = [
     "Reserved",
 ];
 
+pub const ISR_INVALIDATE_TLB: u64 = 252;
+pub const ISR_SCHEDULER: u64 = 254;
 pub static ISR_HANDLERS: Mutex<[Option<fn()>; 256]> = Mutex::new([None; 256]);
 
 #[no_mangle]
@@ -85,10 +94,19 @@ extern "C" fn isr_handler(state: &mut ISRState) {
     if 48 <= state.isr && state.isr <= 254 {
         send_eoi(); // APIC interrupt
     }
-    IN_ISR_HANDLER.store(true, Ordering::SeqCst);
-    if state.isr == 254 {
+    if state.isr == ISR_INVALIDATE_TLB {
+        // TODO: Only invalidate the entries that changed
+        unsafe {
+            asm!("mov rax, cr3; mov cr3, rax", out("rax") _);
+        }
+        return;
+    }
+    let lapic_id = get_current_lapic_id();
+    IN_ISR_HANDLER[lapic_id].store(true, Ordering::SeqCst);
+    if state.isr == ISR_SCHEDULER {
         scheduler(state);
-        IN_ISR_HANDLER.store(false, Ordering::SeqCst);
+        assert_eq!(LOCKS_HELD[lapic_id].load(Ordering::SeqCst), 0);
+        IN_ISR_HANDLER[lapic_id].store(false, Ordering::SeqCst);
         return;
     }
     let handler;
@@ -100,5 +118,6 @@ extern "C" fn isr_handler(state: &mut ISRState) {
     } else {
         warn!("Unhandled interrupt: {}", state.isr);
     }
-    IN_ISR_HANDLER.store(false, Ordering::SeqCst);
+    assert_eq!(LOCKS_HELD[lapic_id].load(Ordering::SeqCst), 0);
+    IN_ISR_HANDLER[lapic_id].store(false, Ordering::SeqCst);
 }

kernel/src/cpu/paging.rs

@@ -13,6 +13,8 @@ use kernel_common::{
 use log::info;
 use spin::Mutex;
 
+use crate::sys::lapic::smp_invalidate_tlb;
+
 extern "C" {
     static _text_start: u8;
     static _text_end: u8;
@@ -42,10 +44,10 @@ fn _get_free_frame() -> u64 {
     panic!("No free memory left");
 }
 fn invlpg(addr: u64) {
-    // TODO: Broadcast to all cores
     unsafe {
         asm!("invlpg [{}]", in(reg) addr);
     }
+    smp_invalidate_tlb();
 }
 pub fn virt_to_phys(virt: u64) -> u64 {
     if !PAGING_ACTIVE.load(Ordering::SeqCst) {

kernel/src/cpu/smp.rs

@@ -1,3 +1,5 @@
+use core::arch::asm;
+
 use kernel_common::instructions::sti;
 
 use crate::{
@@ -15,9 +17,17 @@ extern "C" fn ap_main() -> ! {
     setup_lapic(0);
     // TODO: Also calibrate other cores
     setup_lapic_timer(false);
+    set_cpu_flags();
     unsafe {
         sti();
     }
     yield_task();
     panic!("Yielding to idle task failed");
 }
+
+pub fn set_cpu_flags() {
+    unsafe {
+        // SMAP and SMEP
+        asm!("mov rax, cr4; bts rax, 20; bts rax, 21; mov cr4, rax", out("rax") _);
+    }
+}

kernel/src/main.rs

@@ -15,7 +15,7 @@ use acpi::AcpiTables;
 use acpica_rs::{AcpiEnableSubsystem, AcpiInitializeObjects, AcpiInitializeSubsystem, AcpiInitializeTables, AcpiLoadTables, ACPI_FULL_INITIALIZATION};
 use alloc::format;
 use buddy_system_allocator::LockedHeap;
-use cpu::{gdt::setup_gdt, idt::setup_idt, paging::setup_paging};
+use cpu::{gdt::setup_gdt, idt::setup_idt, paging::setup_paging, smp::set_cpu_flags};
 use kernel_common::{
     instructions::{cli, hlt},
     loader_struct::{FramebufferInfo, LoaderStruct, LoaderStructMemoryRegion, LOADER_STRUCT_MAGIC},
@@ -29,7 +29,7 @@ use sys::{
     acpica_osl::AE_OK,
     early_acpi::EarlyACPIHandler,
     hpet::setup_hpet,
-    lapic::{setup_lapic_timer, start_aps},
+    lapic::{get_current_lapic_id, setup_lapic_timer, start_aps},
     madt::parse_madt,
     pic::disable_pic,
     sync::LOCKS_HELD,
@@ -90,6 +90,7 @@ extern "C" fn early_main(temp_loader_struct: *const LoaderStruct) -> ! {
 }
 fn main() {
     info!("Starting main kernel task...");
+    set_cpu_flags();
     let mut status = unsafe { AcpiInitializeSubsystem() };
     assert_eq!(status, AE_OK);
     status = unsafe { AcpiInitializeTables(null_mut(), 0, 0) };
@@ -105,7 +106,7 @@ fn main() {
 #[panic_handler]
 fn panic(info: &PanicInfo) -> ! {
     cli();
-    LOCKS_HELD.fetch_add(1, Ordering::SeqCst);
+    LOCKS_HELD[get_current_lapic_id()].fetch_add(1, Ordering::SeqCst);
     error!("{}", info);
     let str = format!("{}", info);
     display_print(&str);

kernel/src/sys/ioapic.rs

@@ -27,15 +27,17 @@ bitfield! {
     mask, set_mask: 16, 16;
     destination, set_destination: 63, 56;
 }
-const EMPTY_IOAPIC: IOAPIC = IOAPIC {
-    address: AtomicPtr::new(null_mut()),
-    start_gsi: AtomicUsize::new(0),
-    end_gsi: AtomicUsize::new(0),
-};
+
 const REGISTER_VERSION: u8 = 1;
 const REGISTER_REDIRECTION: u8 = 0x10;
 
-static IOAPICS: [IOAPIC; 32] = [EMPTY_IOAPIC; 32];
+static IOAPICS: [IOAPIC; 32] = [const {
+    IOAPIC {
+        address: AtomicPtr::new(null_mut()),
+        start_gsi: AtomicUsize::new(0),
+        end_gsi: AtomicUsize::new(0),
+    }
+}; 32];
 static NEXT_IOAPIC_ID: AtomicUsize = AtomicUsize::new(0);
 
 fn read_register(apic_i: usize, reg_i: u8) -> u32 {

kernel/src/sys/lapic.rs

@@ -1,4 +1,5 @@
 use core::{
+    arch::asm,
     ptr::{copy, null_mut},
     sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering},
 };
@@ -6,7 +7,10 @@ use core::{
 use alloc::{vec, vec::Vec};
 use kernel_common::{instructions::pause, paging::PageTable};
 
-use crate::cpu::paging::{map_physical, map_range, unmap, virt_to_phys, CURRENT_PML4};
+use crate::cpu::{
+    isr::{ISR_INVALIDATE_TLB, ISR_SCHEDULER},
+    paging::{map_physical, map_range, unmap, virt_to_phys, CURRENT_PML4},
+};
 
 use super::{
     hpet::sleep,
@@ -36,14 +40,14 @@ const REGISTER_TIMER_DIVIDE: usize = 0xf8;
 const IPI_INIT: u32 = 0x500;
 const IPI_STARTUP: u32 = 0x600;
 
-const EMPTY_LAPIC: LAPIC = LAPIC {
-    lapic_id: AtomicUsize::new(0),
-    present: AtomicBool::new(false),
-};
-
 static ADDRESS: AtomicPtr<u32> = AtomicPtr::new(null_mut());
 pub static BSP_LAPIC_ID: AtomicUsize = AtomicUsize::new(0);
-static LAPICS: [LAPIC; 256] = [EMPTY_LAPIC; 256];
+static LAPICS: [LAPIC; 256] = [const {
+    LAPIC {
+        lapic_id: AtomicUsize::new(0),
+        present: AtomicBool::new(false),
+    }
+}; 256];
 static TICKS_PER_MS: AtomicUsize = AtomicUsize::new(0);
 pub static NEXT_LAPIC_ID: AtomicUsize = AtomicUsize::new(0);
 
@@ -55,7 +59,13 @@ pub fn send_eoi() {
 }
 pub fn get_current_lapic_id() -> usize {
     let address = ADDRESS.load(Ordering::SeqCst);
-    unsafe { address.add(REGISTER_ID).read_volatile() as usize >> 24 }
+    let lapic_id = unsafe { address.add(REGISTER_ID).read_volatile() as usize >> 24 };
+    let rflags: u64;
+    unsafe {
+        asm!("pushfq; pop {};", out(reg) rflags);
+    }
+    assert!(rflags & (1 << 9) == 0);
+    lapic_id
 }
 pub fn setup_lapic(phys: u64) {
     if phys != 0 {
@@ -101,10 +111,11 @@ pub fn setup_lapic_timer(calibrate: bool) {
     }
     unsafe {
         address.add(REGISTER_TIMER_INITIAL_COUNT).write_volatile(0);
-        address.add(REGISTER_TIMER_LVT).write_volatile(254);
+        address.add(REGISTER_TIMER_LVT).write_volatile(ISR_SCHEDULER as u32);
     }
 }
 pub fn add_lapic(lapic_id: usize) {
+    assert!(lapic_id < 256);
     let next_id = NEXT_LAPIC_ID.fetch_add(1, Ordering::SeqCst);
     LAPICS[next_id].lapic_id.store(lapic_id, Ordering::SeqCst);
     LAPICS[next_id].present.store(true, Ordering::SeqCst);
@@ -144,8 +155,20 @@ pub fn start_aps() {
             pause();
         }
     }
+    ALL_APS_STARTED.store(true, Ordering::SeqCst);
     unsafe {
         unmap(0x1000);
     }
-    ALL_APS_STARTED.store(true, Ordering::SeqCst);
+}
+pub fn smp_invalidate_tlb() {
+    if !ALL_APS_STARTED.load(Ordering::SeqCst) {
+        return;
+    }
+    for i in 0..NEXT_LAPIC_ID.load(Ordering::SeqCst) {
+        let lapic_id = LAPICS[i].lapic_id.load(Ordering::SeqCst);
+        if lapic_id == BSP_LAPIC_ID.load(Ordering::SeqCst) {
+            continue;
+        }
+        send_ipi(lapic_id, ISR_INVALIDATE_TLB as u32);
+    }
 }

kernel/src/sys/scheduler.rs

@@ -1,6 +1,7 @@
 use core::{arch::asm, sync::atomic::Ordering};
 
 use alloc::{collections::vec_deque::VecDeque, vec::Vec};
+use kernel_common::instructions::cli;
 use spin::Mutex;
 
 use crate::cpu::isr::ISRState;
@@ -63,9 +64,11 @@ pub fn schedule_task(task: Task) {
     scheduler_list.push_back(task);
 }
 pub fn yield_task() {
-    assert!(!IN_ISR_HANDLER.load(Ordering::SeqCst));
-    assert_eq!(LOCKS_HELD.load(Ordering::SeqCst), 0);
+    cli();
+    let lapic_id = get_current_lapic_id();
+    assert!(!IN_ISR_HANDLER[lapic_id].load(Ordering::SeqCst));
+    assert_eq!(LOCKS_HELD[lapic_id].load(Ordering::SeqCst), 0);
     unsafe {
-        asm!("int $254");
+        asm!("sti; int $254");
     }
 }

kernel/src/sys/sync.rs

@@ -12,8 +12,8 @@ use super::{
     task::{Task, TaskState, CURRENT_TASKS, CURRENT_TASK_LOCK},
 };
 
-pub static IN_ISR_HANDLER: AtomicBool = AtomicBool::new(false);
-pub static LOCKS_HELD: AtomicUsize = AtomicUsize::new(0);
+pub static IN_ISR_HANDLER: [AtomicBool; 256] = [const { AtomicBool::new(false) }; 256];
+pub static LOCKS_HELD: [AtomicUsize; 256] = [const { AtomicUsize::new(0) }; 256];
 
 pub struct Semaphore {
     spinlock: Spinlock,
@@ -23,23 +23,30 @@ pub struct Semaphore {
 }
 pub struct Spinlock {
     locked: AtomicBool,
+    lapic_id: AtomicUsize,
 }
 
 impl Spinlock {
     pub const fn new() -> Self {
-        Self { locked: AtomicBool::new(false) }
+        Self {
+            locked: AtomicBool::new(false),
+            lapic_id: AtomicUsize::new(0),
+        }
     }
     pub fn lock(&self) {
         assert!(MULTITASKING_ENABLED.load(Ordering::SeqCst));
         cli();
         while !self.locked.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst).is_ok() {}
-        LOCKS_HELD.fetch_add(1, Ordering::SeqCst);
+        let lapic_id = get_current_lapic_id();
+        LOCKS_HELD[lapic_id].fetch_add(1, Ordering::SeqCst);
+        self.lapic_id.store(lapic_id, Ordering::SeqCst);
     }
     pub fn unlock(&self) {
         assert!(self.locked.load(Ordering::SeqCst));
+        let lapic_id = self.lapic_id.load(Ordering::SeqCst);
         self.locked.store(false, Ordering::SeqCst);
-        LOCKS_HELD.fetch_sub(1, Ordering::SeqCst);
-        if !IN_ISR_HANDLER.load(Ordering::SeqCst) && LOCKS_HELD.load(Ordering::SeqCst) == 0 {
+        LOCKS_HELD[lapic_id].fetch_sub(1, Ordering::SeqCst);
+        if !IN_ISR_HANDLER[lapic_id].load(Ordering::SeqCst) && LOCKS_HELD[lapic_id].load(Ordering::SeqCst) == 0 {
             unsafe {
                 sti();
             }

kernel/src/sys/task.rs

@@ -5,7 +5,7 @@ use core::{
 
 use alloc::{sync::Arc, vec, vec::Vec};
 use hashbrown::HashMap;
-use kernel_common::instructions::{hlt, pause};
+use kernel_common::instructions::{cli, hlt, pause, sti};
 use spin::{Lazy, Mutex};
 
 use crate::{
@@ -142,9 +142,13 @@ extern "C" fn task_entry() -> ! {
 }
 fn idle_main() {
     while !ALL_APS_STARTED.load(Ordering::SeqCst) {
+        cli();
         if STARTING_AP_ID.load(Ordering::SeqCst) == get_current_lapic_id() as i64 {
             let _ = STARTING_AP_ID.compare_exchange(get_current_lapic_id() as i64, -1, Ordering::SeqCst, Ordering::SeqCst);
         }
+        unsafe {
+            sti();
+        }
         pause();
     }
     loop {