From bb6de0ce8d19406268434dbd4ff320c208542269 Mon Sep 17 00:00:00 2001
From: Mathieu Strypsteen <mathieu@strypsteen.me>
Date: Fri, 27 Dec 2024 12:15:38 +0100
Subject: [PATCH] Add stack guard and fix early exceptions

---
 kernel/src/cpu/boot.s     | 4 +++-
 kernel/src/cpu/isr.rs     | 8 ++++++--
 kernel/src/cpu/paging.rs  | 5 +++++
 kernel/src/sys/syscall.rs | 2 +-
 4 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/kernel/src/cpu/boot.s b/kernel/src/cpu/boot.s
index 45f0cb3..f1bafe6 100644
--- a/kernel/src/cpu/boot.s
+++ b/kernel/src/cpu/boot.s
@@ -17,6 +17,8 @@ _start:
 .global __stack_chk_guard
 __stack_chk_guard:
 .skip 8
-.align 16
+.align 0x1000
+.global _stack_guard
+_stack_guard:
 .skip 0x10000
 stack:
diff --git a/kernel/src/cpu/isr.rs b/kernel/src/cpu/isr.rs
index dd96f0c..f0a3e7e 100644
--- a/kernel/src/cpu/isr.rs
+++ b/kernel/src/cpu/isr.rs
@@ -10,6 +10,7 @@ use crate::{
     sys::{
         lapic::{get_current_lapic_id, send_eoi},
         locks::Spinlock,
+        madt::INTERRUPTS_SETUP,
         scheduler::scheduler,
         sync::{IN_ISR_HANDLER, LOCKS_HELD},
         task::{TaskState, CURRENT_TASKS},
@@ -109,8 +110,11 @@ extern "C" fn isr_handler(state: &mut ISRState) {
         }
         return;
     }
-    let lapic_id = get_current_lapic_id();
-    IN_ISR_HANDLER[lapic_id].store(true, Ordering::SeqCst);
+    let mut lapic_id = 0;
+    if INTERRUPTS_SETUP.load(Ordering::SeqCst) {
+        lapic_id = get_current_lapic_id();
+        IN_ISR_HANDLER[lapic_id].store(true, Ordering::SeqCst);
+    }
     if state.isr < 32 {
         if state.cs == 0x23 && state.isr != 2 && state.isr != 8 && state.isr != 18 {
             debug!("Exception in usermode: {}", EXCEPTIONS[state.isr as usize]);
diff --git a/kernel/src/cpu/paging.rs b/kernel/src/cpu/paging.rs
index f0016fd..87d28d2 100644
--- a/kernel/src/cpu/paging.rs
+++ b/kernel/src/cpu/paging.rs
@@ -24,6 +24,7 @@ extern "C" {
     static _data_end: u8;
     static _bss_start: u8;
     static _bss_end: u8;
+    static _stack_guard: u8;
 }
 
 pub struct AddressSpace {
@@ -82,6 +83,7 @@ impl AddressSpace {
     fn map(&mut self, virt: u64, phys: u64, user: bool, write: bool, exec: bool, cache_disable: bool, only_flags: bool) {
         assert!(virt >= 0x1000, "First page shouldn't be mapped");
         assert!(!write || !exec || virt == 0x1000);
+        assert!(user == (virt < USER_END) || virt == 0x1000);
         {
             let mut frames_vec = PHYSICAL_FRAMES.lock();
             let frame = phys as usize / 0x1000;
@@ -249,6 +251,7 @@ pub fn setup_paging(loader_struct: &LoaderStruct, phys_start: u64, heap_start: u
         assert!(!frames_vec.get(1).unwrap());
         frames_vec.set(1, true);
     }
+    let stack_guard;
     let text_start;
     let text_end;
     let rodata_start;
@@ -258,6 +261,7 @@ pub fn setup_paging(loader_struct: &LoaderStruct, phys_start: u64, heap_start: u
     let bss_start;
     let bss_end;
     unsafe {
+        stack_guard = &_stack_guard as *const u8 as u64;
         text_start = &_text_start as *const u8 as u64 / 0x1000;
         text_end = (&_text_end as *const u8 as u64 + 0xfff) / 0x1000;
         rodata_start = &_rodata_start as *const u8 as u64 / 0x1000;
@@ -304,6 +308,7 @@ pub fn setup_paging(loader_struct: &LoaderStruct, phys_start: u64, heap_start: u
         }
     }
     unsafe {
+        address_space.update_flags_range(stack_guard, 0x1000, false, false, false);
         address_space.switch();
     }
     PAGING_ACTIVE.store(true, Ordering::SeqCst);
diff --git a/kernel/src/sys/syscall.rs b/kernel/src/sys/syscall.rs
index c2904a4..4353aeb 100644
--- a/kernel/src/sys/syscall.rs
+++ b/kernel/src/sys/syscall.rs
@@ -39,7 +39,7 @@ fn copy_from_user(start: u64, size: usize) -> Result<Vec<u8>, Box<dyn Error>> {
     let mut process = process.lock();
     for i in start / 0x1000..(end + 0xfff) / 0x1000 {
         let page = process.address_space.get_page(i * 0x1000);
-        if page.is_none() || page.unwrap().present() == 0 {
+        if page.is_none() {
             return Err("Invalid copy requested".into());
         }
     }