From b48da838ec186ad15ef40bbd214f546df8363051 Mon Sep 17 00:00:00 2001
From: Liam Kerr <liam.d.kerr@protonmail.com>
Date: Wed, 11 Feb 2026 22:56:25 +0000
Subject: [PATCH] (interrupt): implement robust trap frame and state
 preservation

---
 src/boot/traps.S        | 23 ++++++++++++-------
 src/kernel/interrupts.c | 10 ++++++---
 src/kernel/interrupts.h | 50 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 11 deletions(-)

diff --git a/src/boot/traps.S b/src/boot/traps.S
index 921de6f..9c2f024 100644
--- a/src/boot/traps.S
+++ b/src/boot/traps.S
@@ -4,11 +4,10 @@
 .global trap_entry
 
 trap_entry:
-    # 1. Create space on the stack for 32 registers (32 * 8 = 256 bytes)
+    # Create space on the stack for 32 registers (32 * 8 = 256 bytes)
     addi sp, sp, -256
 
-    # 2. Save all General Purpose Registers (GPRs)
-    # We don't save x0 (zero) because it's always zero
+    # Save all General Purpose Registers (GPRs)
     sd ra, 0(sp)
     sd gp, 8(sp)
     sd tp, 16(sp)
@@ -40,11 +39,19 @@ trap_entry:
     sd t5, 224(sp)
     sd t6, 232(sp)
 
-    # 3. Call your C handler
-    # The CPU already put the cause in 'mcause', so C can read it
+    # Save mepc to the stack (using offset 240)
+    csrr t0, mepc
+    sd t0, 240(sp)
+
+    # Call trap handler
+    mv a0, sp # put the stack pointer to the first arg
     call handle_trap
 
-    # 4. Restore all GPRs
+    # restore mepc
+    ld t0, 240(sp)
+    csrw mepc, t0
+
+    # Restore all GPRs
     ld ra, 0(sp)
     ld gp, 8(sp)
     ld tp, 16(sp)
@@ -76,8 +83,8 @@ trap_entry:
     ld t5, 224(sp)
     ld t6, 232(sp)
 
-    # 5. Shrink the stack back
+    # Shrink the stack back
     addi sp, sp, 256
 
-    # 6. Return from Machine-mode trap
+    # Return from Machine-mode trap
     mret
\ No newline at end of file
diff --git a/src/kernel/interrupts.c b/src/kernel/interrupts.c
index 83887e8..f6c4e71 100644
--- a/src/kernel/interrupts.c
+++ b/src/kernel/interrupts.c
@@ -24,6 +24,7 @@ void interrupt_init()
 
 void kpanic(const char *reason, ...)
 {
+    asm volatile("csrci mstatus, 8"); // Disable MIE (Machine Interrupt Enable)
     va_list args;
 	va_start(args, reason);
     kputs("\n!!! PANIC !!!\n");
@@ -50,7 +51,7 @@ void kpanic_force()
     }
 }
 
-void handle_trap()
+void handle_trap(trap_frame_t *registers)
 {
     // Read the 'mcause' register to see WHY we trapped
     unsigned long cause;
@@ -60,9 +61,10 @@ void handle_trap()
     // For 64-bit RISC-V, the bit is 63
     int is_interrupt = (cause >> 63) & 1;
 
+    unsigned long code = cause & 0xfff;
+
     if (is_interrupt)
     {
-        unsigned long code = cause & 0xfff;
         handle_interrupt(code);
         return;
     }
@@ -71,7 +73,9 @@ void handle_trap()
         // fault address (if applicable)
         uintptr_t mtval;
         asm volatile("csrr %0, mtval" : "=r"(mtval));
-        kprintf("Faulting Address (if applicable): %x\n", mtval);
+
+        kprintf("\n[EXCEPTION] Code: %d | Instruction: %x | Fault Address: %x\n", code, registers->mepc, mtval);
+
         switch (cause)
         {
         case 0:
diff --git a/src/kernel/interrupts.h b/src/kernel/interrupts.h
index e0db2e4..fd64efd 100644
--- a/src/kernel/interrupts.h
+++ b/src/kernel/interrupts.h
@@ -1,9 +1,59 @@
 #ifndef PANIC_H
 #define PANIC_H
 
+
+// This should match the order of registers in traps.S
+typedef struct {
+    // Return address and pointers
+    uint64_t ra;      // x1
+    uint64_t gp;      // x3
+    uint64_t tp;      // x4
+
+    // Temporary registers
+    uint64_t t0;      // x5
+    uint64_t t1;      // x6
+    uint64_t t2;      // x7
+
+    // Saved registers
+    uint64_t s0;      // x8 (fp)
+    uint64_t s1;      // x9
+
+    // Function arguments / Return values
+    uint64_t a0;      // x10
+    uint64_t a1;      // x11
+    uint64_t a2;      // x12
+    uint64_t a3;      // x13
+    uint64_t a4;      // x14
+    uint64_t a5;      // x15
+    uint64_t a6;      // x16
+    uint64_t a7;      // x17
+
+    // More saved registers
+    uint64_t s2;      // x18
+    uint64_t s3;      // x19
+    uint64_t s4;      // x20
+    uint64_t s5;      // x21
+    uint64_t s6;      // x22
+    uint64_t s7;      // x23
+    uint64_t s8;      // x24
+    uint64_t s9;      // x25
+    uint64_t s10;     // x26
+    uint64_t s11;     // x27
+
+    // More temporary registers
+    uint64_t t3;      // x28
+    uint64_t t4;      // x29
+    uint64_t t5;      // x30
+    uint64_t t6;      // x31
+
+    // Control and Status Register state
+    uint64_t mepc;    // offset 240 (Saved in traps.S)
+} trap_frame_t;
+
 void interrupt_init();
 void kpanic(const char *, ...);
 void kpanic_force();
+void handle_trap(trap_frame_t *registers);
 void handle_interrupt(unsigned long code);
 
 #define KASSERT(cond, msg)                                              \