feat: Enhance Kata memory management with stack growth support and additional user stack properties

8 files changed, 78 insertions, 13 deletions

diff --git a/mirai/asm/entry.zig b/mirai/asm/entry.zig
index 2e14e5c..4f9adf6 100644
--- a/mirai/asm/entry.zig
+++ b/mirai/asm/entry.zig
@@ -67,9 +67,9 @@ comptime {
         \\  # We push in REVERSE order so struct layout matches (stack grows down)
         \\  # NOTE: RCX and R11 already have user_rip and user_rflags from SYSCALL!
         \\  push %r14              # user_rsp (last in struct, push first)
-        \\  push %r11              # user_rflags (R11 saved by SYSCALL)
+        \\  push 32(%r15)          # user_rflags (R11 value before syscall)
         \\  push $0x1B             # user_ss
-        \\  push %rcx              # user_rip (RCX saved by SYSCALL)
+        \\  push 96(%r15)          # user_rip (RCX value before syscall)
         \\  push $0x23             # user_cs
         \\  push 0(%r15)           # r15
         \\  push 8(%r15)           # r14
diff --git a/mirai/boot/tss/tss.zig b/mirai/boot/tss/tss.zig
index 215ad7b..395ef06 100644
--- a/mirai/boot/tss/tss.zig
+++ b/mirai/boot/tss/tss.zig
@@ -23,6 +23,7 @@ const TSS = packed struct {
 
 var tss: TSS align(16) = undefined;
 var kernel_stack: [boot_limits.KERNEL_STACK_SIZE]u8 align(16) = undefined;
+var ist1_stack: [65536]u8 align(16) = undefined; // 64KB for double fault/page fault
 
 pub fn init() void {
     const tss_bytes = @as([*]u8, @ptrCast(&tss));
@@ -31,8 +32,9 @@ pub fn init() void {
     }
 
     tss.rsp0 = @intFromPtr(&kernel_stack) + kernel_stack.len;
+    tss.ist1 = @intFromPtr(&ist1_stack) + ist1_stack.len;
 
-    serial.printf("TSS: addr={x} rsp0={x}\n", .{ @intFromPtr(&tss), tss.rsp0 });
+    serial.printf("TSS: addr={x} rsp0={x} ist1={x}\n", .{ @intFromPtr(&tss), tss.rsp0, tss.ist1 });
 }
 
 pub fn get_address() u64 {
diff --git a/mirai/common/constants/memory.zig b/mirai/common/constants/memory.zig
index 0ec5883..08b7409 100644
--- a/mirai/common/constants/memory.zig
+++ b/mirai/common/constants/memory.zig
@@ -25,8 +25,9 @@ pub const KATA_SPACE_END: u64 = 0x0000800000000000;
 
 // Stack configuration
 pub const USER_STACK_TOP: u64 = 0x00007FFFFFF00000;
-pub const USER_STACK_PAGES: u64 = 512; // 2MB
-pub const USER_STACK_SIZE: u64 = USER_STACK_PAGES * PAGE_SIZE;
+pub const USER_STACK_MAX_PAGES: u64 = 512; // 2MB max
+pub const USER_STACK_INITIAL_PAGES: u64 = 16; // 64KB initial, grows on demand
+pub const USER_STACK_SIZE: u64 = USER_STACK_MAX_PAGES * PAGE_SIZE;
 pub const KERNEL_STACK_PAGES: u64 = 4; // 16KB
 pub const KERNEL_STACK_SIZE: u64 = KERNEL_STACK_PAGES * PAGE_SIZE;
 
diff --git a/mirai/crimson/exception.zig b/mirai/crimson/exception.zig
index 6ba9b96..67147f7 100644
--- a/mirai/crimson/exception.zig
+++ b/mirai/crimson/exception.zig
@@ -1,7 +1,9 @@
 //! CPU exception handler
 
+const kata_memory = @import("../kata/memory.zig");
 const memory = @import("../asm/memory.zig");
 const panic = @import("panic.zig");
+const sensei = @import("../kata/sensei/sensei.zig");
 const serial = @import("../drivers/serial/serial.zig");
 const types = @import("types.zig");
 
@@ -43,6 +45,21 @@ const NAMES = [_][]const u8{
 export fn exception_handler(frame_ptr: u64) void {
     const frame = @as(*types.ExceptionFrame, @ptrFromInt(frame_ptr));
 
+    // Handle page faults - check for stack growth
+    if (frame.int_num == 14) {
+        const cr2 = memory.read_page_fault_address();
+        const is_not_present = (frame.error_code & 1) == 0;
+
+        // Check if fault is in user stack region (demand paging)
+        if (sensei.get_current_kata()) |kata| {
+            if (cr2 >= kata.user_stack_bottom and cr2 < kata.user_stack_top and is_not_present) {
+                if (kata_memory.grow_stack(kata, cr2)) {
+                    return;
+                }
+            }
+        }
+    }
+
     serial.print("\n!!! EXCEPTION: ");
     if (frame.int_num < NAMES.len) {
         serial.print(NAMES[frame.int_num]);
@@ -63,6 +80,10 @@ export fn exception_handler(frame_ptr: u64) void {
         const cr2 = memory.read_page_fault_address();
         serial.printf("CR2 (fault addr): {x}\n", .{cr2});
 
+        if (sensei.get_current_kata()) |kata| {
+            serial.printf("Stack: bottom={x} committed={x} top={x}\n", .{ kata.user_stack_bottom, kata.user_stack_committed, kata.user_stack_top });
+        }
+
         serial.print("Fault type: ");
         if ((frame.error_code & 1) == 0) {
             serial.print("Page not present");
diff --git a/mirai/interrupts/idt.zig b/mirai/interrupts/idt.zig
index 326f441..97fefcf 100644
--- a/mirai/interrupts/idt.zig
+++ b/mirai/interrupts/idt.zig
@@ -30,7 +30,12 @@ pub fn init() void {
 
     for (0..idt_const.NUM_EXCEPTIONS) |i| {
         const handler_addr = @intFromPtr(isr.get_exception_handler(@intCast(i)));
-        set_gate(@intCast(i), handler_addr, gdt_const.KERNEL_CODE, idt_const.GATE_INTERRUPT);
+        // Use IST1 for double fault (8) and page fault (14)
+        if (i == 8 or i == 14) {
+            set_gate_with_ist(@intCast(i), handler_addr, gdt_const.KERNEL_CODE, idt_const.GATE_INTERRUPT, 1);
+        } else {
+            set_gate(@intCast(i), handler_addr, gdt_const.KERNEL_CODE, idt_const.GATE_INTERRUPT);
+        }
     }
 
     set_gate(idt_const.VECTOR_TIMER, @intFromPtr(isr.get_irq_handler(0)), gdt_const.KERNEL_CODE, idt_const.GATE_INTERRUPT);
@@ -51,10 +56,14 @@ pub fn init() void {
 }
 
 fn set_gate(num: u8, handler: u64, selector: u16, type_attr: u8) void {
+    set_gate_with_ist(num, handler, selector, type_attr, 0);
+}
+
+fn set_gate_with_ist(num: u8, handler: u64, selector: u16, type_attr: u8, ist: u8) void {
     table[num] = types.Entry{
         .offset_low = @truncate(handler & 0xFFFF),
         .selector = selector,
-        .ist = 0,
+        .ist = ist,
         .type_attr = type_attr,
         .offset_mid = @truncate((handler >> 16) & 0xFFFF),
         .offset_high = @truncate(handler >> 32),
diff --git a/mirai/kata/memory.zig b/mirai/kata/memory.zig
index a479c25..880cf63 100644
--- a/mirai/kata/memory.zig
+++ b/mirai/kata/memory.zig
@@ -4,7 +4,6 @@ const memory_const = @import("../common/constants/memory.zig");
 const memory_limits = @import("../common/limits/memory.zig");
 const paging = @import("../memory/paging.zig");
 const pmm = @import("../memory/pmm.zig");
-const serial = @import("../drivers/serial/serial.zig");
 const types = @import("types.zig");
 
 const HIGHER_HALF = memory_const.HIGHER_HALF_START;
@@ -77,11 +76,13 @@ pub const VirtualBuffer = struct {
 pub fn setup(kata: *types.Kata, framebuffer_phys: u64, framebuffer_size: u64) !void {
     kata.page_table = try paging.create_page_table();
 
-    const user_stack_base = memory_const.USER_STACK_TOP - (memory_const.USER_STACK_PAGES * PAGE_SIZE);
+    kata.user_stack_top = memory_const.USER_STACK_TOP;
+    kata.user_stack_bottom = memory_const.USER_STACK_TOP - (memory_const.USER_STACK_MAX_PAGES * PAGE_SIZE);
+    kata.user_stack_committed = memory_const.USER_STACK_TOP - (memory_const.USER_STACK_INITIAL_PAGES * PAGE_SIZE);
 
-    for (0..memory_const.USER_STACK_PAGES) |i| {
+    for (0..memory_const.USER_STACK_INITIAL_PAGES) |i| {
         const page = pmm.alloc_page() orelse return error.OutOfMemory;
-        const virt = user_stack_base + (i * PAGE_SIZE);
+        const virt = kata.user_stack_committed + (i * PAGE_SIZE);
         _ = try paging.map_page_in_table(kata.page_table, virt, page, paging.PAGE_WRITABLE | paging.PAGE_USER);
 
         const page_ptr: [*]volatile u8 = @ptrFromInt(page + HIGHER_HALF);
@@ -89,9 +90,7 @@ pub fn setup(kata: *types.Kata, framebuffer_phys: u64, framebuffer_size: u64) !v
             page_ptr[j] = 0;
         }
     }
-    kata.user_stack_top = memory_const.USER_STACK_TOP;
 
-    // Allocate kernel stack pages - must be contiguous for identity mapping
     const first_page = pmm.alloc_page() orelse return error.OutOfMemory;
     const kernel_stack_base = first_page;
 
@@ -225,4 +224,33 @@ pub fn cleanup(kata: *types.Kata) void {
     }
     kata.stack_top = 0;
     kata.user_stack_top = 0;
+    kata.user_stack_bottom = 0;
+    kata.user_stack_committed = 0;
+}
+
+pub fn grow_stack(kata: *types.Kata, fault_addr: u64) bool {
+    const asm_memory = @import("../asm/memory.zig");
+    const page_addr = fault_addr & ~@as(u64, 0xFFF);
+
+    if (page_addr >= kata.user_stack_committed or page_addr < kata.user_stack_bottom) {
+        return false;
+    }
+
+    var addr = kata.user_stack_committed - PAGE_SIZE;
+    while (addr >= page_addr) : (addr -= PAGE_SIZE) {
+        const page = pmm.alloc_page() orelse return false;
+        _ = paging.map_page_in_table(kata.page_table, addr, page, paging.PAGE_WRITABLE | paging.PAGE_USER) catch return false;
+
+        const page_ptr: [*]volatile u8 = @ptrFromInt(page + HIGHER_HALF);
+        for (0..PAGE_SIZE) |j| {
+            page_ptr[j] = 0;
+        }
+
+        asm_memory.invalidate_page(addr);
+
+        if (addr == 0) break;
+    }
+
+    kata.user_stack_committed = page_addr;
+    return true;
 }
diff --git a/mirai/kata/pool.zig b/mirai/kata/pool.zig
index b89fac3..3fff3e0 100644
--- a/mirai/kata/pool.zig
+++ b/mirai/kata/pool.zig
@@ -106,6 +106,8 @@ fn create_empty() types.Kata {
         .page_table = 0,
         .stack_top = 0,
         .user_stack_top = 0,
+        .user_stack_bottom = 0,
+        .user_stack_committed = 0,
         .attachments = [_]?*attachment.Attachment{null} ** kata_limits.MAX_ATTACHMENTS,
         .current_location = undefined,
         .current_location_len = 1,
diff --git a/mirai/kata/types.zig b/mirai/kata/types.zig
index 16a5024..efbbe8d 100644
--- a/mirai/kata/types.zig
+++ b/mirai/kata/types.zig
@@ -68,6 +68,8 @@ pub const Kata = struct {
     page_table: u64,
     stack_top: u64,
     user_stack_top: u64,
+    user_stack_bottom: u64,
+    user_stack_committed: u64,
 
     attachments: [kata_limits.MAX_ATTACHMENTS]?*attachment.Attachment,