From 46ae5de30ac645406ca1bec907d06099b944f97e Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Fri, 24 Nov 2023 17:19:32 -0800
Subject: [PATCH] [Zion] Free paging structures after process exit.

---
 zion/memory/paging_util.cpp   | 29 +++++++++++++++++++++++++++++
 zion/memory/paging_util.h     |  1 +
 zion/object/address_space.cpp |  2 ++
 zion/object/address_space.h   |  4 ++++
 zion/object/process.cpp       |  4 +++-
 5 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/zion/memory/paging_util.cpp b/zion/memory/paging_util.cpp
index 54a3caf..75b4c0a 100644
--- a/zion/memory/paging_util.cpp
+++ b/zion/memory/paging_util.cpp
@@ -87,6 +87,21 @@ uint64_t CurrCr3() {
   return pml4_addr;
 }
 
+void CleanupPageStructure(uint64_t struct_phys, uint64_t level) {
+  uint64_t* struct_virtual =
+      reinterpret_cast<uint64_t*>(boot::GetHigherHalfDirectMap() + struct_phys);
+
+  if (level > 0) {
+    for (uint16_t i = 0; i < 256; i++) {
+      if (struct_virtual[i] & PRESENT_BIT) {
+        CleanupPageStructure(struct_virtual[i] & ~0xFFF, level - 1);
+      }
+    }
+  }
+
+  phys_mem::FreePage(struct_phys);
+}
+
 }  // namespace
 
 void InitializePml4(uint64_t pml4_physical_addr) {
@@ -105,6 +120,20 @@ void InitializePml4(uint64_t pml4_physical_addr) {
   pml4_virtual[Pml4Index(hhdm)] = *Pml4Entry(curr_cr3, hhdm);
 }
 
+void CleanupPml4(uint64_t pml4_physical_addr) {
+  uint64_t* pml4_virtual = reinterpret_cast<uint64_t*>(
+      boot::GetHigherHalfDirectMap() + pml4_physical_addr);
+
+  // Iterate the first half of the pml4 as it contains user-space mappings.
+  for (uint8_t i = 0; i < 128; i++) {
+    if (pml4_virtual[i] & PRESENT_BIT) {
+      CleanupPageStructure(pml4_virtual[i] & ~0xFFF, 2);
+    }
+  }
+
+  phys_mem::FreePage(pml4_physical_addr);
+}
+
 void MapPage(uint64_t cr3, uint64_t vaddr, uint64_t paddr) {
   vaddr = PageAlign(vaddr);
   paddr = PageAlign(paddr);
diff --git a/zion/memory/paging_util.h b/zion/memory/paging_util.h
index 1722221..83c522a 100644
--- a/zion/memory/paging_util.h
+++ b/zion/memory/paging_util.h
@@ -5,6 +5,7 @@
 #include "object/process.h"
 
 void InitializePml4(uint64_t pml4_physical_addr);
+void CleanupPml4(uint64_t pml4_physical_addr);
 
 void MapPage(uint64_t cr3, uint64_t vaddr, uint64_t paddr);
 
diff --git a/zion/object/address_space.cpp b/zion/object/address_space.cpp
index 196593d..791fdbf 100644
--- a/zion/object/address_space.cpp
+++ b/zion/object/address_space.cpp
@@ -18,6 +18,8 @@ AddressSpace::AddressSpace() {
   InitializePml4(cr3_);
 }
 
+AddressSpace::~AddressSpace() { CleanupPml4(cr3_); }
+
 glcr::ErrorOr<uint64_t> AddressSpace::AllocateUserStack() {
   uint64_t base = user_stacks_.NewUserStack();
   auto mem_object = glcr::StaticCastRefPtr<MemoryObject>(
diff --git a/zion/object/address_space.h b/zion/object/address_space.h
index 0be9db1..7c4daf3 100644
--- a/zion/object/address_space.h
+++ b/zion/object/address_space.h
@@ -62,6 +62,10 @@ class AddressSpace : public KernelObject {
   AddressSpace(const AddressSpace&) = delete;
   AddressSpace(AddressSpace&&) = delete;
 
+  // Deconstructing an address space will free all paging structures associated
+  // with this address space.
+  ~AddressSpace();
+
   uint64_t cr3() { return cr3_; }
 
   // User Mappings.
diff --git a/zion/object/process.cpp b/zion/object/process.cpp
index f0c7e86..407f051 100644
--- a/zion/object/process.cpp
+++ b/zion/object/process.cpp
@@ -105,6 +105,8 @@ void Process::Cleanup() {
   PANIC_ON_ERR(vmas_->FreeAddressRange(0, kUserSpaceMax),
                "Failed to cleanup userspace mappings in process exit.");
 
-  // 4. Release paging structures. TODO
+  // 4. Release paging structures.
+  vmas_ = nullptr;
+
   state_ = FINISHED;
 }