Refactor paging

kernel/src/cpu/paging.rs

@@ -28,6 +28,8 @@ extern "C" {
 }
 
 pub struct AddressSpace {
+    pub is_kernel: bool,
+    pub pml4_physical_address: u64,
     pub pml4: OnceCell<Box<PageTable>>,
 }
 
@@ -41,6 +43,7 @@ pub const USER_END: u64 = 0x800000000000 - 0x1000; // Avoid accidentally jumping
 
 impl AddressSpace {
     pub fn get_page(&mut self, virt: u64) -> Option<&mut PageEntry> {
+        assert!(virt < USER_END || self.is_kernel);
         let virt_page = virt as usize / 0x1000;
         let table_i = virt_page % 512;
         let directory_i = virt_page / 512 % 512;
@@ -84,6 +87,7 @@ impl AddressSpace {
         assert!(virt >= 0x1000, "First page shouldn't be mapped");
         assert!(!write || !exec || virt == 0x1000);
         assert!(user == (virt < USER_END) || virt == 0x1000);
+        assert!(virt < USER_END || self.is_kernel);
         {
             let mut frames_vec = PHYSICAL_FRAMES.lock();
             let frame = phys as usize / 0x1000;
@@ -96,9 +100,9 @@ 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;
-        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);
+        self.create_table_entry(pml4_i as i32, -1, -1, true, true, false);
+        self.create_table_entry(pml4_i as i32, pdpt_i as i32, -1, true, true, false);
+        self.create_table_entry(pml4_i as i32, pdpt_i as i32, directory_i as i32, true, true, false);
         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();
@@ -173,12 +177,18 @@ impl AddressSpace {
         }
     }
     pub unsafe fn switch(&self) {
-        let pml4_address = self.pml4.get().unwrap().as_ref() as *const PageTable as u64;
         unsafe {
-            load_cr3(virt_to_phys(pml4_address));
+            load_cr3(self.pml4_physical_address);
         }
     }
-    fn create_table_entry(&mut self, pml4_i: i32, pdpt_i: i32, directory_i: i32, user: bool, exec: bool) {
+    fn create_table_entry(&mut self, pml4_i: i32, pdpt_i: i32, directory_i: i32, user: bool, exec: bool, force_create: bool) {
+        let mut new_table = None;
+        let mut new_table_phys = 0;
+        if force_create {
+            let _new_table = create_page_table();
+            new_table_phys = self.virt_to_phys(_new_table.as_ref() as *const PageTable as u64) / 0x1000;
+            new_table = Some(_new_table);
+        }
         let mut table = self.pml4.get_mut().unwrap();
         let mut i = pml4_i as usize;
         if pdpt_i != -1 {
@@ -189,21 +199,43 @@ impl AddressSpace {
             table = table.entries_virt[i].as_mut().unwrap();
             i = directory_i as usize;
         }
+        let mut retry = false;
         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);
+            if force_create {
+                table.entries_phys[i].set_address(new_table_phys);
+                table.entries_virt[i] = 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);
+            } else {
+                retry = true;
+            }
         }
+        if retry {
+            self.create_table_entry(pml4_i, pdpt_i, directory_i, user, exec, true);
+        }
+    }
+    pub fn virt_to_phys(&mut self, virt: u64) -> u64 {
+        if !self.is_kernel {
+            let kernel_proc = get_kernel_process();
+            return kernel_proc.lock().address_space.virt_to_phys(virt);
+        }
+        if !PAGING_ACTIVE.load(Ordering::SeqCst) {
+            assert!(virt >= KERNEL_HEAP_START);
+            assert!(virt < KERNEL_HEAP_START + KERNEL_HEAP_INITIAL_SIZE as u64);
+            return virt - KERNEL_HEAP_START + HEAP_PHYS_START.load(Ordering::SeqCst);
+        }
+        self.get_page(virt).unwrap().address() * 0x1000 + virt % 0x1000
     }
 }
 
+pub fn create_page_table() -> Box<PageTable> {
+    return Box::new(PageTable {
+        entries_phys: [PageEntry(0); 512],
+        entries_virt: [const { None }; 512],
+    });
+}
 fn get_free_frame() -> u64 {
     let frames_vec = PHYSICAL_FRAMES.lock();
     let frames_vec = frames_vec.get().unwrap();
@@ -220,15 +252,6 @@ fn invlpg(addr: u64) {
     }
     smp_invalidate_tlb();
 }
-pub fn virt_to_phys(virt: u64) -> u64 {
-    assert!(virt >= KERNEL_HEAP_START);
-    assert!(virt < KERNEL_HEAP_START + KERNEL_HEAP_INITIAL_SIZE as u64);
-    if !PAGING_ACTIVE.load(Ordering::SeqCst) {
-        return virt - KERNEL_HEAP_START + HEAP_PHYS_START.load(Ordering::SeqCst);
-    }
-    let heap_map = HEAP_PHYS_MAPPING.lock();
-    heap_map[(virt as usize - KERNEL_HEAP_START as usize) / 0x1000] + virt % 0x1000
-}
 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;
@@ -274,18 +297,14 @@ pub fn setup_paging(loader_struct: &LoaderStruct, phys_start: u64, heap_start: u
     let kernel_proc = get_kernel_process();
     let mut kernel_proc = kernel_proc.lock();
     let address_space = &mut kernel_proc.address_space;
-    address_space
-        .pml4
-        .set(Box::new(PageTable {
-            entries_phys: [PageEntry(0); 512],
-            entries_virt: [const { None }; 512],
-        }))
-        .unwrap_or_else(|_| panic!());
+    let pml4 = create_page_table();
+    address_space.pml4_physical_address = address_space.virt_to_phys(pml4.as_ref() as *const PageTable as u64);
+    address_space.pml4.set(pml4).unwrap_or_else(|_| panic!());
     for i in 256..512 {
-        if i != 511 {
-            address_space.create_table_entry(i as i32, -1, -1, false, false);
+        if i == 511 {
+            address_space.create_table_entry(i as i32, -1, -1, false, true, false);
         } else {
-            address_space.create_table_entry(i as i32, -1, -1, false, true);
+            address_space.create_table_entry(i as i32, -1, -1, false, false, false);
         }
     }
     for i in text_start..text_end {

kernel/src/sys/process.rs

@@ -1,9 +1,13 @@
 use core::cell::{LazyCell, OnceCell};
 
-use alloc::{boxed::Box, sync::Arc};
-use kernel_common::paging::{PageEntry, PageTable};
+use alloc::sync::Arc;
+use kernel_common::paging::PageTable;
 
-use crate::{cpu::paging::AddressSpace, misc::elf::load_binary, INIT_BINARY};
+use crate::{
+    cpu::paging::{create_page_table, AddressSpace},
+    misc::elf::load_binary,
+    INIT_BINARY,
+};
 
 use super::{locks::Spinlock, scheduler::SCHEDULER, task::create_task};
 
@@ -14,7 +18,11 @@ pub struct Process {
 
 static KERNEL_PROCESS: Spinlock<LazyCell<Arc<Spinlock<Process>>>> = Spinlock::new(LazyCell::new(|| {
     let process = Process {
-        address_space: AddressSpace { pml4: OnceCell::new() },
+        address_space: AddressSpace {
+            is_kernel: true,
+            pml4: OnceCell::new(),
+            pml4_physical_address: 0,
+        },
         binary: None,
     };
     Arc::new(Spinlock::new(process))
@@ -24,20 +32,22 @@ pub fn get_kernel_process() -> Arc<Spinlock<Process>> {
     KERNEL_PROCESS.lock().clone()
 }
 pub fn create_process(binary: Option<&'static [u8]>) -> Arc<Spinlock<Process>> {
-    let process = Process {
-        address_space: AddressSpace { pml4: OnceCell::new() },
+    let mut process = Process {
+        address_space: AddressSpace {
+            is_kernel: false,
+            pml4: OnceCell::new(),
+            pml4_physical_address: 0,
+        },
         binary,
     };
-    let mut pml4 = Box::new(PageTable {
-        entries_phys: [PageEntry(0); 512],
-        entries_virt: [const { None }; 512],
-    });
+    let mut pml4 = create_page_table();
     let kernel_proc = get_kernel_process();
-    let kernel_proc = kernel_proc.lock();
+    let mut kernel_proc = kernel_proc.lock();
     let kernel_pml4 = kernel_proc.address_space.pml4.get().unwrap();
     for i in 256..512 {
         pml4.entries_phys[i] = kernel_pml4.entries_phys[i];
     }
+    process.address_space.pml4_physical_address = kernel_proc.address_space.virt_to_phys(pml4.as_ref() as *const PageTable as u64);
     process.address_space.pml4.set(pml4).unwrap_or_else(|_| panic!());
     Arc::new(Spinlock::new(process))
 }

kernel/src/sys/smp.rs

@@ -3,10 +3,7 @@ use core::{ptr::copy, sync::atomic::Ordering};
 use alloc::{vec, vec::Vec};
 use kernel_common::{instructions::pause, paging::PageTable};
 
-use crate::{
-    cpu::{isr::ISR_INVALIDATE_TLB, paging::virt_to_phys},
-    BROADCASTED_PANIC,
-};
+use crate::{cpu::isr::ISR_INVALIDATE_TLB, BROADCASTED_PANIC};
 
 use super::{
     hpet::sleep,
@@ -36,7 +33,7 @@ pub fn start_aps() {
         let mut kernel_proc = kernel_proc.lock();
         let address_space = &mut kernel_proc.address_space;
         let pml4 = &address_space.pml4;
-        let pml4_phys_addr_u64 = virt_to_phys(pml4.get().unwrap().as_ref() as *const PageTable as u64);
+        let pml4_phys_addr_u64 = address_space.virt_to_phys(pml4.get().unwrap().as_ref() as *const PageTable as u64);
         pml4_phys_addr = u32::try_from(pml4_phys_addr_u64).unwrap();
         unsafe {
             address_space.map_range(0x1000, 0x1000, 0x1000, false, true, false, false);