From bb7c9b6bfa2d5f521084a0a547907902386b5f83 Mon Sep 17 00:00:00 2001
From: Mathieu Strypsteen <mathieu@strypsteen.me>
Date: Fri, 27 Dec 2024 11:45:58 +0100
Subject: [PATCH] Refactor create_table_entry

---
 kernel/src/cpu/paging.rs | 54 ++++++++++++++++++++++++----------------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/kernel/src/cpu/paging.rs b/kernel/src/cpu/paging.rs
index 5ccb0bf..f0016fd 100644
--- a/kernel/src/cpu/paging.rs
+++ b/kernel/src/cpu/paging.rs
@@ -94,9 +94,12 @@ impl AddressSpace {
         let directory_i = virt_page / 512 % 512;
         let pdpt_i = virt_page / 512 / 512 % 512;
         let pml4_i = virt_page / 512 / 512 / 512 % 512;
-        let pdpt = get_table_entry(self.pml4.get_mut().unwrap(), pml4_i);
-        let directory = get_table_entry(pdpt, pdpt_i);
-        let table = get_table_entry(directory, directory_i);
+        self.create_table_entry(pml4_i as i32, -1, -1, true, true);
+        self.create_table_entry(pml4_i as i32, pdpt_i as i32, -1, true, true);
+        self.create_table_entry(pml4_i as i32, pdpt_i as i32, directory_i as i32, true, true);
+        let table = self.pml4.get_mut().unwrap().entries_virt[pml4_i].as_mut().unwrap().entries_virt[pdpt_i].as_mut().unwrap().entries_virt[directory_i]
+            .as_mut()
+            .unwrap();
         let should_invalidate = table.entries_phys[table_i].present() == 1;
         if !only_flags {
             table.entries_phys[table_i].set_address(phys / 0x1000);
@@ -173,6 +176,30 @@ impl AddressSpace {
             load_cr3(virt_to_phys(pml4_address));
         }
     }
+    fn create_table_entry(&mut self, pml4_i: i32, pdpt_i: i32, directory_i: i32, user: bool, exec: bool) {
+        let mut table = self.pml4.get_mut().unwrap();
+        let mut i = pml4_i as usize;
+        if pdpt_i != -1 {
+            table = table.entries_virt[i].as_mut().unwrap();
+            i = pdpt_i as usize;
+        }
+        if directory_i != -1 {
+            table = table.entries_virt[i].as_mut().unwrap();
+            i = directory_i as usize;
+        }
+        if table.entries_virt[i].is_none() {
+            let new_table = Box::new(PageTable {
+                entries_phys: [PageEntry(0); 512],
+                entries_virt: [const { None }; 512],
+            });
+            table.entries_phys[i].set_address(virt_to_phys(new_table.as_ref() as *const PageTable as u64) / 0x1000);
+            table.entries_virt[i] = Some(new_table);
+            table.entries_phys[i].set_write(1);
+            table.entries_phys[i].set_user(user as u64);
+            table.entries_phys[i].set_execute_disable(!exec as u64);
+            table.entries_phys[i].set_present(1);
+        }
+    }
 }
 
 fn get_free_frame() -> u64 {
@@ -200,20 +227,6 @@ pub fn virt_to_phys(virt: u64) -> u64 {
     let heap_map = HEAP_PHYS_MAPPING.lock();
     heap_map[(virt as usize - KERNEL_HEAP_START as usize) / 0x1000] + virt % 0x1000
 }
-fn get_table_entry(table: &mut PageTable, i: usize) -> &mut PageTable {
-    if table.entries_virt[i].is_none() {
-        let new_table = Box::new(PageTable {
-            entries_phys: [PageEntry(0); 512],
-            entries_virt: [const { None }; 512],
-        });
-        table.entries_phys[i].set_address(virt_to_phys(new_table.as_ref() as *const PageTable as u64) / 0x1000);
-        table.entries_virt[i] = Some(new_table);
-        table.entries_phys[i].set_write(1);
-        table.entries_phys[i].set_user(1);
-        table.entries_phys[i].set_present(1);
-    }
-    table.entries_virt[i].as_mut().unwrap()
-}
 pub fn setup_paging(loader_struct: &LoaderStruct, phys_start: u64, heap_start: u64) {
     HEAP_PHYS_START.store(heap_start, Ordering::SeqCst);
     let mut memory_size = 0;
@@ -264,12 +277,11 @@ pub fn setup_paging(loader_struct: &LoaderStruct, phys_start: u64, heap_start: u
             entries_virt: [const { None }; 512],
         }))
         .unwrap_or_else(|_| panic!());
-    let pml4 = address_space.pml4.get_mut().unwrap();
     for i in 256..512 {
-        get_table_entry(pml4, i);
-        pml4.entries_phys[i].set_user(0);
         if i != 511 {
-            pml4.entries_phys[i].set_execute_disable(1);
+            address_space.create_table_entry(i as i32, -1, -1, false, false);
+        } else {
+            address_space.create_table_entry(i as i32, -1, -1, false, true);
         }
     }
     for i in text_start..text_end {