From fe1d55e70c2ff44b6d6d049c5cd3f47f0bf3303c Mon Sep 17 00:00:00 2001
From: jewelcodes <omarmelghoul01@gmail.com>
Date: Sun, 6 Oct 2024 03:19:31 -0400
Subject: [PATCH] x86_64: paging and context fixes

---
 src/platform/x86_64/cpu/paging.c    |  8 ++++++--
 src/platform/x86_64/sched/context.c | 22 +++++++++++++++-------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/platform/x86_64/cpu/paging.c b/src/platform/x86_64/cpu/paging.c
index b86f152..0ec221d 100644
--- a/src/platform/x86_64/cpu/paging.c
+++ b/src/platform/x86_64/cpu/paging.c
@@ -262,7 +262,7 @@ uint64_t clonePagingLayer(uint64_t ptr, int layer) {
                 newPhys = pmmAllocate();
                 if(!newPhys) return 0;
 
-                oldPhys = parent[i] & ~(PAGE_SIZE-1) & ~PT_PAGE_NXE;
+                oldPhys = parent[i] & ~((PAGE_SIZE-1) | PT_PAGE_NXE);
                 memcpy((void *)vmmMMIO(newPhys, true), (const void *)vmmMMIO(oldPhys, true), PAGE_SIZE);
 
                 clone[i] = newPhys | (parent[i] & ((uint64_t)PT_PAGE_LOW_FLAGS | PT_PAGE_NXE));   // copy the parent's permissions
@@ -275,6 +275,8 @@ uint64_t clonePagingLayer(uint64_t ptr, int layer) {
                 clone[i] = clonePagingLayer(oldPhys, layer+1);
                 clone[i] |= parent[i] & PT_PAGE_LOW_FLAGS;  // copy permissions again
             }
+        } else {
+            clone[i] = 0;
         }
     }
 
@@ -296,8 +298,10 @@ void *platformCloneUserSpace(uintptr_t parent) {
     for(int i = 0; i < 256; i++) {
         uint64_t ptr = oldPML4[i] & ~(PAGE_SIZE-1);
         uint64_t flags = oldPML4[i] & PT_PAGE_LOW_FLAGS;
-        if(oldPML4[i]) {
+        if((flags & PT_PAGE_PRESENT) && ptr) {
             newPML4[i] = clonePagingLayer(ptr, 0) | flags;
+        } else {
+            newPML4[i] = 0;
         }
     }
 
diff --git a/src/platform/x86_64/sched/context.c b/src/platform/x86_64/sched/context.c
index 7134daf..35b7e1e 100644
--- a/src/platform/x86_64/sched/context.c
+++ b/src/platform/x86_64/sched/context.c
@@ -289,14 +289,17 @@ void setLocalSched(bool sched) {
 
 void freePT(uint64_t *base, int depth, int maxdepth) {
     if(depth < 0 || depth > maxdepth) return;
+
+    PhysicalMemoryStatus st;
+    pmmStatus(&st);
+
     for(int i = 0; i < 512; i++) {
         uint64_t entry = base[i];
         uint64_t phys = entry & ~((PAGE_SIZE-1) | PT_PAGE_NXE);
-        if((entry & PT_PAGE_PRESENT) && phys) {
+        if((entry & PT_PAGE_PRESENT) && (phys) && (phys < st.highestUsableAddress)) {
             if(depth < maxdepth)
                 freePT((uint64_t *) vmmMMIO(phys, true), depth+1, maxdepth);
-            
-            //KDEBUG("freeing 0x%X at depth %d\n", phys, depth);
+
             pmmFree(phys);
         }
     }
@@ -315,10 +318,15 @@ void platformCleanThread(void *ptr, uintptr_t highest) {
     if(!ctx->cr3) return;
     
     // free the page tables themselves and all associated physical mem
-    uint64_t saved = readCR3();
-    writeCR3(ctx->cr3);
-    vmmFree(USER_BASE_ADDRESS, (highest-USER_BASE_ADDRESS) / PAGE_SIZE);
-    writeCR3(saved);
+    //uint64_t saved = readCR3();
+    //writeCR3(ctx->cr3);
+    //vmmFree(USER_BASE_ADDRESS, (highest-USER_BASE_ADDRESS) / PAGE_SIZE);
+    //writeCR3(saved);
+
+    uint64_t *pml4 = (uint64_t *) vmmMMIO(ctx->cr3, true);
+    for(int i = 0; i < 256; i++) {
+        freePT((uint64_t *) vmmMMIO(pml4[i] & ~(PAGE_SIZE-1), true), 1, 3);
+    }
 
     pmmFree(ctx->cr3);
 }