Re: [PATCH] rust: mm: add abstractions for mm_struct and vm_area_struct

[Benno Lossin]

On 23.07.24 16:32, Alice Ryhl wrote:
> This is a follow-up to the page abstractions [1] that were recently
> merged in 6.11. Rust Binder will need these abstractions to manipulate
> the vma in its implementation of the mmap fop on the Binder file.
> 
> The ARef wrapper is not used for mm_struct because there are several
> different types of refcounts.

I am confused, why can't you use the `ARef` wrapper for the different
types that you create below?

> This patch is based on Wedson's implementation on the old rust branch,
> but has been changed significantly. All mistakes are Alice's.
> 
> Link: https://lore.kernel.org/r/20240528-alice-mm-v7-4-78222c31b8f4@xxxxxxxxxx [1]
> Co-developed-by: Wedson Almeida Filho <wedsonaf@xxxxxxxxx>
> Signed-off-by: Wedson Almeida Filho <wedsonaf@xxxxxxxxx>
> Signed-off-by: Alice Ryhl <aliceryhl@xxxxxxxxxx>
> ---
>  rust/helpers.c         |  49 ++++++++++
>  rust/kernel/lib.rs     |   1 +
>  rust/kernel/mm.rs      | 259 +++++++++++++++++++++++++++++++++++++++++++++++++
>  rust/kernel/mm/virt.rs | 193 ++++++++++++++++++++++++++++++++++++
>  4 files changed, 502 insertions(+)
> 
> diff --git a/rust/helpers.c b/rust/helpers.c
> index 305f0577fae9..9aa5150ebe26 100644
> --- a/rust/helpers.c
> +++ b/rust/helpers.c
> @@ -199,6 +199,55 @@ rust_helper_krealloc(const void *objp, size_t new_size, gfp_t flags)
>  }
>  EXPORT_SYMBOL_GPL(rust_helper_krealloc);
> 
> +void rust_helper_mmgrab(struct mm_struct *mm)
> +{
> +	mmgrab(mm);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_mmgrab);
> +
> +void rust_helper_mmdrop(struct mm_struct *mm)
> +{
> +	mmdrop(mm);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_mmdrop);
> +
> +bool rust_helper_mmget_not_zero(struct mm_struct *mm)
> +{
> +	return mmget_not_zero(mm);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_mmget_not_zero);
> +
> +bool rust_helper_mmap_read_trylock(struct mm_struct *mm)
> +{
> +	return mmap_read_trylock(mm);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_mmap_read_trylock);
> +
> +void rust_helper_mmap_read_unlock(struct mm_struct *mm)
> +{
> +	mmap_read_unlock(mm);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_mmap_read_unlock);
> +
> +void rust_helper_mmap_write_lock(struct mm_struct *mm)
> +{
> +	mmap_write_lock(mm);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_mmap_write_lock);
> +
> +void rust_helper_mmap_write_unlock(struct mm_struct *mm)
> +{
> +	mmap_write_unlock(mm);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_mmap_write_unlock);
> +
> +struct vm_area_struct *rust_helper_vma_lookup(struct mm_struct *mm,
> +					      unsigned long addr)
> +{
> +	return vma_lookup(mm, addr);
> +}
> +EXPORT_SYMBOL_GPL(rust_helper_vma_lookup);
> +
>  /*
>   * `bindgen` binds the C `size_t` type as the Rust `usize` type, so we can
>   * use it in contexts where Rust expects a `usize` like slice (array) indices.
> diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
> index 5d310e79485f..3cbc4cf847a2 100644
> --- a/rust/kernel/lib.rs
> +++ b/rust/kernel/lib.rs
> @@ -33,6 +33,7 @@
>  pub mod ioctl;
>  #[cfg(CONFIG_KUNIT)]
>  pub mod kunit;
> +pub mod mm;
>  #[cfg(CONFIG_NET)]
>  pub mod net;
>  pub mod page;
> diff --git a/rust/kernel/mm.rs b/rust/kernel/mm.rs
> new file mode 100644
> index 000000000000..7fa1e2431944
> --- /dev/null
> +++ b/rust/kernel/mm.rs
> @@ -0,0 +1,259 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +// Copyright (C) 2024 Google LLC.
> +
> +//! Memory management.
> +//!
> +//! C header: [`include/linux/mm.h`](../../../../include/linux/mm.h)
> +
> +use crate::bindings;
> +
> +use core::{marker::PhantomData, mem::ManuallyDrop, ptr::NonNull};
> +
> +pub mod virt;
> +
> +/// A smart pointer that references a `struct mm` and owns an `mmgrab` refcount.
> +///
> +/// # Invariants
> +///
> +/// An `MmGrab` owns an `mmgrab` refcount to the inner `struct mm_struct`.

You also need that `mm` is a valid pointer.

> +pub struct MmGrab {
> +    mm: NonNull<bindings::mm_struct>,
> +}
> +
> +impl MmGrab {
> +    /// Call `mmgrab` on `current.mm`.
> +    #[inline]
> +    pub fn mmgrab_current() -> Option<Self> {
> +        // SAFETY: It's safe to get the `mm` field from current.
> +        let mm = unsafe {
> +            let current = bindings::get_current();
> +            (*current).mm
> +        };
> +
> +        let mm = NonNull::new(mm)?;
> +
> +        // SAFETY: We just checked that `mm` is not null.
> +        unsafe { bindings::mmgrab(mm.as_ptr()) };
> +
> +        // INVARIANT: We just created an `mmgrab` refcount.
> +        Some(Self { mm })
> +    }
> +
> +    /// Check whether this vma is associated with this mm.
> +    #[inline]
> +    pub fn is_same_mm(&self, area: &virt::Area) -> bool {
> +        // SAFETY: The `vm_mm` field of the area is immutable, so we can read it without
> +        // synchronization.
> +        let vm_mm = unsafe { (*area.as_ptr()).vm_mm };
> +
> +        vm_mm == self.mm.as_ptr()
> +    }
> +
> +    /// Calls `mmget_not_zero` and returns a handle if it succeeds.
> +    #[inline]
> +    pub fn mmget_not_zero(&self) -> Option<MmGet> {
> +        // SAFETY: We know that `mm` is still valid since we hold an `mmgrab` refcount.
> +        let success = unsafe { bindings::mmget_not_zero(self.mm.as_ptr()) };
> +
> +        if success {
> +            Some(MmGet { mm: self.mm })
> +        } else {
> +            None
> +        }
> +    }
> +}
> +
> +// SAFETY: It is safe to call `mmdrop` on another thread than where `mmgrab` was called.
> +unsafe impl Send for MmGrab {}
> +// SAFETY: All methods on this struct are safe to call in parallel from several threads.
> +unsafe impl Sync for MmGrab {}
> +
> +impl Drop for MmGrab {
> +    #[inline]
> +    fn drop(&mut self) {
> +        // SAFETY: This gives up an `mmgrab` refcount to a valid `struct mm_struct`.
> +        // INVARIANT: We own an `mmgrab` refcount, so we can give it up.

This INVARIANT comment seems out of place and the SAFETY comment should
probably be "By the type invariant of `Self`, we own an `mmgrab`
refcount and `self.mm` is valid.".

> +        unsafe { bindings::mmdrop(self.mm.as_ptr()) };
> +    }
> +}
> +
> +/// A smart pointer that references a `struct mm` and owns an `mmget` refcount.
> +///
> +/// Values of this type are created using [`MmGrab::mmget_not_zero`].
> +///
> +/// # Invariants
> +///
> +/// An `MmGet` owns an `mmget` refcount to the inner `struct mm_struct`.

Ditto with the valid pointer here and below.

> +pub struct MmGet {
> +    mm: NonNull<bindings::mm_struct>,
> +}
> +
> +impl MmGet {
> +    /// Lock the mmap read lock.
> +    #[inline]
> +    pub fn mmap_write_lock(&self) -> MmapWriteLock<'_> {
> +        // SAFETY: The pointer is valid since we hold a refcount.
> +        unsafe { bindings::mmap_write_lock(self.mm.as_ptr()) };
> +
> +        // INVARIANT: We just acquired the write lock, so we can transfer to this guard.
> +        //
> +        // The signature of this function ensures that the `MmapWriteLock` will not outlive this
> +        // `mmget` refcount.
> +        MmapWriteLock {
> +            mm: self.mm,
> +            _lifetime: PhantomData,
> +        }
> +    }
> +
> +    /// When dropping this refcount, use `mmput_async` instead of `mmput`.

I don't get this comment.

> +    #[inline]
> +    pub fn use_async_put(self) -> MmGetAsync {
> +        // Disable destructor of `self`.
> +        let me = ManuallyDrop::new(self);
> +
> +        MmGetAsync { mm: me.mm }
> +    }
> +}
> +
> +impl Drop for MmGet {
> +    #[inline]
> +    fn drop(&mut self) {
> +        // SAFETY: We acquired a refcount when creating this object.

You can just copy-paste the SAFETY comment from above (if you don't use
the `ARef` pattern). Ditto below.

---
Cheers,
Benno

> +        unsafe { bindings::mmput(self.mm.as_ptr()) };
> +    }
> +}
> +