From 7d00c2670fb43f3610156ce76e5c19a9515aaad0 Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Thu, 13 Jul 2023 14:28:53 +0300
Subject: [PATCH] Kernel: Support execute disable bit

We will now map executable memory explicitly as executable.
---
 kernel/arch/x86_64/PageTable.cpp         | 36 ++++++++++++++++++++++--
 kernel/arch/x86_64/linker.ld             |  2 ++
 kernel/include/kernel/Memory/PageTable.h |  1 +
 kernel/kernel/Process.cpp                |  2 ++
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/kernel/arch/x86_64/PageTable.cpp b/kernel/arch/x86_64/PageTable.cpp
index 161c4424c..26a2c9390 100644
--- a/kernel/arch/x86_64/PageTable.cpp
+++ b/kernel/arch/x86_64/PageTable.cpp
@@ -1,5 +1,6 @@
 #include <BAN/Errors.h>
 #include <kernel/Arch.h>
+#include <kernel/CPUID.h>
 #include <kernel/LockGuard.h>
 #include <kernel/Memory/kmalloc.h>
 #include <kernel/Memory/PageTable.h>
@@ -7,11 +8,15 @@
 extern uint8_t g_kernel_start[];
 extern uint8_t g_kernel_end[];
 
+extern uint8_t g_kernel_execute_start[];
+extern uint8_t g_kernel_execute_end[];
+
 namespace Kernel
 {
 	
 	static PageTable* s_kernel = nullptr;
 	static PageTable* s_current = nullptr;
+	static bool s_has_nxe = false;
 
 	// Page Directories for kernel memory (KERNEL_OFFSET -> 0xFFFFFFFFFFFFFFFF)
 	static paddr_t s_global[(0xFFFFFFFFFFFFFFFF - KERNEL_OFFSET + 1) / (4096ull * 512ull * 512ull)] { };
@@ -40,6 +45,17 @@ namespace Kernel
 
 	void PageTable::initialize()
 	{
+		if (CPUID::has_nxe())
+		{
+			asm volatile(
+				"movl $0xC0000080, %ecx;"
+				"rdmsr;"
+				"orl $0x800, %eax;"
+				"wrmsr"
+			);
+			s_has_nxe = true;
+		}
+
 		ASSERT(s_kernel == nullptr);
 		s_kernel = new PageTable();
 		ASSERT(s_kernel);
@@ -78,6 +94,14 @@ namespace Kernel
 
 		// Map (0 -> phys_kernel_end) to (KERNEL_OFFSET -> virt_kernel_end)
 		map_range_at(0, KERNEL_OFFSET, (uintptr_t)g_kernel_end - KERNEL_OFFSET, Flags::ReadWrite | Flags::Present);
+		
+		// Map executable kernel memory as executable
+		map_range_at(
+			V2P(g_kernel_execute_start),
+			(vaddr_t)g_kernel_execute_start,
+			g_kernel_execute_end - g_kernel_execute_start,
+			Flags::Execute | Flags::ReadWrite | Flags::Present
+		);
 	}
 
 	BAN::ErrorOr<PageTable*> PageTable::create_userspace()
@@ -207,6 +231,10 @@ namespace Kernel
 	{
 		LockGuard _(m_lock);
 
+		uint64_t xd_bit = 0;
+		if (s_has_nxe && !(flags & Flags::Execute))
+			xd_bit = 1ull << 63;
+
 		if (vaddr && (vaddr >= KERNEL_OFFSET) != (this == s_kernel))
 			Kernel::panic("mapping {8H} to {8H}, kernel: {}", paddr, vaddr, this == s_kernel);
 
@@ -248,7 +276,7 @@ namespace Kernel
 		}
 
 		uint64_t* pt = (uint64_t*)P2V(pd[pde] & PAGE_ADDR_MASK);
-		pt[pte] = paddr | flags;
+		pt[pte] = xd_bit | paddr | flags;
 
 		invalidate(canonicalize(vaddr));
 	}
@@ -304,12 +332,14 @@ namespace Kernel
 
 	PageTable::flags_t PageTable::get_page_flags(vaddr_t addr) const
 	{
-		return get_page_data(addr) & PAGE_FLAG_MASK;
+		uint64_t page_data = get_page_data(addr);
+		return (page_data & (1ull << 63) ? Flags::Execute : 0) | (page_data & PAGE_FLAG_MASK);
 	}
 
 	paddr_t PageTable::physical_address_of(vaddr_t addr) const
 	{
-		return get_page_data(addr) & PAGE_ADDR_MASK;
+		uint64_t page_data = get_page_data(addr);
+		return (page_data & PAGE_ADDR_MASK) & ~(1ull << 63);
 	}
 
 	vaddr_t PageTable::get_free_page(vaddr_t first_address) const
diff --git a/kernel/arch/x86_64/linker.ld b/kernel/arch/x86_64/linker.ld
index b6b1fa90d..4e37c00c0 100644
--- a/kernel/arch/x86_64/linker.ld
+++ b/kernel/arch/x86_64/linker.ld
@@ -9,12 +9,14 @@ SECTIONS
 	g_kernel_start = .;
 	.text ALIGN(4K) : AT(ADDR(.text) - KERNEL_OFFSET)
 	{
+		g_kernel_execute_start = .;
 		*(.multiboot)
 		*(.text)
 	}
 	.rodata ALIGN(4K) : AT(ADDR(.rodata) - KERNEL_OFFSET)
 	{
 		*(.rodata.*)
+		g_kernel_execute_end = .;
 	}
 	.data ALIGN(4K) : AT(ADDR(.data) - KERNEL_OFFSET)
 	{
diff --git a/kernel/include/kernel/Memory/PageTable.h b/kernel/include/kernel/Memory/PageTable.h
index b936f4a2c..a1e2bd880 100644
--- a/kernel/include/kernel/Memory/PageTable.h
+++ b/kernel/include/kernel/Memory/PageTable.h
@@ -16,6 +16,7 @@ namespace Kernel
 			Present = 1,
 			ReadWrite = 2,
 			UserSupervisor = 4,
+			Execute = 8,
 		};
 
 	public:
diff --git a/kernel/kernel/Process.cpp b/kernel/kernel/Process.cpp
index a45cafbc8..3b6bf7de5 100644
--- a/kernel/kernel/Process.cpp
+++ b/kernel/kernel/Process.cpp
@@ -458,6 +458,8 @@ namespace Kernel
 				uint8_t flags = PageTable::Flags::UserSupervisor | PageTable::Flags::Present;
 				if (elf_program_header.p_flags & LibELF::PF_W)
 					flags |= PageTable::Flags::ReadWrite;
+				if (elf_program_header.p_flags & LibELF::PF_X)
+					flags |= PageTable::Flags::Execute;
 
 				size_t page_start = elf_program_header.p_vaddr / PAGE_SIZE;
 				size_t page_end = BAN::Math::div_round_up<size_t>(elf_program_header.p_vaddr + elf_program_header.p_memsz, PAGE_SIZE);