From 68403dc029f3a1c1a98e3af89f6b01262283bd0d Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Sat, 15 Oct 2022 17:40:33 +0200
Subject: [PATCH] Kernel: Make AddressSpaces reference-counted

---
 kernel/include/memory/AddressSpace.h |  5 +++--
 kernel/src/memory/AddressSpace.cpp   | 18 +++++++++++++-----
 kernel/src/sys/exec.cpp              |  6 ++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/kernel/include/memory/AddressSpace.h b/kernel/include/memory/AddressSpace.h
index b6158a07..ef40b774 100644
--- a/kernel/include/memory/AddressSpace.h
+++ b/kernel/include/memory/AddressSpace.h
@@ -18,10 +18,11 @@ struct AddressSpace
 
     bool is_cloned()
     {
-        return m_cloned;
+        return *m_refs > 1;
     }
 
   private:
     PageTable* m_pml4;
-    bool m_cloned;
+
+    int* m_refs;
 };
\ No newline at end of file
diff --git a/kernel/src/memory/AddressSpace.cpp b/kernel/src/memory/AddressSpace.cpp
index 5b45f02a..bef29b71 100644
--- a/kernel/src/memory/AddressSpace.cpp
+++ b/kernel/src/memory/AddressSpace.cpp
@@ -4,22 +4,25 @@
 #include "log/Log.h"
 #include "memory/PMM.h"
 #include "memory/VMM.h"
+#include "std/stdlib.h"
 #include "utils/move.h"
 
 AddressSpace AddressSpace::create()
 {
     AddressSpace result;
     result.m_pml4 = (PageTable*)PMM::request_page();
-    result.m_cloned = false;
+    result.m_refs = (int*)kmalloc(sizeof(int));
+    *result.m_refs = 1;
     VMM::install_kernel_page_directory_into_address_space(result);
     return move(result);
 }
 
 void AddressSpace::destroy()
 {
-    if (m_cloned)
+    if (is_cloned())
     {
         kdbgln("Will not destroy a cloned address space, I don't own it");
+        (*m_refs)--;
         return;
     }
     uint64_t pages_freed = 0;
@@ -75,21 +78,26 @@ void AddressSpace::destroy()
     pages_freed++;
     PMM::free_page(m_pml4);
 
+    kfree(m_refs);
+
     kdbgln("Reclaimed %ld pages from address space!", pages_freed);
 }
 
 void AddressSpace::detach()
 {
-    if (!m_cloned) return;
+    if (!is_cloned()) return;
+    (*m_refs)--;
+    m_refs = (int*)kmalloc(sizeof(int));
+    *m_refs = 1;
     m_pml4 = (PageTable*)PMM::request_page();
     VMM::install_kernel_page_directory_into_address_space(*this);
-    m_cloned = false;
 }
 
 AddressSpace AddressSpace::clone()
 {
     AddressSpace result;
     result.m_pml4 = m_pml4;
-    result.m_cloned = true;
+    result.m_refs = m_refs;
+    *m_refs = *m_refs + 1;
     return result;
 }
\ No newline at end of file
diff --git a/kernel/src/sys/exec.cpp b/kernel/src/sys/exec.cpp
index b4d2335d..1f3f3c0e 100644
--- a/kernel/src/sys/exec.cpp
+++ b/kernel/src/sys/exec.cpp
@@ -73,6 +73,12 @@ void sys_exec(Context* context, const char* pathname)
     Task* task = Scheduler::current_task();
     ASSERT(task);
 
+    if (task->address_space.is_cloned())
+    {
+        task->address_space.detach();
+        VMM::switch_to_user_address_space(task->address_space);
+    }
+
     // At this point, pretty much nothing can fail.
 
     ELFImage* image = ELFLoader::load_elf_from_vfs(program);