From b4527786d4a12b4a8a5e0188124f20f5d6545e25 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Wed, 2 Aug 2023 22:17:32 +0200
Subject: [PATCH] kernel/MMU: Don't copy pages on clone and allow manipulating
 other page directories

---
 kernel/src/arch/MMU.h          | 12 ++++++------
 kernel/src/arch/x86_64/MMU.cpp | 35 ++++++++++++++++------------------
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/kernel/src/arch/MMU.h b/kernel/src/arch/MMU.h
index 2211ee5a..2db1230a 100644
--- a/kernel/src/arch/MMU.h
+++ b/kernel/src/arch/MMU.h
@@ -29,15 +29,15 @@ namespace MMU
 
     u64 translate_physical_address(u64 phys);
 
-    Result<void> map(u64 virt, u64 phys, int flags, UseHugePages use_huge_pages);
-    Result<u64> unmap(u64 virt);
-    Result<u64> get_physical(u64 virt);
-    Result<int> get_flags(u64 virt);
-    Result<void> remap(u64 virt, int flags);
-
     void switch_page_directory(PageDirectory* dir);
     PageDirectory* get_page_directory();
 
+    Result<void> map(u64 virt, u64 phys, int flags, UseHugePages use_huge_pages, PageDirectory* directory = nullptr);
+    Result<u64> unmap(u64 virt, PageDirectory* directory = nullptr);
+    Result<u64> get_physical(u64 virt, PageDirectory* directory = nullptr);
+    Result<int> get_flags(u64 virt, PageDirectory* directory = nullptr);
+    Result<void> remap(u64 virt, int flags);
+
     void flush_all();
 
     Result<PageDirectory*> create_page_directory_for_userspace();
diff --git a/kernel/src/arch/x86_64/MMU.cpp b/kernel/src/arch/x86_64/MMU.cpp
index a56156f8..4914eb04 100644
--- a/kernel/src/arch/x86_64/MMU.cpp
+++ b/kernel/src/arch/x86_64/MMU.cpp
@@ -110,10 +110,11 @@ namespace MMU
         return result;
     }
 
-    PageTableEntry& l4_entry(u64 virt)
+    PageTableEntry& l4_entry(u64 virt, PageDirectory* directory = nullptr)
     {
         auto index = l4_index(virt);
-        return get_virtual_page_directory()->entries[index];
+        auto* vdir = directory ? translate_physical(directory) : get_virtual_page_directory();
+        return vdir->entries[index];
     }
 
     PageDirectory& page_table(const PageTableEntry& entry)
@@ -139,9 +140,9 @@ namespace MMU
         return page_table(entry).entries[index];
     }
 
-    Result<PageTableEntry*> find_entry(u64 virt)
+    Result<PageTableEntry*> find_entry(u64 virt, PageDirectory* directory = nullptr)
     {
-        const auto& l4 = l4_entry(virt);
+        const auto& l4 = l4_entry(virt, directory);
         if (!l4.present) return err(EFAULT);
         auto& l3 = l3_entry(l4, virt);
         if (!l3.present) return err(EFAULT);
@@ -152,9 +153,9 @@ namespace MMU
         return &l1_entry(l2, virt);
     }
 
-    Result<PageTableEntry*> apply_cascading_flags(u64 virt, int flags)
+    Result<PageTableEntry*> apply_cascading_flags(u64 virt, int flags, PageDirectory* directory = nullptr)
     {
-        auto& l4 = l4_entry(virt);
+        auto& l4 = l4_entry(virt, directory);
         if (!l4.present) return err(EFAULT);
         if (flags & Flags::ReadWrite) l4.read_write = true;
         if (flags & Flags::User) l4.user = true;
@@ -183,9 +184,9 @@ namespace MMU
         entry.set_address(phys);
     }
 
-    Result<void> map(u64 virt, u64 phys, int flags, UseHugePages use_huge_pages)
+    Result<void> map(u64 virt, u64 phys, int flags, UseHugePages use_huge_pages, PageDirectory* directory)
     {
-        auto& l4 = l4_entry(virt);
+        auto& l4 = l4_entry(virt, directory);
         if (!l4.present)
         {
             const u64 addr = TRY(MemoryManager::alloc_frame());
@@ -252,9 +253,9 @@ namespace MMU
         return {};
     }
 
-    Result<u64> unmap(u64 virt)
+    Result<u64> unmap(u64 virt, PageDirectory* directory)
     {
-        auto& l1 = *TRY(find_entry(virt));
+        auto& l1 = *TRY(find_entry(virt, directory));
         if (!l1.present) return err(EFAULT);
         const u64 address = l1.get_address();
         l1.clear();
@@ -262,16 +263,16 @@ namespace MMU
         return address;
     }
 
-    Result<u64> get_physical(u64 virt)
+    Result<u64> get_physical(u64 virt, PageDirectory* directory)
     {
-        const auto& l1 = *TRY(find_entry(virt));
+        const auto& l1 = *TRY(find_entry(virt, directory));
         if (!l1.present) return err(EFAULT);
         return l1.get_address();
     }
 
-    Result<int> get_flags(u64 virt)
+    Result<int> get_flags(u64 virt, PageDirectory* directory)
     {
-        const auto& l1 = *TRY(find_entry(virt));
+        const auto& l1 = *TRY(find_entry(virt, directory));
         if (!l1.present) return err(EFAULT);
         return arch_flags_to_mmu(l1);
     }
@@ -440,12 +441,8 @@ namespace MMU
 
                     for (u64 l = 0; l < 512; l++)
                     {
-                        PageTableEntry& old_l1 = old_pt->entries[l];
-                        if (!old_l1.present) continue;
                         PageTableEntry& new_l1 = new_pt->entries[l];
-                        new_l1.set_address(TRY(MemoryManager::alloc_frame()));
-
-                        memcpy(&page_table(new_l1), &page_table(old_l1), ARCH_PAGE_SIZE);
+                        new_l1.present = false; // The entry must be filled in by the caller.
                     }
                 }
             }