From 0cb53efa01bf60f41c765bbaf7e1190d85ac9183 Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Fri, 5 May 2023 14:18:18 +0300
Subject: [PATCH] Kernel: 64 bit MMU now properly identity maps kernel

---
 kernel/arch/x86_64/MMU.cpp         | 34 +++++++-----------------------
 kernel/include/kernel/Memory/MMU.h | 18 ++++++++++------
 2 files changed, 19 insertions(+), 33 deletions(-)

diff --git a/kernel/arch/x86_64/MMU.cpp b/kernel/arch/x86_64/MMU.cpp
index 8fded3d77..c16333067 100644
--- a/kernel/arch/x86_64/MMU.cpp
+++ b/kernel/arch/x86_64/MMU.cpp
@@ -44,30 +44,12 @@ namespace Kernel
 
 	void MMU::initialize_kernel()
 	{
-		// FIXME: We should just identity map until g_kernel_end
-
-		ASSERT((uintptr_t)g_kernel_end <= 6 * (1 << 20));
-
-		// Identity map from 0 -> 6 MiB
 		m_highest_paging_struct = allocate_page_aligned_page();
-		
-		uint64_t* pdpt = allocate_page_aligned_page();
-		m_highest_paging_struct[0] = (uint64_t)pdpt | Flags::ReadWrite | Flags::Present;
+		memset(m_highest_paging_struct, 0, PAGE_SIZE);
 
-		uint64_t* pd = allocate_page_aligned_page();
-		pdpt[0] = (uint64_t)pd | Flags::ReadWrite | Flags::Present;
-
-		for (uint32_t i = 0; i < 3; i++)
-		{
-			uint64_t* pt = allocate_page_aligned_page();
-			for (uint64_t j = 0; j < 512; j++)
-				pt[j] = (i << 21) | (j << 12) | Flags::ReadWrite | Flags::Present;
-			pd[i] = (uint64_t)pt | Flags::ReadWrite | Flags::Present;
-		}
-
-		// Unmap 0 -> 4 KiB
-		uint64_t* pt1 = (uint64_t*)(pd[0] & PAGE_MASK);
-		pt1[0] = 0;
+		// Identity map 4 KiB -> kernel end. We don't map the first page since nullptr derefs should
+		// page fault. Also there isn't anything useful in that memory.
+		identity_map_range(PAGE_SIZE, (uintptr_t)g_kernel_end, Flags::ReadWrite | Flags::Present);
 	}
 
 	MMU::MMU()
@@ -150,13 +132,13 @@ namespace Kernel
 		asm volatile("movq %0, %%cr3" :: "r"(m_highest_paging_struct));
 	}
 
-	void MMU::identity_map_page(paddr_t address, uint8_t flags)
+	void MMU::identity_map_page(paddr_t address, flags_t flags)
 	{
 		address &= PAGE_MASK;
 		map_page_at(address, address, flags);
 	}
 
-	void MMU::identity_map_range(paddr_t address, ptrdiff_t size, uint8_t flags)
+	void MMU::identity_map_range(paddr_t address, size_t size, flags_t flags)
 	{
 		paddr_t s_page = address & PAGE_MASK;
 		paddr_t e_page = (address + size - 1) & PAGE_MASK;
@@ -201,7 +183,7 @@ namespace Kernel
 		pml4[pml4e] = 0;
 	}
 
-	void MMU::unmap_range(vaddr_t address, ptrdiff_t size)
+	void MMU::unmap_range(vaddr_t address, size_t size)
 	{
 		vaddr_t s_page = address & PAGE_MASK;
 		vaddr_t e_page = (address + size - 1) & PAGE_MASK;
@@ -209,7 +191,7 @@ namespace Kernel
 			unmap_page(page);
 	}
 
-	void MMU::map_page_at(paddr_t paddr, vaddr_t vaddr, uint8_t flags)
+	void MMU::map_page_at(paddr_t paddr, vaddr_t vaddr, flags_t flags)
 	{
 		ASSERT((paddr >> 48) == 0);
 		ASSERT((vaddr >> 48) == 0);
diff --git a/kernel/include/kernel/Memory/MMU.h b/kernel/include/kernel/Memory/MMU.h
index 41aece4c1..7a0d1a615 100644
--- a/kernel/include/kernel/Memory/MMU.h
+++ b/kernel/include/kernel/Memory/MMU.h
@@ -4,11 +4,12 @@
 
 namespace Kernel
 {
-	
+
 	class MMU
 	{
 	public:
-		enum Flags : uint8_t 
+		using flags_t = uint8_t;
+		enum Flags : flags_t 
 		{
 			Present = 1,
 			ReadWrite = 2,
@@ -22,15 +23,18 @@ namespace Kernel
 		MMU();
 		~MMU();
 
-		void identity_map_page(paddr_t, uint8_t);
-		void identity_map_range(paddr_t, ptrdiff_t, uint8_t);
+		void identity_map_page(paddr_t, flags_t);
+		void identity_map_range(paddr_t, size_t, flags_t);
 
 		void unmap_page(vaddr_t);
-		void unmap_range(vaddr_t, ptrdiff_t);
+		void unmap_range(vaddr_t, size_t);
 
-		void map_page_at(paddr_t, vaddr_t, uint8_t);
+		void map_page_at(paddr_t, vaddr_t, flags_t);
 
-		uint8_t get_page_flags(vaddr_t) const;
+		flags_t get_page_flags(vaddr_t) const;
+
+		vaddr_t get_free_page() const;
+		vaddr_t get_free_contiguous_pages(uint32_t) const;
 
 		void load();