From 84e1b4cdb4f82c86aa018fff874ced8167855f2f Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Thu, 23 Nov 2023 07:12:23 -0800
Subject: [PATCH] [Zion] Plumb user stack to free function on thread exit.

---
 zion/object/address_space.cpp | 4 ++++
 zion/object/address_space.h   | 1 +
 zion/object/thread.cpp        | 9 +++++----
 zion/object/thread.h          | 1 +
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/zion/object/address_space.cpp b/zion/object/address_space.cpp
index 7299118..b41e4c7 100644
--- a/zion/object/address_space.cpp
+++ b/zion/object/address_space.cpp
@@ -22,6 +22,10 @@ uint64_t AddressSpace::AllocateUserStack() {
   return user_stacks_.NewUserStack();
 }
 
+void AddressSpace::FreeUserStack(uint64_t rsp) {
+  return user_stacks_.FreeUserStack(rsp);
+}
+
 uint64_t AddressSpace::GetNextMemMapAddr(uint64_t size) {
   if (size == 0) {
     panic("Zero size memmap");
diff --git a/zion/object/address_space.h b/zion/object/address_space.h
index cd5c1a7..56d58e7 100644
--- a/zion/object/address_space.h
+++ b/zion/object/address_space.h
@@ -66,6 +66,7 @@ class AddressSpace : public KernelObject {
 
   // User Mappings.
   uint64_t AllocateUserStack();
+  void FreeUserStack(uint64_t);
   uint64_t GetNextMemMapAddr(uint64_t size);
 
   // Maps in a memory object at a specific address.
diff --git a/zion/object/thread.cpp b/zion/object/thread.cpp
index e3440e6..1adf607 100644
--- a/zion/object/thread.cpp
+++ b/zion/object/thread.cpp
@@ -67,14 +67,14 @@ void Thread::Init() {
 #if K_THREAD_DEBUG
   dbgln("Thread start.", pid(), id_);
 #endif
-  uint64_t rsp = process_.vmas()->AllocateUserStack();
+  uint64_t rsp_ = process_.vmas()->AllocateUserStack();
   // TODO: Investigate this further but without this GCC
   // will emit movaps calls to non-16-bit-aligned stack
   // addresses.
-  rsp -= 0x8;
-  *reinterpret_cast<uint64_t*>(rsp) = kStackBaseSentinel;
+  rsp_ -= 0x8;
+  *reinterpret_cast<uint64_t*>(rsp_) = kStackBaseSentinel;
   SetRsp0(rsp0_start_);
-  jump_user_space(rip_, rsp, arg1_, arg2_);
+  jump_user_space(rip_, rsp_, arg1_, arg2_);
 }
 
 void Thread::Exit() {
@@ -102,6 +102,7 @@ void Thread::Cleanup() {
   // TODO: Race condition when called from exit, once kernel stack manager
   // actually reuses stacks this will cause an issue
   KernelVmm::FreeKernelStack(rsp0_start_);
+  process_.vmas()->FreeUserStack(rsp_);
 }
 
 void Thread::Wait() {
diff --git a/zion/object/thread.h b/zion/object/thread.h
index c678d99..1fae0b1 100644
--- a/zion/object/thread.h
+++ b/zion/object/thread.h
@@ -77,6 +77,7 @@ class Thread : public KernelObject, public glcr::IntrusiveListNode<Thread> {
   State state_ = CREATED;
 
   // Startup Context for the thread.
+  uint64_t rsp_;
   uint64_t rip_;
   uint64_t arg1_;
   uint64_t arg2_;