From 5aa73d3bcff8a549176d4366c64ee8bc56de9e6d Mon Sep 17 00:00:00 2001
From: Ruihan Li <lrh2000@pku.edu.cn>
Date: Mon, 23 Dec 2024 23:09:19 +0800
Subject: [PATCH] Use `REF_COUNT_MAX` to avoid overflowing

---
 ostd/src/mm/page/meta.rs | 26 +++++++++++++++++++++++++-
 ostd/src/mm/page/mod.rs  | 12 ++++++++----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/ostd/src/mm/page/meta.rs b/ostd/src/mm/page/meta.rs
index 752965a6f..649a7b986 100644
--- a/ostd/src/mm/page/meta.rs
+++ b/ostd/src/mm/page/meta.rs
@@ -53,6 +53,7 @@ use crate::{
         kspace::LINEAR_MAPPING_BASE_VADDR, paddr_to_vaddr, page_size, page_table::boot_pt,
         CachePolicy, Paddr, PageFlags, PageProperty, PrivilegedPageFlags, Vaddr, PAGE_SIZE,
     },
+    panic::abort,
 };
 
 /// The maximum number of bytes of the metadata of a page.
@@ -81,7 +82,10 @@ pub(in crate::mm) struct MetaSlot {
     /// * `REF_COUNT_UNUSED`: The page is not in use.
     /// * `0`: The page is being constructed ([`Page::from_unused`])
     ///   or destructured ([`drop_last_in_place`]).
-    /// * `1..u32::MAX`: The page is in use.
+    /// * `1..REF_COUNT_MAX`: The page is in use.
+    /// * `REF_COUNT_MAX..REF_COUNT_UNUSED`: Illegal values to
+    ///   prevent the reference count from overflowing. Otherwise,
+    ///   overflowing the reference count will cause soundness issue.
     ///
     /// [`Page::from_unused`]: super::Page::from_unused
     pub(super) ref_count: AtomicU32,
@@ -90,6 +94,7 @@ pub(in crate::mm) struct MetaSlot {
 }
 
 pub(super) const REF_COUNT_UNUSED: u32 = u32::MAX;
+const REF_COUNT_MAX: u32 = i32::MAX as u32;
 
 type PageMetaVtablePtr = core::ptr::DynMetadata<dyn PageMeta>;
 
@@ -132,6 +137,25 @@ macro_rules! impl_page_meta {
 
 pub use impl_page_meta;
 
+impl MetaSlot {
+    /// Increases the page reference count by one.
+    ///
+    /// # Safety
+    ///
+    /// The caller must have already held a reference to the page.
+    pub(super) unsafe fn inc_ref_count(&self) {
+        let last_ref_cnt = self.ref_count.fetch_add(1, Ordering::Relaxed);
+        debug_assert!(last_ref_cnt != 0 && last_ref_cnt != REF_COUNT_UNUSED);
+
+        if last_ref_cnt >= REF_COUNT_MAX {
+            // This follows the same principle as the `Arc::clone` implementation to prevent the
+            // reference count from overflowing. See also
+            // <https://doc.rust-lang.org/std/sync/struct.Arc.html#method.clone>.
+            abort();
+        }
+    }
+}
+
 /// An internal routine in dropping implementations.
 ///
 /// # Safety
diff --git a/ostd/src/mm/page/mod.rs b/ostd/src/mm/page/mod.rs
index 0a261553a..0ab6374e9 100644
--- a/ostd/src/mm/page/mod.rs
+++ b/ostd/src/mm/page/mod.rs
@@ -189,7 +189,9 @@ impl<M: PageMeta> Page<M> {
 
 impl<M: PageMeta> Clone for Page<M> {
     fn clone(&self) -> Self {
-        self.slot().ref_count.fetch_add(1, Ordering::Relaxed);
+        // SAFETY: We have already held a reference to the page.
+        unsafe { self.slot().inc_ref_count() };
+
         Self {
             ptr: self.ptr,
             _marker: PhantomData,
@@ -331,7 +333,9 @@ impl From<Frame> for DynPage {
 
 impl Clone for DynPage {
     fn clone(&self) -> Self {
-        self.slot().ref_count.fetch_add(1, Ordering::Relaxed);
+        // SAFETY: We have already held a reference to the page.
+        unsafe { self.slot().inc_ref_count() };
+
         Self { ptr: self.ptr }
     }
 }
@@ -370,6 +374,6 @@ pub(in crate::mm) unsafe fn inc_page_ref_count(paddr: Paddr) {
     // an immutable reference to it is always safe.
     let slot = unsafe { &*(vaddr as *const MetaSlot) };
 
-    let old = slot.ref_count.fetch_add(1, Ordering::Relaxed);
-    debug_assert!(old > 0);
+    // SAFETY: We have already held a reference to the page.
+    unsafe { slot.inc_ref_count() };
 }