5 files changed, 51 insertions, 18 deletions
diff --git a/arch/x86_64/CMakeLists.txt b/arch/x86_64/CMakeLists.txt
index 22f9ad4..1f15c79 100644
--- a/arch/x86_64/CMakeLists.txt
+++ b/arch/x86_64/CMakeLists.txt
@@ -53,6 +53,7 @@ target_sources("_memory" PRIVATE
   "src/memory/paging/active_page_table.cpp"
   "src/memory/paging/inactive_page_table.cpp"
   "src/memory/paging/kernel_mapper.cpp"
+  "src/memory/paging/tlb.cpp"
 )
 
 #[============================================================================[
diff --git a/arch/x86_64/include/arch/memory/paging/active_page_table.hpp b/arch/x86_64/include/arch/memory/paging/active_page_table.hpp
index e2c205b..30ae455 100644
--- a/arch/x86_64/include/arch/memory/paging/active_page_table.hpp
+++ b/arch/x86_64/include/arch/memory/paging/active_page_table.hpp
@@ -3,6 +3,7 @@
 
 #include "arch/exception_handling/assert.hpp"
 #include "arch/memory/allocator/concept.hpp"
+#include "arch/memory/paging/tlb.hpp"
 #include "arch/memory/paging/virtual_page.hpp"
 
 #include <array>
@@ -10,8 +11,6 @@
 
 namespace teachos::arch::memory::paging
 {
-  std::size_t constexpr PAGE_TABLE_LEVEL_4_ADDRESS = 0xffffffff'fffff000;
-
   /**
    * @brief Currently active level 4 page table, is used to ensure there is only ever one valid instance and it cannot
    * be copied or constructed again.
@@ -55,18 +54,6 @@ namespace teachos::arch::memory::paging
     auto translate_huge_page(virtual_page page) -> std::optional<allocator::physical_frame>;
 
     /**
-     * @brief Invalidates any translation lookaside buffer (TLB) entry for the page table the given address is cotained
-     * in. See https://www.felixcloutier.com/x86/invlpg for more information on the used x86 instruction.
-     *
-     * @param address Memory address, which will be used to determine the contained page and flush the TLB entry for
-     * that page.
-     */
-    static auto invalidate_page_cache(virtual_address address) -> void
-    {
-      asm volatile("invlpg (%0)" ::"r"(address) : "memory");
-    }
-
-    /**
      * @brief Maps a virtual page to a physical frame in the page table with the specified flags.
      *
      * @note Allocates and maps an entry in every page level if it does not exists yet down to level 1. If the level 1
@@ -171,9 +158,12 @@ namespace teachos::arch::memory::paging
           break;
         }
       }
-      invalidate_page_cache(page.start_address());
+
+      tlb_flush(page.start_address());
     }
 
+    page_table_handle active_handle;
+
   private:
     /**
      * @brief Private constructor should only be used by create or get method, which ensures to create only ever one
@@ -214,8 +204,6 @@ namespace teachos::arch::memory::paging
       entry.set_unused();
       allocator.deallocate_frame(frame.value());
     }
-
-    page_table_handle active_handle;
   };
 
 }  // namespace teachos::arch::memory::paging
diff --git a/arch/x86_64/include/arch/memory/paging/tlb.hpp b/arch/x86_64/include/arch/memory/paging/tlb.hpp
new file mode 100644
index 0000000..e74eb74
--- /dev/null
+++ b/arch/x86_64/include/arch/memory/paging/tlb.hpp
@@ -0,0 +1,26 @@
+#ifndef TEACHOS_ARCH_X86_64_MEMORY_PAGING_TLB_HPP
+#define TEACHOS_ARCH_X86_64_MEMORY_PAGING_TLB_HPP
+
+#include "arch/memory/paging/virtual_page.hpp"
+
+namespace teachos::arch::memory::paging
+{
+  std::size_t constexpr PAGE_TABLE_LEVEL_4_ADDRESS = 0xffffffff'fffff000;
+
+  /**
+   * @brief Invalidates any translation lookaside buffer (TLB) entry for the page table the given address is cotained
+   * in. See https://www.felixcloutier.com/x86/invlpg for more information on the used x86 instruction.
+   *
+   * @param address Memory address, which will be used to determine the contained page and flush the TLB entry for
+   * that page.
+   */
+  auto tlb_flush(virtual_address address) -> void;
+
+  /**
+   * @brief Invalidates the translation lookaside buffer (TLB) entry for all page tables
+   */
+  auto tlb_flush_all() -> void;
+
+}  // namespace teachos::arch::memory::paging
+
+#endif
+\ No newline at end of file
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