From f28b850db259a0f0c8f4977a00b563d790bfa745 Mon Sep 17 00:00:00 2001
From: Daniel Hammer <daniel@nexlink.cloud>
Date: Thu, 26 Feb 2026 20:12:56 +0100
Subject: [PATCH] feat: implement SYS_FREE

---
 kernel/src/Api/Heap.hpp      | 65 ++++++++++++++++++++++++++++++++++--
 kernel/src/Memory/Paging.cpp | 28 ++++++++++++++++
 kernel/src/Memory/Paging.hpp |  3 ++
 3 files changed, 93 insertions(+), 3 deletions(-)

diff --git a/kernel/src/Api/Heap.hpp b/kernel/src/Api/Heap.hpp
index ec355a4..b462663 100644
--- a/kernel/src/Api/Heap.hpp
+++ b/kernel/src/Api/Heap.hpp
@@ -11,6 +11,27 @@
 #include <Memory/PageFrameAllocator.hpp>
 
 namespace Zenith {
+
+    // Per-process heap allocation tracking (separate from Process struct to avoid bloating it)
+    struct HeapAlloc {
+        uint64_t va;
+        uint64_t numPages;
+    };
+
+    static constexpr int MaxHeapAllocs = 128;
+
+    static HeapAlloc g_heapAllocs[Sched::MaxProcesses][MaxHeapAllocs];
+    static int g_heapAllocCount[Sched::MaxProcesses];
+
+    // Get the process table slot index for the current process
+    static int GetCurrentSlot() {
+        auto* proc = Sched::GetCurrentProcessPtr();
+        if (proc == nullptr) return -1;
+        // Process slot index = pointer offset from slot 0
+        auto* slot0 = Sched::GetProcessSlot(0);
+        return (int)(proc - slot0);
+    }
+
     static uint64_t Sys_Alloc(uint64_t size) {
         auto* proc = Sched::GetCurrentProcessPtr();
         if (proc == nullptr) return 0;
@@ -30,10 +51,48 @@ namespace Zenith {
         }
 
         proc->heapNext += size;
+
+        // Track the allocation so Sys_Free can release it
+        int slot = GetCurrentSlot();
+        if (slot >= 0 && g_heapAllocCount[slot] < MaxHeapAllocs) {
+            g_heapAllocs[slot][g_heapAllocCount[slot]++] = { userVa, numPages };
+        }
+
         return userVa;
     }
 
-    static void Sys_Free(uint64_t) {
-        // TODO No-op for now (pages leak)
+    static void Sys_Free(uint64_t addr) {
+        auto* proc = Sched::GetCurrentProcessPtr();
+        if (proc == nullptr) return;
+
+        int slot = GetCurrentSlot();
+        if (slot < 0) return;
+
+        // Find the allocation record matching this address
+        int idx = -1;
+        for (int i = 0; i < g_heapAllocCount[slot]; i++) {
+            if (g_heapAllocs[slot][i].va == addr) {
+                idx = i;
+                break;
+            }
+        }
+        if (idx < 0) return;  // Unknown address — ignore
+
+        uint64_t va = g_heapAllocs[slot][idx].va;
+        uint64_t numPages = g_heapAllocs[slot][idx].numPages;
+
+        // Free each physical page and unmap the virtual address
+        for (uint64_t i = 0; i < numPages; i++) {
+            uint64_t pageVa = va + i * 0x1000;
+            uint64_t physAddr = Memory::VMM::Paging::GetPhysAddr(proc->pml4Phys, pageVa);
+            if (physAddr != 0) {
+                Memory::g_pfa->Free((void*)Memory::HHDM(physAddr));
+            }
+            Memory::VMM::Paging::UnmapUserIn(proc->pml4Phys, pageVa);
+        }
+
+        // Remove tracking entry by swapping with the last element
+        g_heapAllocs[slot][idx] = g_heapAllocs[slot][g_heapAllocCount[slot] - 1];
+        g_heapAllocCount[slot]--;
     }
-};
\ No newline at end of file
+};
diff --git a/kernel/src/Memory/Paging.cpp b/kernel/src/Memory/Paging.cpp
index b0177b5..2520a0b 100644
--- a/kernel/src/Memory/Paging.cpp
+++ b/kernel/src/Memory/Paging.cpp
@@ -240,6 +240,34 @@ namespace Memory::VMM {
         pageEntry->Address = physicalAddress >> 12;
     }
 
+    void Paging::UnmapUserIn(std::uint64_t pml4Phys, std::uint64_t virtualAddress) {
+        VirtualAddress va(virtualAddress);
+
+        // Walk without allocating — if any level is absent, nothing is mapped.
+        auto walkRead = [](PageTable* table, uint64_t index) -> PageTable* {
+            PageTableEntry* entry = (PageTableEntry*)Memory::HHDM(&table->entries[index]);
+            if (!entry->Present) return nullptr;
+            return (PageTable*)(entry->Address << 12);
+        };
+
+        PageTable* pml4 = (PageTable*)pml4Phys;
+        auto pml3 = walkRead(pml4, va.GetL4Index());
+        if (!pml3) return;
+        auto pml2 = walkRead(pml3, va.GetL3Index());
+        if (!pml2) return;
+        auto pml1 = walkRead(pml2, va.GetL2Index());
+        if (!pml1) return;
+
+        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&pml1->entries[va.GetPageIndex()]);
+        if (!pageEntry->Present) return;
+
+        // Clear the entire 8-byte PTE
+        *(uint64_t*)pageEntry = 0;
+
+        // Invalidate TLB for this virtual address
+        asm volatile("invlpg (%0)" :: "r"(virtualAddress) : "memory");
+    }
+
     std::uint64_t Paging::GetPhysAddr(std::uint64_t pml4, std::uint64_t virtualAddress, bool use40BitL1) {
         VirtualAddress virtualAddressObj(virtualAddress);
 
diff --git a/kernel/src/Memory/Paging.hpp b/kernel/src/Memory/Paging.hpp
index cb7d0c6..73b33b3 100644
--- a/kernel/src/Memory/Paging.hpp
+++ b/kernel/src/Memory/Paging.hpp
@@ -109,6 +109,9 @@ public:
         // Map a page into an arbitrary PML4 with User + Write-Combining attributes.
         static void MapUserInWC(std::uint64_t pml4Phys, std::uint64_t physicalAddress, std::uint64_t virtualAddress);
 
+        // Unmap a single page from an arbitrary PML4 (clears PTE + invalidates TLB).
+        static void UnmapUserIn(std::uint64_t pml4Phys, std::uint64_t virtualAddress);
+
         // Identity-map EFI runtime service regions so firmware code can
         // reference its own data at physical addresses.
         void MapEfiRuntime(limine_efi_memmap_response* efiMemmap);