From d4d14c78c304f1c3ae9628b5305ec3fa342c1409 Mon Sep 17 00:00:00 2001
From: Jeremy Soller <jeremy@system76.com>
Date: Wed, 12 Feb 2020 20:34:49 -0700
Subject: [PATCH] Ensure page table locking at runtime

---
 src/arch/x86_64/paging/mod.rs | 57 +++++++++++++++++++++++++++++++++--
 src/syscall/process.rs        |  7 ++++-
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/src/arch/x86_64/paging/mod.rs b/src/arch/x86_64/paging/mod.rs
index f5642fc2..2b4cce24 100644
--- a/src/arch/x86_64/paging/mod.rs
+++ b/src/arch/x86_64/paging/mod.rs
@@ -3,6 +3,7 @@
 
 use core::{mem, ptr};
 use core::ops::{Deref, DerefMut};
+use spin::Mutex;
 use x86::shared::{control_regs, msr, tlb};
 
 use crate::memory::{allocate_frames, Frame};
@@ -22,6 +23,39 @@ pub const ENTRY_COUNT: usize = 512;
 /// Size of pages
 pub const PAGE_SIZE: usize = 4096;
 
+//TODO: This is a rudimentary recursive mutex used to naively fix multi_core issues, replace it!
+pub struct PageTableLock {
+    cpu_id: usize,
+    count: usize,
+}
+
+pub static PAGE_TABLE_LOCK: Mutex<PageTableLock> = Mutex::new(PageTableLock {
+    cpu_id: 0,
+    count: 0,
+});
+
+fn page_table_lock() {
+    let cpu_id = crate::cpu_id();
+    loop {
+        {
+            let mut lock = PAGE_TABLE_LOCK.lock();
+            if lock.count == 0 || lock.cpu_id == cpu_id {
+                lock.cpu_id = cpu_id;
+                lock.count += 1;
+                return;
+            }
+        }
+        unsafe {
+            crate::arch::interrupt::pause();
+        }
+    }
+}
+
+fn page_table_unlock() {
+    let mut lock = PAGE_TABLE_LOCK.lock();
+    lock.count -= 1;
+}
+
 /// Setup page attribute table
 unsafe fn init_pat() {
     let uncacheable = 0;
@@ -108,7 +142,7 @@ pub unsafe fn init(cpu_id: usize, kernel_start: usize, kernel_end: usize, stack_
 
     init_pat();
 
-    let mut active_table = ActivePageTable::new();
+    let mut active_table = ActivePageTable::new_unlocked();
 
     let mut temporary_page = TemporaryPage::new(Page::containing_address(VirtualAddress::new(crate::USER_TMP_MISC_OFFSET)));
 
@@ -210,7 +244,7 @@ pub unsafe fn init_ap(cpu_id: usize, bsp_table: usize, stack_start: usize, stack
 
     init_pat();
 
-    let mut active_table = ActivePageTable::new();
+    let mut active_table = ActivePageTable::new_unlocked();
 
     let mut new_table = InactivePageTable::from_address(bsp_table);
 
@@ -259,6 +293,7 @@ pub unsafe fn init_ap(cpu_id: usize, bsp_table: usize, stack_start: usize, stack
 
 pub struct ActivePageTable {
     mapper: Mapper,
+    locked: bool,
 }
 
 impl Deref for ActivePageTable {
@@ -277,8 +312,17 @@ impl DerefMut for ActivePageTable {
 
 impl ActivePageTable {
     pub unsafe fn new() -> ActivePageTable {
+        page_table_lock();
         ActivePageTable {
             mapper: Mapper::new(),
+            locked: true,
+        }
+    }
+
+    pub unsafe fn new_unlocked() -> ActivePageTable {
+        ActivePageTable {
+            mapper: Mapper::new(),
+            locked: false,
         }
     }
 
@@ -331,6 +375,15 @@ impl ActivePageTable {
     }
 }
 
+impl Drop for ActivePageTable {
+    fn drop(&mut self) {
+        if self.locked {
+            page_table_unlock();
+            self.locked = false;
+        }
+    }
+}
+
 pub struct InactivePageTable {
     p4_frame: Frame,
 }
diff --git a/src/syscall/process.rs b/src/syscall/process.rs
index f06b6c82..25c37fee 100644
--- a/src/syscall/process.rs
+++ b/src/syscall/process.rs
@@ -380,7 +380,12 @@ pub fn clone(flags: CloneFlags, stack_base: usize) -> Result<ContextId> {
             context.sigmask = sigmask;
             context.umask = umask;
 
-            context.cpu_id = cpu_id_opt;
+            //TODO: Better CPU balancing
+            if let Some(cpu_id) = cpu_id_opt {
+                context.cpu_id = Some(cpu_id);
+            } else {
+                context.cpu_id = Some(pid.into() % crate::cpu_count());
+            }
 
             context.status = context::Status::Runnable;
 
-- 
GitLab