From 5dfa42665376dd42c36c2f44a2de72f932726595 Mon Sep 17 00:00:00 2001
From: Sarah <sarah.simon@rwth-aachen.de>
Date: Wed, 20 Sep 2023 16:38:37 +0200
Subject: [PATCH] shift making rootpage writable to paging.rs, adding function
 for re-setting WP bit after making rootpage writable

---
 src/arch/x86_64/kernel/processor.rs |  4 ----
 src/arch/x86_64/mm/paging.rs        | 34 ++++++++++++++++++++---------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/arch/x86_64/kernel/processor.rs b/src/arch/x86_64/kernel/processor.rs
index b786de6ec8..b2cba0b110 100644
--- a/src/arch/x86_64/kernel/processor.rs
+++ b/src/arch/x86_64/kernel/processor.rs
@@ -784,10 +784,6 @@ pub fn configure() {
 		wrmsr(IA32_EFER, rdmsr(IA32_EFER) | EFER_LMA | EFER_SCE | EFER_NXE);
 	}
 
-	// Check rootpagetable and make it writable if it isn't
-	// Necessary for memory manipulation
-	crate::arch::x86_64::mm::paging::check_root_pagetable();
-
 	//
 	// CR0 CONFIGURATION
 	//
diff --git a/src/arch/x86_64/mm/paging.rs b/src/arch/x86_64/mm/paging.rs
index a27879784e..563c89f3da 100644
--- a/src/arch/x86_64/mm/paging.rs
+++ b/src/arch/x86_64/mm/paging.rs
@@ -160,9 +160,9 @@ pub fn map<S>(
 					flush.flush();
 					debug!("Had to unmap page {page:?} before mapping.");
 				}
-				let result = pt.identity_map(frame, flags, &mut physicalmem::FrameAlloc);
-				// .unwrap()
-				// .flush();
+				pt.identity_map(frame, flags, &mut physicalmem::FrameAlloc)
+					.unwrap()
+					.flush();
 			}
 		}
 	}
@@ -256,7 +256,11 @@ pub fn get_application_page_size() -> usize {
 	LargePageSize::SIZE as usize
 }
 
-pub fn init() {}
+pub fn init() {
+	if crate::arch::x86_64::kernel::is_uefi().is_ok() {
+		check_root_pagetable();
+	}
+}
 
 pub fn init_page_tables() {
 	if env::is_uhyve() {
@@ -284,12 +288,12 @@ pub fn init_page_tables() {
 }
 
 /// Checks the address stored in the CR3 register and if necessary, makes its page writable.
-pub fn check_root_pagetable() {
+fn check_root_pagetable() {
 	use x86_64::structures::paging::mapper::TranslateResult;
 	let level_4_table_addr = Cr3::read().0.start_address().as_u64();
 	let virt_lvl_4_addr = x86_64::VirtAddr::new(level_4_table_addr);
 	let mut pt = unsafe { identity_mapped_page_table() };
-	set_cr0_flags();
+	clear_WP_bit();
 	match pt.translate(virt_lvl_4_addr) {
 		TranslateResult::Mapped {
 			frame,
@@ -309,21 +313,31 @@ pub fn check_root_pagetable() {
 		TranslateResult::NotMapped => todo!(),
 		TranslateResult::InvalidFrameAddress(_) => todo!(),
 	};
+	set_WP_bit();
 }
 
-/// Clears the WRITE_PROTECT bit in order to write into read-only Pages in supervisor mode
-fn set_cr0_flags() {
+/// Clears the WRITE_PROTECT bit in order to write into read-only Pages in supervisor mode.
+fn clear_WP_bit() {
 	let mut cr0 = Cr0::read();
 
-	//Clear WRITE_PROTECT in order to write in read-only pages in supervisor mode
 	if cr0.contains(Cr0Flags::WRITE_PROTECT) {
 		trace!("clear WRITE_PROTECT bit temporarily");
 		unsafe {
 			cr0.remove(Cr0Flags::WRITE_PROTECT);
 			Cr0::write(cr0);
 		}
-		trace!("Cr0 flags: {:?}", Cr0::read());
+		debug!("Cr0 flags: {:?}", Cr0::read());
+	}
+}
+
+/// Sets the WRITE_PROTECT bit in order to prevent writing into read-only Pages even in supervisor mode.
+fn set_WP_bit() {
+	let mut cr0 = Cr0::read();
+	unsafe {
+		cr0.insert(Cr0Flags::WRITE_PROTECT);
+		Cr0::write(cr0);
 	}
+	debug!("Cr0 flags: {:?}", Cr0::read());
 }
 
 /// This function takes the rootpage and depending on its size (1GiB, 2MiB, 4KiB), changes the flags to make it writable and then flushes the TLB.