[PATCH v3 1/2] rust: page: use the page's reference count to decide when to free the allocation

[Abdiel Janulgue <abdiel.janulgue@gmail.com>]

Ensure pages returned by the constructor are always reference counted.
This requires that we replace the page pointer wrapper with Opaque instead
of NonNull to make it possible to cast to a Page pointer from a raw struct
page pointer which is needed to create an ARef instance.

Signed-off-by: Abdiel Janulgue <abdiel.janulgue@gmail.com>
---
 rust/bindings/bindings_helper.h |  1 +
 rust/helpers/page.c             | 10 +++++++++
 rust/kernel/page.rs             | 38 ++++++++++++++++++++++-----------
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index a80783fcbe04..daa3225a185f 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -15,6 +15,7 @@
 #include <linux/firmware.h>
 #include <linux/jiffies.h>
 #include <linux/mdio.h>
+#include <linux/mm.h>
 #include <linux/phy.h>
 #include <linux/refcount.h>
 #include <linux/sched.h>
diff --git a/rust/helpers/page.c b/rust/helpers/page.c
index b3f2b8fbf87f..48d4481c1e33 100644
--- a/rust/helpers/page.c
+++ b/rust/helpers/page.c
@@ -17,3 +17,13 @@ void rust_helper_kunmap_local(const void *addr)
 {
 	kunmap_local(addr);
 }
+
+void rust_helper_put_page(struct page *page)
+{
+	put_page(page);
+}
+
+void rust_helper_get_page(struct page *page)
+{
+	get_page(page);
+}
diff --git a/rust/kernel/page.rs b/rust/kernel/page.rs
index fdac6c375fe4..fdf7ee203597 100644
--- a/rust/kernel/page.rs
+++ b/rust/kernel/page.rs
@@ -8,6 +8,7 @@
     error::code::*,
     error::Result,
     uaccess::UserSliceReader,
+    types::{Opaque, ARef},
 };
 use core::ptr::{self, NonNull};
 
@@ -30,13 +31,14 @@ pub const fn page_align(addr: usize) -> usize {
     (addr + (PAGE_SIZE - 1)) & PAGE_MASK
 }
 
-/// A pointer to a page that owns the page allocation.
+/// A pointer to a reference-counted page.
 ///
 /// # Invariants
 ///
-/// The pointer is valid, and has ownership over the page.
+/// The pointer is valid.
+#[repr(transparent)]
 pub struct Page {
-    page: NonNull<bindings::page>,
+    page: Opaque<bindings::page>,
 }
 
 // SAFETY: Pages have no logic that relies on them staying on a given thread, so moving them across
@@ -71,19 +73,23 @@ impl Page {
     /// let page = Page::alloc_page(GFP_KERNEL | __GFP_ZERO)?;
     /// # Ok(()) }
     /// ```
-    pub fn alloc_page(flags: Flags) -> Result<Self, AllocError> {
+    pub fn alloc_page(flags: Flags) -> Result<ARef<Self>, AllocError> {
         // SAFETY: Depending on the value of `gfp_flags`, this call may sleep. Other than that, it
         // is always safe to call this method.
         let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
-        let page = NonNull::new(page).ok_or(AllocError)?;
-        // INVARIANT: We just successfully allocated a page, so we now have ownership of the newly
-        // allocated page. We transfer that ownership to the new `Page` object.
-        Ok(Self { page })
+        if page.is_null() {
+            return Err(AllocError);
+        }
+        // CAST: Self` is a `repr(transparent)` wrapper around `bindings::page`.
+        let ptr = page.cast::<Self>();
+        // INVARIANT: We just successfully allocated a page, ptr points to the new `Page` object.
+        // SAFETY: According to invariant above ptr is valid.
+        Ok(unsafe { ARef::from_raw(NonNull::new_unchecked(ptr)) })
     }
 
     /// Returns a raw pointer to the page.
     pub fn as_ptr(&self) -> *mut bindings::page {
-        self.page.as_ptr()
+        self.page.get()
     }
 
     /// Runs a piece of code with this page mapped to an address.
@@ -252,9 +258,15 @@ pub unsafe fn copy_from_user_slice_raw(
     }
 }
 
-impl Drop for Page {
-    fn drop(&mut self) {
-        // SAFETY: By the type invariants, we have ownership of the page and can free it.
-        unsafe { bindings::__free_pages(self.page.as_ptr(), 0) };
+// SAFETY: Instances of `Page` are always reference-counted.
+unsafe impl crate::types::AlwaysRefCounted for Page {
+    fn inc_ref(&self) {
+        // SAFETY: The existence of a shared reference guarantees that the refcount is non-zero.
+        unsafe { bindings::get_page(self.as_ptr()) };
+    }
+
+    unsafe fn dec_ref(obj: ptr::NonNull<Self>) {
+        // SAFETY: The safety requirements guarantee that the refcount is non-zero.
+        unsafe { bindings::put_page(obj.cast().as_ptr()) }
     }
 }
-- 
2.43.0