From b360307f41d0dd6d94915cee0f6f04d73cb827bb Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Thu, 13 Oct 2022 17:58:13 +0200
Subject: [PATCH] VMM: Make it so much gooder

There are still some fixes to be made, but I think this is already way cleaner than before.
---
 kernel/include/memory/VMM.h |   8 +-
 kernel/src/memory/VMM.cpp   | 485 +++++++++++++++---------------------
 2 files changed, 213 insertions(+), 280 deletions(-)

diff --git a/kernel/include/memory/VMM.h b/kernel/include/memory/VMM.h
index bca59d0b..0c3e8b31 100644
--- a/kernel/include/memory/VMM.h
+++ b/kernel/include/memory/VMM.h
@@ -10,7 +10,6 @@ enum Flags
 namespace VMM
 {
     void init(); // fetch page table from cr3
-    void init(PageTable* cr3);
 
     void map(uint64_t virtualAddress, uint64_t physicalAddress, int flags);
     void remap(uint64_t virtualAddress, int flags);
@@ -18,4 +17,11 @@ namespace VMM
     uint64_t getPhysical(uint64_t virtualAddress);
     uint64_t getFlags(uint64_t virtualAddress);
 
+    PageDirectoryEntry* find_pde(PageTable* root, uint64_t virtualAddress);
+    PageDirectoryEntry* create_pde_if_not_exists(PageTable* root, uint64_t virtualAddress);
+
+    void propagate_read_write(PageTable* root, uint64_t virtualAddress);
+    void propagate_user(PageTable* root, uint64_t virtualAddress);
+
+    void flush_tlb(uint64_t addr);
 };
\ No newline at end of file
diff --git a/kernel/src/memory/VMM.cpp b/kernel/src/memory/VMM.cpp
index c04ee44f..4652be4b 100644
--- a/kernel/src/memory/VMM.cpp
+++ b/kernel/src/memory/VMM.cpp
@@ -1,6 +1,10 @@
+#define MODULE "vmm"
+
 #include "memory/VMM.h"
 #include "assert.h"
+#include "log/Log.h"
 #include "memory/PMM.h"
+#include "misc/utils.h"
 #include "std/string.h"
 
 // FIXME: There is a lot of duplicate code in this file. This should probably be refactored.
@@ -12,322 +16,245 @@ void VMM::init()
     asm volatile("mov %%cr3, %0" : "=r"(PML4));
 }
 
-void VMM::init(PageTable* cr3)
-{
-    PML4 = cr3;
-}
-
 void VMM::unmap(uint64_t virtualAddress)
 {
-    virtualAddress >>= 12;
-    uint64_t P_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PT_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PD_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PDP_i = virtualAddress & 0x1ff;
+    virtualAddress = Utilities::round_down_to_nearest_page(virtualAddress);
 
-    PageDirectoryEntry PDE;
+    PageDirectoryEntry* pde = find_pde(PML4, virtualAddress);
+    if (!pde) return; // Already unmapped
 
-    PDE = PML4->entries[PDP_i];
-    PageTable* PDP;
-    if (!PDE.Present)
-    {
-        return; // Already unmapped
-    }
-    else
-    {
-        if (PDE.LargerPages)
-        {
-            PDE.Present = false;
-            PDE.LargerPages = false;
-            PML4->entries[PDP_i] = PDE;
-            goto invalidate;
-        }
-        PDP = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PDP->entries[PD_i];
-    PageTable* PD;
-    if (!PDE.Present)
-    {
-        return; // Already unmapped
-    }
-    else
-    {
-        if (PDE.LargerPages)
-        {
-            PDE.Present = false;
-            PDE.LargerPages = false;
-            PDP->entries[PD_i] = PDE;
-            goto invalidate;
-        }
-        PD = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PD->entries[PT_i];
-    PageTable* PT;
-    if (!PDE.Present)
-    {
-        return; // Already unmapped
-    }
-    else
-    {
-        if (PDE.LargerPages)
-        {
-            PDE.LargerPages = false;
-            PDE.Present = false;
-            PD->entries[PT_i] = PDE;
-            goto invalidate;
-        }
-        PT = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PT->entries[P_i];
-    PDE.Present = false;
-    PT->entries[P_i] = PDE;
-invalidate:
-    asm volatile("invlpg (%0)" : : "r"(virtualAddress) : "memory");
+    memset(pde, 0, sizeof(PageDirectoryEntry));
+    flush_tlb(virtualAddress);
 }
 
 uint64_t VMM::getPhysical(uint64_t virtualAddress)
 {
-    virtualAddress >>= 12;
-    uint64_t P_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PT_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PD_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PDP_i = virtualAddress & 0x1ff;
+    PageDirectoryEntry* pde = find_pde(PML4, Utilities::round_down_to_nearest_page(virtualAddress));
+    if (!pde) return UINT64_MAX; // Not mapped
 
-    PageDirectoryEntry PDE;
-
-    PDE = PML4->entries[PDP_i];
-    PageTable* PDP;
-    if (!PDE.Present)
-    {
-        return UINT64_MAX; // Not mapped
-    }
-    else
-    {
-        if (PDE.LargerPages) return PDE.Address << 12 | (virtualAddress & PAGE_SIZE);
-        PDP = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PDP->entries[PD_i];
-    PageTable* PD;
-    if (!PDE.Present)
-    {
-        return UINT64_MAX; // Not mapped
-    }
-    else
-    {
-        if (PDE.LargerPages) return PDE.Address << 12 | (virtualAddress & PAGE_SIZE);
-        PD = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PD->entries[PT_i];
-    PageTable* PT;
-    if (!PDE.Present)
-    {
-        return UINT64_MAX; // Not mapped
-    }
-    else
-    {
-        if (PDE.LargerPages) return PDE.Address << 12 | (virtualAddress & PAGE_SIZE);
-        PT = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PT->entries[P_i];
-    if (!PDE.Present) return UINT64_MAX;
-    return PDE.Address << 12 | (virtualAddress & PAGE_SIZE);
+    return pde->Address << 12 | (virtualAddress % PAGE_SIZE);
 }
 
-uint64_t VMM::getFlags(uint64_t virtualAddress) // FIXME: Add support for larger pages to getFlags.
+uint64_t VMM::getFlags(uint64_t virtualAddress)
 {
-    virtualAddress >>= 12;
-    uint64_t P_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PT_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PD_i = virtualAddress & 0x1ff;
-    virtualAddress >>= 9;
-    uint64_t PDP_i = virtualAddress & 0x1ff;
+    PageDirectoryEntry* pde = find_pde(PML4, Utilities::round_down_to_nearest_page(virtualAddress));
+    if (!pde) return 0; // Not mapped
 
-    PageDirectoryEntry PDE;
-
-    PDE = PML4->entries[PDP_i];
-    PageTable* PDP;
-    if (!PDE.Present)
-    {
-        return 0; // Not mapped
-    }
-    else
-    {
-        if (PDE.LargerPages) goto read_flags;
-        PDP = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PDP->entries[PD_i];
-    PageTable* PD;
-    if (!PDE.Present)
-    {
-        return 0; // Not mapped
-    }
-    else
-    {
-        if (PDE.LargerPages) goto read_flags;
-        PD = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PD->entries[PT_i];
-    PageTable* PT;
-    if (!PDE.Present)
-    {
-        return 0; // Not mapped
-    }
-    else
-    {
-        if (PDE.LargerPages) goto read_flags;
-        PT = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-
-    PDE = PT->entries[P_i];
-
-read_flags:
     uint64_t flags = 0;
-    if (PDE.UserSuper) flags |= User;
-    if (PDE.ReadWrite) flags |= ReadWrite;
+    if (pde->UserSuper) flags |= User;
+    if (pde->ReadWrite) flags |= ReadWrite;
     return flags;
 }
 
 void VMM::map(uint64_t virtualAddress, uint64_t physicalAddress, int flags)
+{
+    virtualAddress = Utilities::round_down_to_nearest_page(virtualAddress);
+    PageDirectoryEntry* pde = find_pde(PML4, virtualAddress);
+    bool will_flush_tlb = true;
+    if (!pde)
+    {
+        pde = create_pde_if_not_exists(PML4, virtualAddress);
+        will_flush_tlb = false;
+    }
+    else if (pde->LargerPages)
+    {
+        unmap(virtualAddress);
+        pde = create_pde_if_not_exists(PML4, virtualAddress);
+        will_flush_tlb = false;
+    }
+
+    pde->set_address(Utilities::round_down_to_nearest_page(physicalAddress));
+    if (flags & User) propagate_user(PML4, virtualAddress);
+    if (flags & ReadWrite) propagate_read_write(PML4, virtualAddress);
+    if (will_flush_tlb) flush_tlb(virtualAddress);
+}
+
+PageDirectoryEntry* VMM::find_pde(PageTable* root, uint64_t virtualAddress)
 {
     virtualAddress >>= 12;
-    uint64_t P_i = virtualAddress & 0x1ff;
+    uint64_t page_index = virtualAddress & 0x1ff;
     virtualAddress >>= 9;
-    uint64_t PT_i = virtualAddress & 0x1ff;
+    uint64_t pt_index = virtualAddress & 0x1ff;
     virtualAddress >>= 9;
-    uint64_t PD_i = virtualAddress & 0x1ff;
+    uint64_t pd_index = virtualAddress & 0x1ff;
     virtualAddress >>= 9;
-    uint64_t PDP_i = virtualAddress & 0x1ff;
+    uint64_t pdp_index = virtualAddress & 0x1ff;
+    PageDirectoryEntry* pde;
+    PageTable* pt;
 
-    PageDirectoryEntry PDE;
+    pde = &root->entries[pdp_index]; // PML4
+    if (!pde->Present) return nullptr;
+    else if (pde->LargerPages)
+        return pde;
+    else { pt = (PageTable*)((uint64_t)pde->Address << 12); }
 
-    PDE = PML4->entries[PDP_i];
-    PageTable* PDP;
-    if (!PDE.Present)
+    pde = &pt->entries[pd_index]; // PDP
+    if (!pde->Present) return nullptr;
+    else if (pde->LargerPages)
+        return pde;
+    else { pt = (PageTable*)((uint64_t)pde->Address << 12); }
+
+    pde = &pt->entries[pt_index]; // PD
+    if (!pde->Present) return nullptr;
+    else if (pde->LargerPages)
+        return pde;
+    else { pt = (PageTable*)((uint64_t)pde->Address << 12); }
+
+    pde = &pt->entries[page_index]; // PT
+    if (!pde->Present) return nullptr;
+    return pde;
+}
+
+PageDirectoryEntry* VMM::create_pde_if_not_exists(PageTable* root, uint64_t virtualAddress)
+{
+    virtualAddress >>= 12;
+    uint64_t page_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pt_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pd_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pdp_index = virtualAddress & 0x1ff;
+    PageDirectoryEntry* pde;
+    PageTable* pt;
+
+    pde = &root->entries[pdp_index]; // PML4
+    if (!pde->Present)
     {
-        PDP = (PageTable*)PMM::request_page();
-        ASSERT(!(PMM_DID_FAIL(PDP)));
-        memset(PDP, 0, PAGE_SIZE);
-        PDE.set_address((uint64_t)PDP);
-        PDE.Present = true;
-        PDE.ReadWrite = true;
-        if (flags & User) PDE.UserSuper = true;
-        PML4->entries[PDP_i] = PDE;
+        pt = (PageTable*)PMM::request_page();
+        ASSERT(!(PMM_DID_FAIL(pt)));
+        memset(pt, 0, PAGE_SIZE);
+        pde->set_address((uint64_t)pt);
+        pde->Present = true;
     }
+    else if (pde->LargerPages)
+        return pde;
+    else { pt = (PageTable*)((uint64_t)pde->Address << 12); }
+
+    pde = &pt->entries[pd_index]; // PDP
+    if (!pde->Present)
+    {
+        pt = (PageTable*)PMM::request_page();
+        ASSERT(!(PMM_DID_FAIL(pt)));
+        memset(pt, 0, PAGE_SIZE);
+        pde->set_address((uint64_t)pt);
+        pde->Present = true;
+    }
+    else if (pde->LargerPages)
+        return pde;
+    else { pt = (PageTable*)((uint64_t)pde->Address << 12); }
+
+    pde = &pt->entries[pt_index]; // PD
+    if (!pde->Present)
+    {
+        pt = (PageTable*)PMM::request_page();
+        ASSERT(!(PMM_DID_FAIL(pt)));
+        memset(pt, 0, PAGE_SIZE);
+        pde->set_address((uint64_t)pt);
+        pde->Present = true;
+    }
+    else if (pde->LargerPages)
+        return pde;
+    else { pt = (PageTable*)((uint64_t)pde->Address << 12); }
+
+    pde = &pt->entries[page_index]; // PT
+    if (!pde->Present) { pde->Present = true; }
+    return pde;
+}
+
+void VMM::propagate_read_write(PageTable* root, uint64_t virtualAddress)
+{
+    virtualAddress >>= 12;
+    uint64_t page_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pt_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pd_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pdp_index = virtualAddress & 0x1ff;
+    PageDirectoryEntry* pde;
+    PageTable* pt;
+
+    pde = &root->entries[pdp_index]; // PML4
+    if (!pde->Present) return;
     else
     {
-        if (PDE.LargerPages)
-        {
-            unmap(virtualAddress);
-            PDE.LargerPages = false;
-            PDP = (PageTable*)PMM::request_page();
-            ASSERT(!(PMM_DID_FAIL(PDP)));
-            memset(PDP, 0, PAGE_SIZE);
-            PDE.set_address((uint64_t)PDP);
-            PDE.Present = true;
-            PDE.ReadWrite = true;
-            if (flags & User) PDE.UserSuper = true;
-            PML4->entries[PDP_i] = PDE;
-        }
-        PDP = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-    if ((flags & User) && !PDE.UserSuper)
-    {
-        PDE.UserSuper = true;
-        PML4->entries[PDP_i] = PDE;
+        pde->ReadWrite = true;
+        if (pde->LargerPages) return;
+        pt = (PageTable*)((uint64_t)pde->Address << 12);
     }
 
-    PDE = PDP->entries[PD_i];
-    PageTable* PD;
-    if (!PDE.Present)
-    {
-        PD = (PageTable*)PMM::request_page();
-        ASSERT(!(PMM_DID_FAIL(PD)));
-        memset(PD, 0, PAGE_SIZE);
-        PDE.set_address((uint64_t)PD);
-        PDE.Present = true;
-        PDE.ReadWrite = true;
-        if (flags & User) PDE.UserSuper = true;
-        PDP->entries[PD_i] = PDE;
-    }
+    pde = &pt->entries[pd_index]; // PDP
+    if (!pde->Present) return;
     else
     {
-        if (PDE.LargerPages)
-        {
-            unmap(virtualAddress);
-            PDE.LargerPages = false;
-            PD = (PageTable*)PMM::request_page();
-            ASSERT(!(PMM_DID_FAIL(PD)));
-            memset(PD, 0, PAGE_SIZE);
-            PDE.set_address((uint64_t)PD);
-            PDE.Present = true;
-            PDE.ReadWrite = true;
-            if (flags & User) PDE.UserSuper = true;
-            PDP->entries[PD_i] = PDE;
-        }
-        PD = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-    if ((flags & User) && !PDE.UserSuper)
-    {
-        PDE.UserSuper = true;
-        PDP->entries[PD_i] = PDE;
+        pde->ReadWrite = true;
+        if (pde->LargerPages) return;
+        pt = (PageTable*)((uint64_t)pde->Address << 12);
     }
 
-    PDE = PD->entries[PT_i];
-    PageTable* PT;
-    if (!PDE.Present)
-    {
-        PT = (PageTable*)PMM::request_page();
-        ASSERT(!(PMM_DID_FAIL(PT)));
-        memset(PT, 0, PAGE_SIZE);
-        PDE.set_address((uint64_t)PT);
-        PDE.Present = true;
-        PDE.ReadWrite = true;
-        if (flags & User) PDE.UserSuper = true;
-        PD->entries[PT_i] = PDE;
-    }
+    pde = &pt->entries[pt_index]; // PD
+    if (!pde->Present) return;
     else
     {
-        if (PDE.LargerPages)
-        {
-            unmap(virtualAddress);
-            PDE.LargerPages = false;
-            PT = (PageTable*)PMM::request_page();
-            ASSERT(!(PMM_DID_FAIL(PT)));
-            memset(PT, 0, PAGE_SIZE);
-            PDE.set_address((uint64_t)PT);
-            PDE.Present = true;
-            PDE.ReadWrite = true;
-            if (flags & User) PDE.UserSuper = true;
-            PD->entries[PT_i] = PDE;
-        }
-        PT = (PageTable*)((uint64_t)PDE.Address << 12);
-    }
-    if ((flags & User) && !PDE.UserSuper)
-    {
-        PDE.UserSuper = true;
-        PD->entries[PT_i] = PDE;
+        pde->ReadWrite = true;
+        if (pde->LargerPages) return;
+        pt = (PageTable*)((uint64_t)pde->Address << 12);
     }
 
-    PDE = PT->entries[P_i];
-    PDE.Present = true;
-    PDE.ReadWrite = flags & ReadWrite;
-    PDE.UserSuper = flags & User;
-    PDE.set_address(physicalAddress);
-    PT->entries[P_i] = PDE;
+    pde = &pt->entries[page_index];
+    if (!pde->Present) return;
+    else
+        pde->ReadWrite = true;
+}
+
+void VMM::propagate_user(PageTable* root, uint64_t virtualAddress)
+{
+    virtualAddress >>= 12;
+    uint64_t page_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pt_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pd_index = virtualAddress & 0x1ff;
+    virtualAddress >>= 9;
+    uint64_t pdp_index = virtualAddress & 0x1ff;
+    PageDirectoryEntry* pde;
+    PageTable* pt;
+
+    pde = &root->entries[pdp_index]; // PML4
+    if (!pde->Present) return;
+    else
+    {
+        pde->UserSuper = true;
+        if (pde->LargerPages) return;
+        pt = (PageTable*)((uint64_t)pde->Address << 12);
+    }
+
+    pde = &pt->entries[pd_index]; // PDP
+    if (!pde->Present) return;
+    else
+    {
+        pde->UserSuper = true;
+        if (pde->LargerPages) return;
+        pt = (PageTable*)((uint64_t)pde->Address << 12);
+    }
+
+    pde = &pt->entries[pt_index]; // PD
+    if (!pde->Present) return;
+    else
+    {
+        pde->UserSuper = true;
+        if (pde->LargerPages) return;
+        pt = (PageTable*)((uint64_t)pde->Address << 12);
+    }
+
+    pde = &pt->entries[page_index];
+    if (!pde->Present) return;
+    else
+        pde->UserSuper = true;
+}
+
+void VMM::flush_tlb(uint64_t addr)
+{
+    asm volatile("invlpg (%0)" : : "r"(addr) : "memory");
 }
\ No newline at end of file