2 files changed, 19 insertions, 1 deletions
diff --git a/arch/x86_64/src/memory/paging/kernel_mapper.cpp b/arch/x86_64/src/memory/paging/kernel_mapper.cpp
index 59930fc..07d55ff 100644
--- a/arch/x86_64/src/memory/paging/kernel_mapper.cpp
+++ b/arch/x86_64/src/memory/paging/kernel_mapper.cpp
@@ -1,5 +1,8 @@
 #include "arch/memory/paging/kernel_mapper.hpp"
 
+#include "arch/memory/allocator/physical_frame.hpp"
+#include "arch/memory/paging/tlb.hpp"
+
 namespace teachos::arch::memory::paging
 {
   kernel_mapper::kernel_mapper(active_page_table & active_table)
@@ -11,10 +14,18 @@ namespace teachos::arch::memory::paging
   auto kernel_mapper::with(inactive_page_table inactive_page_table, temporary_page temporary_page,
                            active_page_table::function f) -> void
   {
+    auto backup = allocator::physical_frame::containing_address(PAGE_TABLE_LEVEL_4_ADDRESS);
+    auto page_table_level4 = temporary_page.map_table_frame(backup, active_table);
+
     active_table.active_handle[511].set_entry(inactive_page_table.page_table_level_4_frame,
                                               entry::PRESENT | entry::WRITABLE);
-    active_page_table::invalidate_page_cache(PAGE_TABLE_LEVEL_4_ADDRESS);
+    tlb_flush_all();
 
     f(active_table);
+
+    page_table_level4[511].set_entry(backup, entry::PRESENT | entry::WRITABLE);
+    tlb_flush_all();
+
+    temporary_page.unmap(active_table);
   }
 }  // namespace teachos::arch::memory::paging
diff --git a/arch/x86_64/src/memory/paging/tlb.cpp b/arch/x86_64/src/memory/paging/tlb.cpp
new file mode 100644
index 0000000..b3255e9
--- /dev/null
+++ b/arch/x86_64/src/memory/paging/tlb.cpp
@@ -0,0 +1,7 @@
+#include "arch/memory/paging/tlb.hpp"
+
+namespace teachos::arch::memory::paging
+{
+  auto tlb_flush(virtual_address address) -> void { asm volatile("invlpg (%0)" ::"r"(address) : "memory"); }
+  auto tlb_flush_all() -> void { flush(PAGE_TABLE_LEVEL_4_ADDRESS); }
+}  // namespace teachos::arch::memory::paging