From d842443869f01de7e0b61b4ca5ad49bcaef4517a Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Sat, 18 Feb 2023 00:18:42 +0100
Subject: [PATCH] kernel: Switch to a stack we control for the main task as
 soon as we leave early boot

This is probably pointless as the main task only exists up until the point where the Scheduler starts running, but some init code was running into a stack overflow due to the fact that we only had one bootloader-provided page of stack.

This should probably be changed afterwards.
---
 kernel/src/arch/CPU.h          |  2 ++
 kernel/src/arch/x86_64/CPU.cpp |  9 +++++++++
 kernel/src/main.cpp            | 14 +++++++++++---
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/kernel/src/arch/CPU.h b/kernel/src/arch/CPU.h
index 54422463..3e56d105 100644
--- a/kernel/src/arch/CPU.h
+++ b/kernel/src/arch/CPU.h
@@ -26,5 +26,7 @@ namespace CPU
     void get_stack_trace_at(Registers* regs, void (*callback)(u64, void*), void* arg);
     void print_stack_trace_at(Registers* regs);
 
+    [[noreturn]] void bootstrap_switch_stack(u64 stack, void* function);
+
     void pause();
 }
diff --git a/kernel/src/arch/x86_64/CPU.cpp b/kernel/src/arch/x86_64/CPU.cpp
index 5a56a7c8..b703071c 100644
--- a/kernel/src/arch/x86_64/CPU.cpp
+++ b/kernel/src/arch/x86_64/CPU.cpp
@@ -295,6 +295,15 @@ namespace CPU
             &frame_index);
     }
 
+    [[noreturn]] void bootstrap_switch_stack(u64 stack, void* function)
+    {
+        asm volatile("mov %0, %%rsp\n"
+                     "jmp *%1"
+                     :
+                     : "r"(stack), "r"(function));
+        __builtin_unreachable();
+    }
+
     void pause()
     {
         asm volatile("pause");
diff --git a/kernel/src/main.cpp b/kernel/src/main.cpp
index aad0a78d..2b3e3ffc 100644
--- a/kernel/src/main.cpp
+++ b/kernel/src/main.cpp
@@ -90,11 +90,19 @@ Result<void> init()
     return {};
 }
 
-extern "C" [[noreturn]] void _start()
+[[noreturn]] void init_wrapper()
 {
-    Init::check_magic();
-    Init::early_init();
     auto rc = init();
     if (rc.has_error()) kerrorln("Runtime error: %s", rc.error_string());
     CPU::idle_loop();
 }
+
+// FIXME: Add a guard page to make sure the stack doesn't end up in random kernel memory. Also reclaim this memory after
+// leaving the init task.
+extern "C" [[noreturn]] void _start()
+{
+    Init::check_magic();
+    Init::early_init();
+    Stack stack { MemoryManager::alloc_for_kernel(8, MMU::ReadWrite | MMU::NoExecute).value(), 8 * ARCH_PAGE_SIZE };
+    CPU::bootstrap_switch_stack(stack.top(), (void*)init_wrapper);
+}