From 862b64b786e236df1166c2b48845f922587cd56e Mon Sep 17 00:00:00 2001
From: Zhang Junyang <junyang@stu.pku.edu.cn>
Date: Tue, 23 Jul 2024 03:26:14 +0000
Subject: [PATCH] Sanitize the mutability of page table metadata

---
 ostd/src/mm/page/meta.rs       | 20 +++++++-------------
 ostd/src/mm/page_table/node.rs | 34 ++++++++++++++++++----------------
 2 files changed, 25 insertions(+), 29 deletions(-)

diff --git a/ostd/src/mm/page/meta.rs b/ostd/src/mm/page/meta.rs
index e1b84fda6..fc10f4c19 100644
--- a/ostd/src/mm/page/meta.rs
+++ b/ostd/src/mm/page/meta.rs
@@ -175,13 +175,6 @@ pub struct FrameMeta {
 
 impl Sealed for FrameMeta {}
 
-#[derive(Debug, Default)]
-#[repr(C)]
-pub struct PageTablePageMetaInner {
-    pub level: PagingLevel,
-    pub nr_children: u16,
-}
-
 /// The metadata of any kinds of page table pages.
 /// Make sure the the generic parameters don't effect the memory layout.
 #[derive(Debug)]
@@ -192,9 +185,12 @@ pub struct PageTablePageMeta<
 > where
     [(); C::NR_LEVELS as usize]:,
 {
+    pub level: PagingLevel,
+    /// The lock for the page table page.
     pub lock: AtomicU8,
-    pub inner: UnsafeCell<PageTablePageMetaInner>,
-    _phantom: PhantomData<(E, C)>,
+    /// The number of valid PTEs. It is mutable if the lock is held.
+    pub nr_children: UnsafeCell<u16>,
+    _phantom: core::marker::PhantomData<(E, C)>,
 }
 
 impl<E: PageTableEntryTrait, C: PagingConstsTrait> PageTablePageMeta<E, C>
@@ -203,11 +199,9 @@ where
 {
     pub fn new_locked(level: PagingLevel) -> Self {
         Self {
+            level,
             lock: AtomicU8::new(1),
-            inner: UnsafeCell::new(PageTablePageMetaInner {
-                level,
-                nr_children: 0,
-            }),
+            nr_children: UnsafeCell::new(0),
             _phantom: PhantomData,
         }
     }
diff --git a/ostd/src/mm/page_table/node.rs b/ostd/src/mm/page_table/node.rs
index 191f569a2..9b489ba17 100644
--- a/ostd/src/mm/page_table/node.rs
+++ b/ostd/src/mm/page_table/node.rs
@@ -36,7 +36,7 @@ use crate::{
         paddr_to_vaddr,
         page::{
             self,
-            meta::{PageMeta, PageTablePageMeta, PageTablePageMetaInner, PageUsage},
+            meta::{PageMeta, PageTablePageMeta, PageUsage},
             DynPage, Page,
         },
         page_prop::PageProperty,
@@ -462,8 +462,16 @@ where
         unsafe { self.as_ptr().add(idx).read() }
     }
 
-    fn start_paddr(&self) -> Paddr {
-        self.page.paddr()
+    /// The number of valid PTEs.
+    pub(super) fn nr_children(&self) -> u16 {
+        // SAFETY: The lock is held so there is no mutable reference to it.
+        // It would be safe to read.
+        unsafe { *self.meta().nr_children.get() }
+    }
+
+    fn nr_children_mut(&mut self) -> &mut u16 {
+        // SAFETY: The lock is held so we have an exclusive access.
+        unsafe { &mut *self.meta().nr_children.get() }
     }
 
     /// Replaces a page table entry at a given index.
@@ -504,13 +512,13 @@ where
 
             // Update the child count.
             if pte.is_none() {
-                self.meta_mut().nr_children -= 1;
+                *self.nr_children_mut() -= 1;
             }
         } else if let Some(e) = pte {
             // SAFETY: This is safe as described in the above branch.
             unsafe { (self.as_ptr() as *mut E).add(idx).write(e) };
 
-            self.meta_mut().nr_children += 1;
+            *self.nr_children_mut() += 1;
         }
     }
 
@@ -518,16 +526,12 @@ where
         paddr_to_vaddr(self.start_paddr()) as *const E
     }
 
-    fn meta(&self) -> &PageTablePageMetaInner {
-        // SAFETY: We have exclusively locked the page, so we can derive an immutable reference
-        // from an immutable reference to the lock guard.
-        unsafe { &*self.page.meta().inner.get() }
+    fn start_paddr(&self) -> Paddr {
+        self.page.paddr()
     }
 
-    fn meta_mut(&mut self) -> &mut PageTablePageMetaInner {
-        // SAFETY: We have exclusively locked the page, so we can derive a mutable reference from a
-        // mutable reference to the lock guard.
-        unsafe { &mut *self.page.meta().inner.get() }
+    fn meta(&self) -> &PageTablePageMeta<E, C> {
+        self.page.meta()
     }
 }
 
@@ -549,9 +553,7 @@ where
 
     fn on_drop(page: &mut Page<Self>) {
         let paddr = page.paddr();
-
-        let inner = unsafe { &mut *page.meta().inner.get() };
-        let level = inner.level;
+        let level = page.meta().level;
 
         // Drop the children.
         for i in 0..nr_subpage_per_huge::<C>() {