From 44045b258cf83a8aa987ca7b2255c1382af060be Mon Sep 17 00:00:00 2001
From: Andreas Kling <awesomekling@gmail.com>
Date: Tue, 23 Oct 2018 15:53:11 +0200
Subject: [PATCH] Teach MM to flush individual TLB entries only.

---
 Kernel/MemoryManager.cpp | 19 +++++++++++--------
 Kernel/MemoryManager.h   |  3 ++-
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/Kernel/MemoryManager.cpp b/Kernel/MemoryManager.cpp
index 0587e51d1c..1ad4c738ff 100644
--- a/Kernel/MemoryManager.cpp
+++ b/Kernel/MemoryManager.cpp
@@ -92,8 +92,8 @@ void MemoryManager::protectMap(LinearAddress linearAddress, size_t length)
         pte.setUserAllowed(false);
         pte.setPresent(false);
         pte.setWritable(false);
+        flushTLB(pteAddress);
     }
-    flushTLB();
 }
 
 void MemoryManager::identityMap(LinearAddress linearAddress, size_t length)
@@ -106,8 +106,8 @@ void MemoryManager::identityMap(LinearAddress linearAddress, size_t length)
         pte.setUserAllowed(true);
         pte.setPresent(true);
         pte.setWritable(true);
+        flushTLB(pteAddress);
     }
-    flushTLB();
 }
 
 void MemoryManager::initialize()
@@ -155,11 +155,11 @@ byte* MemoryManager::quickMapOnePage(PhysicalAddress physicalAddress)
     pte.setPhysicalPageBase(physicalAddress.pageBase());
     pte.setPresent(true);
     pte.setWritable(true);
-    flushTLB();
+    flushTLB(LinearAddress(4 * MB));
     return (byte*)(4 * MB);
 }
 
-void MemoryManager::flushTLB()
+void MemoryManager::flushEntireTLB()
 {
     asm volatile(
         "mov %cr3, %eax\n"
@@ -167,6 +167,11 @@ void MemoryManager::flushTLB()
     );
 }
 
+void MemoryManager::flushTLB(LinearAddress laddr)
+{
+    asm volatile("invlpg %0": :"m" (*(char*)laddr.get()));
+}
+
 bool MemoryManager::unmapRegion(Task& task, Task::Region& region)
 {
     auto& zone = *region.zone;
@@ -177,10 +182,9 @@ bool MemoryManager::unmapRegion(Task& task, Task::Region& region)
         pte.setPresent(false);
         pte.setWritable(false);
         pte.setUserAllowed(false);
-
+        flushTLB(laddr);
 //        kprintf("MM: >> Unmapped L%x => P%x <<\n", laddr, zone.m_pages[i].get());
     }
-    flushTLB();
     return true;
 }
 
@@ -203,10 +207,9 @@ bool MemoryManager::mapRegion(Task& task, Task::Region& region)
         pte.setPresent(true);
         pte.setWritable(true);
         pte.setUserAllowed(!task.isRing0());
-
+        flushTLB(laddr);
         //kprintf("MM: >> Mapped L%x => P%x <<\n", laddr, zone.m_pages[i].get());
     }
-    flushTLB();
     return true;
 }
 
diff --git a/Kernel/MemoryManager.h b/Kernel/MemoryManager.h
index 916ae45337..7b1a607914 100644
--- a/Kernel/MemoryManager.h
+++ b/Kernel/MemoryManager.h
@@ -60,7 +60,8 @@ private:
     ~MemoryManager();
 
     void initializePaging();
-    void flushTLB();
+    void flushEntireTLB();
+    void flushTLB(LinearAddress);
 
     void protectMap(LinearAddress, size_t length);
     void identityMap(LinearAddress, size_t length);