ANDROID: ashmem_rust: add ASHMEM_PIN/UNPIN

Add the PIN/UNPIN ioctls that let the user pin and unpin pages in the
file. The ashmem file will store information about which pages are
unpinned in a linked list of ranges.

Instead of using the mutex of the ashmem file, the ranges are protected
by a new global mutex. This mutex will allow us to store all of the
ranges in a global LRU list.

The ranges are sorted in descending order to match ashmem.

Bug: 370906207
Change-Id: I01ffdf63dc8bbfcc64ccb01200a9f3c5cd2cdcbe
Signed-off-by: Alice Ryhl <aliceryhl@google.com>

drivers/staging/android/ashmem_range.rs

@@ -0,0 +1,379 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2024 Google LLC.
+
+//! Keeps track of unpinned ranges in an ashmem file.
+
+use crate::shmem::ShmemFile;
+use core::{mem::MaybeUninit, pin::Pin};
+use kernel::{
+    list::{List, ListArc, ListLinks},
+    prelude::*,
+    sync::{GlobalGuard, GlobalLockedBy, UniqueArc},
+};
+
+pub(crate) struct AshmemLru {}
+
+/// Represents ownership of the `ASHMEM_MUTEX` lock.
+///
+/// Using a wrapper struct around `GlobalGuard` so we can add our own methods to the guard.
+pub(crate) struct AshmemGuard(pub(crate) GlobalGuard<ASHMEM_MUTEX>);
+
+// These make `AshmemGuard` inherit the behavior of `GlobalGuard`.
+impl core::ops::Deref for AshmemGuard {
+    type Target = GlobalGuard<ASHMEM_MUTEX>;
+    fn deref(&self) -> &GlobalGuard<ASHMEM_MUTEX> {
+        &self.0
+    }
+}
+impl core::ops::DerefMut for AshmemGuard {
+    fn deref_mut(&mut self) -> &mut GlobalGuard<ASHMEM_MUTEX> {
+        &mut self.0
+    }
+}
+
+impl AshmemGuard {
+    fn shrink_range(&mut self, range: &Range, pgstart: usize, pgend: usize) {
+        let inner = range.inner.as_mut(self);
+        inner.pgstart = pgstart;
+        inner.pgend = pgend;
+    }
+}
+
+kernel::sync::global_lock! {
+    // SAFETY: We call `init` as the very first thing in the initialization of this module, so
+    // there are no calls to `lock` before `init` is called.
+    pub(crate) unsafe(uninit) static ASHMEM_MUTEX: Mutex<AshmemLru> = AshmemLru {};
+}
+
+#[pin_data]
+pub(crate) struct Range {
+    /// prev/next pointers for `Area::unpinned_list`.
+    ///
+    /// Note that "unpinned" here refers to the ASHMEM_PIN/UNPIN ioctls, which is unrelated to
+    /// Rust's concept of pinning.
+    #[pin]
+    unpinned: ListLinks<1>,
+    pub(crate) inner: GlobalLockedBy<RangeInner, ASHMEM_MUTEX>,
+}
+
+pub(crate) struct RangeInner {
+    pub(crate) pgstart: usize,
+    pub(crate) pgend: usize,
+    pub(crate) purged: bool,
+}
+
+impl Range {
+    pub(crate) fn purged(&self, guard: &AshmemGuard) -> bool {
+        self.inner.as_ref(guard).purged
+    }
+
+    pub(crate) fn pgstart(&self, guard: &AshmemGuard) -> usize {
+        self.inner.as_ref(guard).pgstart
+    }
+
+    pub(crate) fn pgend(&self, guard: &AshmemGuard) -> usize {
+        self.inner.as_ref(guard).pgend
+    }
+
+    pub(crate) fn is_before_page(&self, page: usize, guard: &AshmemGuard) -> bool {
+        let inner = self.inner.as_ref(guard);
+        inner.pgend < page
+    }
+
+    pub(crate) fn contains_page(&self, page: usize, guard: &AshmemGuard) -> bool {
+        let inner = self.inner.as_ref(guard);
+        inner.pgstart <= page && inner.pgend >= page
+    }
+
+    pub(crate) fn is_superset_of_range(
+        &self,
+        pgstart: usize,
+        pgend: usize,
+        guard: &AshmemGuard,
+    ) -> bool {
+        let inner = self.inner.as_ref(guard);
+        inner.pgstart <= pgstart && inner.pgend >= pgend
+    }
+
+    pub(crate) fn is_subset_of_range(
+        &self,
+        pgstart: usize,
+        pgend: usize,
+        guard: &AshmemGuard,
+    ) -> bool {
+        let inner = self.inner.as_ref(guard);
+        inner.pgstart >= pgstart && inner.pgend <= pgend
+    }
+
+    pub(crate) fn overlaps_with_range(
+        &self,
+        pgstart: usize,
+        pgend: usize,
+        guard: &AshmemGuard,
+    ) -> bool {
+        self.contains_page(pgstart, guard)
+            || self.contains_page(pgend, guard)
+            || self.is_subset_of_range(pgstart, pgend, guard)
+    }
+}
+
+kernel::list::impl_has_list_links! {
+    impl HasListLinks<1> for Range { self.unpinned }
+}
+
+kernel::list::impl_list_arc_safe! {
+    impl ListArcSafe<1> for Range { untracked; }
+}
+
+kernel::list::impl_list_item! {
+    impl ListItem<1> for Range { using ListLinks; }
+}
+
+pub(crate) struct Area {
+    /// List of page ranges that have been unpinned by `ASHMEM_UNPIN`.
+    ///
+    /// The ranges are sorted in descending order.
+    unpinned_list: List<Range, 1>,
+}
+
+impl Area {
+    pub(crate) fn new() -> Self {
+        Self {
+            unpinned_list: List::new(),
+        }
+    }
+
+    /// Mark the given range of pages as unpinned so they can be reclaimed.
+    ///
+    /// The `new_range` argument must be `Some` when calling this method. If this call needs an
+    /// allocation, it will take it from the option. Otherwise, the allocation is left in the
+    /// option so that the caller can free it after releasing the mutex.
+    pub(crate) fn unpin(
+        &mut self,
+        mut pgstart: usize,
+        mut pgend: usize,
+        new_range: &mut Option<NewRange<'_>>,
+        guard: &mut AshmemGuard,
+    ) {
+        let mut purged = false;
+        let mut cursor = self.unpinned_list.cursor_front();
+        while let Some(next) = cursor.peek_next() {
+            // Short-circuit: this is our insertion point.
+            if next.is_before_page(pgstart, guard) {
+                break;
+            }
+
+            // If the entire range is already unpinned, just return.
+            if next.is_superset_of_range(pgstart, pgend, guard) {
+                return;
+            }
+
+            if next.overlaps_with_range(pgstart, pgend, guard) {
+                pgstart = usize::min(pgstart, next.pgstart(guard));
+                pgend = usize::max(pgend, next.pgend(guard));
+                purged |= next.purged(guard);
+                next.remove();
+
+                // restart loop
+                cursor = self.unpinned_list.cursor_front();
+                continue;
+            }
+
+            cursor.move_next();
+        }
+
+        let new_range = new_range.take().unwrap().init(RangeInner {
+            pgstart,
+            pgend,
+            purged,
+        });
+
+        let new_range = ListArc::from(new_range);
+        cursor.insert(new_range);
+    }
+
+    /// Mark the given range of pages as pinned so they can't be reclaimed.
+    ///
+    /// Returns whether any of the pages have been reclaimed.
+    ///
+    /// The `new_range` argument must be `Some` when calling this method. If this call needs an
+    /// allocation, it will take it from the option. Otherwise, the allocation is left in the
+    /// option so that the caller can free it after releasing the mutex.
+    pub(crate) fn pin(
+        &mut self,
+        pgstart: usize,
+        pgend: usize,
+        new_range: &mut Option<NewRange<'_>>,
+        guard: &mut AshmemGuard,
+    ) -> bool {
+        let mut purged = false;
+        let mut cursor = self.unpinned_list.cursor_front();
+        while let Some(next) = cursor.peek_next() {
+            // moved past last applicable page; we can short circuit
+            if next.is_before_page(pgstart, guard) {
+                break;
+            }
+
+            // The user can ask us to pin pages that span multiple ranges,
+            // or to pin pages that aren't even unpinned, so this is messy.
+            //
+            // Four cases:
+            // 1. The requested range subsumes an existing range, so we
+            //    just remove the entire matching range.
+            // 2. The requested range overlaps the start of an existing
+            //    range, so we just update that range.
+            // 3. The requested range overlaps the end of an existing
+            //    range, so we just update that range.
+            // 4. The requested range punches a hole in an existing range,
+            //    so we have to update one side of the range and then
+            //    create a new range for the other side.
+            if next.overlaps_with_range(pgstart, pgend, guard) {
+                purged |= next.purged(guard);
+
+                let curr_pgstart = next.pgstart(guard);
+                let curr_pgend = next.pgend(guard);
+
+                if next.is_subset_of_range(pgstart, pgend, guard) {
+                    // Case #1: Easy. Just nuke the whole thing.
+                    next.remove();
+                    continue;
+                } else if curr_pgstart >= pgstart {
+                    // Case #2: We overlap from the start, so adjust it.
+                    guard.shrink_range(&next, pgend + 1, curr_pgend);
+                } else if curr_pgend <= pgend {
+                    // Case #3: We overlap from the rear, so adjust it.
+                    guard.shrink_range(&next, curr_pgstart, pgstart - 1);
+                } else {
+                    // Case #4: We eat a chunk out of the middle. A bit
+                    // more complicated, we allocate a new range for the
+                    // second half and adjust the first chunk's endpoint.
+                    guard.shrink_range(&next, curr_pgstart, pgstart - 1);
+                    let purged = next.purged(guard);
+
+                    let new_range = new_range.take().unwrap().init(RangeInner {
+                        pgstart: pgend + 1,
+                        pgend: curr_pgend,
+                        purged,
+                    });
+
+                    let new_range = ListArc::from(new_range);
+                    cursor.insert(new_range);
+                    break;
+                }
+            }
+
+            cursor.move_next();
+        }
+        purged
+    }
+
+    pub(crate) fn range_has_unpinned_page(
+        &self,
+        pgstart: usize,
+        pgend: usize,
+        guard: &mut AshmemGuard,
+    ) -> bool {
+        for range in &self.unpinned_list {
+            if range.overlaps_with_range(pgstart, pgend, guard) {
+                return true;
+            }
+        }
+        false
+    }
+}
+
+pub(crate) struct NewRange<'a> {
+    #[allow(dead_code)]
+    pub(crate) file: &'a ShmemFile,
+    pub(crate) alloc: UniqueArc<MaybeUninit<Range>>,
+}
+
+impl<'a> NewRange<'a> {
+    fn init(self, inner: RangeInner) -> Pin<UniqueArc<Range>> {
+        let new_range = self.alloc.pin_init_with(pin_init!(Range {
+            unpinned <- ListLinks::new(),
+            inner: GlobalLockedBy::new(inner),
+        }));
+
+        match new_range {
+            Ok(new_range) => new_range,
+            Err(infallible) => match infallible {},
+        }
+    }
+}
+
+#[cfg(test)]
+fn range_test() -> Result {
+    fn get_random(max: usize) -> usize {
+        let rng = unsafe { kernel::bindings::get_random_u64() };
+        (rng % max as u64) as usize
+    }
+
+    fn memset(slice: &mut [bool], value: bool) {
+        for ptr in slice {
+            *ptr = value;
+        }
+    }
+
+    const SIZE: usize = 16;
+
+    let file = ShmemFile::new(c_str!("test_file"), SIZE * PAGE_SIZE, 0)?;
+    let mut area = Area::new();
+    let mut unpinned = [false; SIZE];
+
+    let mut new_range = None;
+    for _ in 0..SIZE {
+        let start = get_random(SIZE);
+        let end = get_random(SIZE - start) + start;
+        let op = get_random(2) == 0;
+
+        if new_range.is_none() {
+            new_range = Some(NewRange {
+                file: &file,
+                alloc: UniqueArc::new_uninit(GFP_KERNEL)?,
+            });
+        }
+        let mut lock = AshmemGuard(ASHMEM_MUTEX.lock());
+        if op {
+            pr_err!("Unpinning {start} to {end}.");
+            area.unpin(start, end, &mut new_range, &mut lock);
+            memset(&mut unpinned[start..=end], true);
+        } else {
+            pr_err!("Pinning {start} to {end}.");
+            area.pin(start, end, &mut new_range, &mut lock);
+            memset(&mut unpinned[start..=end], false);
+        }
+
+        for item in &area.unpinned_list {
+            pr_err!(
+                "Seeing range {} to {}.",
+                item.pgstart(&lock),
+                item.pgend(&lock)
+            );
+        }
+
+        let mut cursor = area.unpinned_list.cursor_back();
+        let mut fail = false;
+        for i in 0..SIZE {
+            let mut target = false;
+            while let Some(prev) = cursor.peek_prev() {
+                if prev.pgend(&lock) < i {
+                    cursor.move_prev();
+                    continue;
+                }
+                target = prev.pgstart(&lock) <= i;
+                break;
+            }
+            if target != unpinned[i] {
+                pr_err!("Mismatch on {i}!");
+                fail = true;
+            }
+        }
+        if fail {
+            return Err(EINVAL);
+        }
+    }
+    pr_err!("Test completed successfully!");
+    Ok(())
+}

drivers/staging/android/ashmem_rust.rs

@@ -12,16 +12,17 @@
 
 use core::{ffi::c_int, pin::Pin};
 use kernel::{
-    bindings, c_str,
+    bindings::{self, ASHMEM_GET_PIN_STATUS, ASHMEM_PIN, ASHMEM_UNPIN},
+    c_str,
     error::Result,
     fs::{File, LocalFile},
     ioctl::_IOC_SIZE,
     miscdevice::{loff_t, IovIter, Kiocb, MiscDevice, MiscDeviceOptions, MiscDeviceRegistration},
     mm::virt::{flags as vma_flags, VmAreaNew},
-    page::page_align,
+    page::{page_align, PAGE_MASK, PAGE_SIZE},
     prelude::*,
     seq_file::{seq_print, SeqFile},
-    sync::{new_mutex, Mutex},
+    sync::{new_mutex, Mutex, UniqueArc},
     task::Task,
     uaccess::{UserSlice, UserSliceReader, UserSliceWriter},
 };
@@ -36,6 +37,9 @@ const PROT_EXEC: usize = bindings::PROT_EXEC as usize;
 const PROT_WRITE: usize = bindings::PROT_WRITE as usize;
 const PROT_MASK: usize = PROT_EXEC | PROT_READ | PROT_WRITE;
 
+mod ashmem_range;
+use ashmem_range::{Area, AshmemGuard, NewRange, ASHMEM_MUTEX};
+
 mod shmem;
 use shmem::ShmemFile;
 
@@ -62,6 +66,8 @@ impl kernel::Module for AshmemModule {
     fn init(_module: &'static kernel::ThisModule) -> Result<Self> {
         // SAFETY: Called once since this is the module initializer.
         unsafe { shmem::SHMEM_FOPS_ONCE.init() };
+        // SAFETY: Called once since this is the module initializer.
+        unsafe { ASHMEM_MUTEX.init() };
 
         pr_info!("Using Rust implementation.");
 
@@ -89,6 +95,7 @@ struct AshmemInner {
     /// If set, then this holds the ashmem name without the dev/ashmem/ prefix. No zero terminator.
     name: Option<Vec<u8>>,
     file: Option<ShmemFile>,
+    area: Area,
 }
 
 #[vtable]
@@ -104,6 +111,7 @@ impl MiscDevice for Ashmem {
                         prot_mask: PROT_MASK,
                         name: None,
                         file: None,
+                        area: Area::new(),
                     }),
                 }
             },
@@ -210,6 +218,9 @@ impl MiscDevice for Ashmem {
             bindings::ASHMEM_SET_PROT_MASK => me.set_prot_mask(arg),
             bindings::ASHMEM_GET_PROT_MASK => me.get_prot_mask(),
             bindings::ASHMEM_GET_FILE_ID => me.get_file_id(UserSlice::new(arg, size).writer()),
+            ASHMEM_PIN | ASHMEM_UNPIN | ASHMEM_GET_PIN_STATUS => {
+                me.pin_unpin(cmd, UserSlice::new(arg, size).reader())
+            }
             _ => Err(EINVAL),
         }
     }
@@ -325,6 +336,78 @@ impl Ashmem {
         writer.write(&ino)?;
         Ok(0)
     }
+
+    fn pin_unpin(&self, cmd: u32, mut reader: UserSliceReader) -> Result<isize> {
+        let (offset, cmd_len) = {
+            #[allow(dead_code)] // spurious warning because it is never explicitly constructed
+            #[repr(transparent)]
+            struct AshmemPin(bindings::ashmem_pin);
+            // SAFETY: All bit-patterns are valid for `ashmem_pin`.
+            unsafe impl kernel::types::FromBytes for AshmemPin {}
+            let AshmemPin(pin) = reader.read()?;
+            (pin.offset as usize, pin.len as usize)
+        };
+
+        // If `pin`/`unpin` needs a new range, they will take it from this `Option`. Otherwise,
+        // they will leave it here, and it gets dropped after the mutexes are released.
+        let new_range = if cmd == ASHMEM_GET_PIN_STATUS {
+            None
+        } else {
+            Some(UniqueArc::new_uninit(GFP_KERNEL)?)
+        };
+
+        let mut guard = AshmemGuard(ASHMEM_MUTEX.lock());
+        // C ashmem waits for in-flight shrinkers here using a separate mechanism, but we don't
+        // release the lock when calling `punch_hole` in the shrinker, so we don't need to do that.
+        let asma = &mut *self.inner.lock();
+        let mut new_range = match asma.file.as_ref() {
+            Some(file) => new_range.map(|alloc| NewRange { file, alloc }),
+            None => return Err(EINVAL),
+        };
+
+        // Per custom, you can pass zero for len to mean "everything onward".
+        let len = if cmd_len == 0 {
+            page_align(asma.size) - offset
+        } else {
+            cmd_len
+        };
+
+        if (offset | len) & !PAGE_MASK != 0 {
+            return Err(EINVAL);
+        }
+        let len_plus_offset = offset.checked_add(len).ok_or(EINVAL)?;
+        if page_align(asma.size) < len_plus_offset {
+            return Err(EINVAL);
+        }
+
+        let pgstart = offset / PAGE_SIZE;
+        let pgend = pgstart + (len / PAGE_SIZE) - 1;
+
+        match cmd {
+            ASHMEM_PIN => {
+                if asma.area.pin(pgstart, pgend, &mut new_range, &mut guard) {
+                    Ok(bindings::ASHMEM_WAS_PURGED as isize)
+                } else {
+                    Ok(bindings::ASHMEM_NOT_PURGED as isize)
+                }
+            }
+            ASHMEM_UNPIN => {
+                asma.area.unpin(pgstart, pgend, &mut new_range, &mut guard);
+                Ok(0)
+            }
+            ASHMEM_GET_PIN_STATUS => {
+                if asma
+                    .area
+                    .range_has_unpinned_page(pgstart, pgend, &mut guard)
+                {
+                    Ok(bindings::ASHMEM_IS_UNPINNED as isize)
+                } else {
+                    Ok(bindings::ASHMEM_IS_PINNED as isize)
+                }
+            }
+            _ => unreachable!(),
+        }
+    }
 }
 
 impl AshmemInner {