From 624b1e7e69aedf9110179ebec1f1e2e78b6fc491 Mon Sep 17 00:00:00 2001
From: Daniel Hammer <daniel@nexlink.cloud>
Date: Wed, 4 Mar 2026 07:41:40 +0100
Subject: [PATCH] fix: fix memory leaks

---
 kernel/src/Api/Heap.hpp               |  4 +++
 kernel/src/Api/Process.hpp            |  3 +--
 kernel/src/Api/WinServer.cpp          | 35 +++++++++++++++++++++++++++
 kernel/src/Sched/Scheduler.hpp        |  3 +++
 programs/src/desktop/app_terminal.cpp |  1 +
 5 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/kernel/src/Api/Heap.hpp b/kernel/src/Api/Heap.hpp
index dbc053d..5d85e76 100644
--- a/kernel/src/Api/Heap.hpp
+++ b/kernel/src/Api/Heap.hpp
@@ -54,6 +54,7 @@ namespace Montauk {
 
         // Track the allocation so Sys_Free can release it
         int slot = GetCurrentSlot();
+        if (slot >= 0) Sched::g_allocatedPages[slot] += numPages;
         if (slot >= 0 && g_heapAllocCount[slot] < MaxHeapAllocs) {
             g_heapAllocs[slot][g_heapAllocCount[slot]++] = { userVa, numPages };
         }
@@ -66,6 +67,7 @@ namespace Montauk {
     static void CleanupHeapForSlot(int slot, uint64_t /*pml4Phys*/) {
         if (slot < 0 || slot >= Sched::MaxProcesses) return;
         g_heapAllocCount[slot] = 0;
+        Sched::g_allocatedPages[slot] = 0;
     }
 
     static void Sys_Free(uint64_t addr) {
@@ -98,6 +100,8 @@ namespace Montauk {
             Memory::VMM::Paging::UnmapUserIn(proc->pml4Phys, pageVa);
         }
 
+        if (slot >= 0) Sched::g_allocatedPages[slot] -= numPages;
+
         // Remove tracking entry by swapping with the last element
         g_heapAllocs[slot][idx] = g_heapAllocs[slot][g_heapAllocCount[slot] - 1];
         g_heapAllocCount[slot]--;
diff --git a/kernel/src/Api/Process.hpp b/kernel/src/Api/Process.hpp
index b1e850b..023c0ae 100644
--- a/kernel/src/Api/Process.hpp
+++ b/kernel/src/Api/Process.hpp
@@ -92,8 +92,7 @@ namespace Montauk {
                     buf[count].name[j] = proc->name[j];
                 buf[count].name[j] = '\0';
             }
-            buf[count].heapUsed = (proc->heapNext > Sched::UserHeapBase)
-                ? proc->heapNext - Sched::UserHeapBase : 0;
+            buf[count].heapUsed = Sched::g_allocatedPages[i] * 0x1000;
             count++;
         }
         return count;
diff --git a/kernel/src/Api/WinServer.cpp b/kernel/src/Api/WinServer.cpp
index fdfee85..b71ea8d 100644
--- a/kernel/src/Api/WinServer.cpp
+++ b/kernel/src/Api/WinServer.cpp
@@ -85,6 +85,25 @@ namespace WinServer {
         WindowSlot& slot = g_slots[windowId];
         if (!slot.used || slot.ownerPid != callerPid) return -1;
 
+        // Unmap pixel pages from desktop's address space
+        if (slot.desktopVa != 0 && slot.desktopPid != 0) {
+            auto* desktopProc = Sched::GetProcessByPid(slot.desktopPid);
+            if (desktopProc) {
+                for (int p = 0; p < slot.pixelNumPages; p++) {
+                    Memory::VMM::Paging::UnmapUserIn(
+                        desktopProc->pml4Phys,
+                        slot.desktopVa + (uint64_t)p * 0x1000);
+                }
+            }
+        }
+
+        // Free physical pixel pages
+        for (int i = 0; i < slot.pixelNumPages; i++) {
+            if (slot.pixelPhysPages[i] != 0) {
+                Memory::g_pfa->Free((void*)Memory::HHDM(slot.pixelPhysPages[i]));
+            }
+        }
+
         slot.used = false;
         return 0;
     }
@@ -179,6 +198,15 @@ namespace WinServer {
         int numPages = (int)((bufSize + 0xFFF) / 0x1000);
         if (numPages > MaxPixelPages) return -1;
 
+        // Free old pixel pages before allocating new ones
+        int oldNumPages = slot.pixelNumPages;
+        for (int i = 0; i < oldNumPages; i++) {
+            if (slot.pixelPhysPages[i] != 0) {
+                Memory::g_pfa->Free((void*)Memory::HHDM(slot.pixelPhysPages[i]));
+                slot.pixelPhysPages[i] = 0;
+            }
+        }
+
         // Allocate new pages and map into owner's address space
         uint64_t userVa = heapNext;
         for (int i = 0; i < numPages; i++) {
@@ -243,6 +271,13 @@ namespace WinServer {
                     }
                 }
 
+                // Free physical pixel pages
+                for (int p = 0; p < g_slots[i].pixelNumPages; p++) {
+                    if (g_slots[i].pixelPhysPages[p] != 0) {
+                        Memory::g_pfa->Free((void*)Memory::HHDM(g_slots[i].pixelPhysPages[p]));
+                    }
+                }
+
                 g_slots[i].used = false;
             }
         }
diff --git a/kernel/src/Sched/Scheduler.hpp b/kernel/src/Sched/Scheduler.hpp
index 4d42c49..3632b9c 100644
--- a/kernel/src/Sched/Scheduler.hpp
+++ b/kernel/src/Sched/Scheduler.hpp
@@ -88,4 +88,7 @@ namespace Sched {
     // Get a pointer to slot i in the process table (for enumeration)
     Process* GetProcessSlot(int slot);
 
+    // Per-process allocated page count (tracked by Heap syscalls, separate from Process struct)
+    inline uint64_t g_allocatedPages[MaxProcesses] = {};
+
 }
diff --git a/programs/src/desktop/app_terminal.cpp b/programs/src/desktop/app_terminal.cpp
index 6c6e525..77f04a6 100644
--- a/programs/src/desktop/app_terminal.cpp
+++ b/programs/src/desktop/app_terminal.cpp
@@ -39,6 +39,7 @@ static void terminal_on_key(Window* win, const Montauk::KeyEvent& key) {
 static void terminal_on_close(Window* win) {
     TerminalState* ts = (TerminalState*)win->app_data;
     if (ts) {
+        if (ts->child_pid > 0) montauk::kill(ts->child_pid);
         if (ts->cells) montauk::free(ts->cells);
         if (ts->alt_cells) montauk::free(ts->alt_cells);
         montauk::mfree(ts);