From 5e6bdcc52d10ce55f0ea6f0fd64e502169fa70a3 Mon Sep 17 00:00:00 2001
From: kamkow1 <kamkow256@gmail.com>
Date: Thu, 1 Jan 2026 18:42:53 +0100
Subject: [PATCH] Handle swapgs in interrupts and scheduling

---
 kernel/amd64/intr_stub.S | 12 ++++++++++++
 kernel/amd64/msr.h       |  3 ---
 kernel/amd64/sched.S     |  6 ++++++
 kernel/amd64/smp.c       |  4 ++--
 kernel/proc/proc.c       |  7 ++++---
 5 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/kernel/amd64/intr_stub.S b/kernel/amd64/intr_stub.S
index 9f833f2..1085dc7 100644
--- a/kernel/amd64/intr_stub.S
+++ b/kernel/amd64/intr_stub.S
@@ -48,6 +48,12 @@
   amd64_intr ## n:; \
   x(n); \
   cli; \
+  movq %rsp, %rax; \
+  movq 144(%rax), %rax; \
+  testb $3, %al; \
+  jz 1f; \
+  swapgs; \
+1:; \
   push_regs; \
   cld; \
   movq %rsp, %rdi; \
@@ -58,6 +64,12 @@
   callq amd64_intr_handler; \
   movq (%rsp), %rsp; \
   pop_regs; \
+  movq %rsp, %rax; \
+  movq 144(%rax), %rax; \
+  testb $3, %al; \
+  jz 2f; \
+  swapgs; \
+2:; \
   addq $16, %rsp; \
   iretq;
   
diff --git a/kernel/amd64/msr.h b/kernel/amd64/msr.h
index 5c3d325..d1ccfdf 100644
--- a/kernel/amd64/msr.h
+++ b/kernel/amd64/msr.h
@@ -3,9 +3,6 @@
 
 #include <libk/std.h>
 
-#define MSR_FS_BASE 0xC0000100
-#define MSR_GS_BASE 0xC0000101
-
 uint64_t amd64_rdmsr (uint32_t msr);
 void amd64_wrmsr (uint32_t msr, uint64_t value);
 
diff --git a/kernel/amd64/sched.S b/kernel/amd64/sched.S
index 816937f..e21941c 100644
--- a/kernel/amd64/sched.S
+++ b/kernel/amd64/sched.S
@@ -20,5 +20,11 @@ amd64_do_sched:
   movq %rsi, %cr3
   movq %rdi, %rsp
   pop_regs
+
+  movq 144(%rsp), %rax
+  testb $3, %al
+  jz 1f
+  swapgs
+1:
   add $16, %rsp
   iretq
diff --git a/kernel/amd64/smp.c b/kernel/amd64/smp.c
index 036e340..6b07909 100644
--- a/kernel/amd64/smp.c
+++ b/kernel/amd64/smp.c
@@ -41,10 +41,10 @@ struct cpu* cpu_get (uint32_t id) {
 }
 
 /// Get ID of current running CPU
-uint32_t cpu_id (void) { return (uint32_t)amd64_rdmsr (MSR_GS_BASE); }
+uint32_t cpu_id (void) { return (uint32_t)amd64_rdmsr (MSR_SHADOW_GS_BASE); }
 
 /// Assign an ID to the current running CPU
-void cpu_assign (uint32_t id) { amd64_wrmsr (MSR_GS_BASE, (uint64_t)id); }
+void cpu_assign (uint32_t id) { amd64_wrmsr (MSR_SHADOW_GS_BASE, (uint64_t)id); }
 
 /// Bootstrap code for non-BSP CPUs
 static void amd64_smp_bootstrap (struct limine_mp_info* mp_info) {
diff --git a/kernel/proc/proc.c b/kernel/proc/proc.c
index 14c8514..8b6e84c 100644
--- a/kernel/proc/proc.c
+++ b/kernel/proc/proc.c
@@ -17,7 +17,9 @@
 #include <sys/smp.h>
 
 #if defined(__x86_64__)
-#include <amd64/intr_defs.h>
+  #include <amd64/intr_defs.h>
+  #include <amd64/msr.h>
+  #include <amd64/msr-index.h>
 #endif
 
 struct elf_aux {
@@ -216,9 +218,8 @@ void proc_sched (void) {
 
   next = proc_find_sched ();
 
-  if (next != NULL) {
+  if (next != NULL)
     thiscpu->proc_current = next;
-  }
 
   spin_unlock (&thiscpu->lock);