lib: rust_support: impl mutex_t wrapper

Implement `Mutex` by wrapping the C `mutex_t` API.
`Mutex` is modeled on `std::sync::Mutex`'s API.

Test: build.py generic-x86_64-test
Bug: 298705967
Change-Id: I2443da0757796f7859bc6ff38d99b104d40e8da0

6 files changed, 233 insertions, 7 deletions

diff --git a/include/kernel/mutex.h b/include/kernel/mutex.h
index 1adfa47a..97b59787 100644
--- a/include/kernel/mutex.h
+++ b/include/kernel/mutex.h
@@ -69,6 +69,8 @@ static bool is_mutex_held(mutex_t *m)
     return m->holder == get_current_thread();
 }
 
+bool extern_is_mutex_held(mutex_t *m);
+
 __END_CDECLS;
 #endif
 
diff --git a/kernel/mutex.c b/kernel/mutex.c
index 575f7b9c..505140fb 100644
--- a/kernel/mutex.c
+++ b/kernel/mutex.c
@@ -143,3 +143,6 @@ status_t mutex_release(mutex_t *m)
     return NO_ERROR;
 }
 
+bool extern_is_mutex_held(mutex_t *m) {
+    return is_mutex_held(m);
+}
diff --git a/lib/rust_support/bindings.h b/lib/rust_support/bindings.h
index 2192bee1..e628e4ea 100644
--- a/lib/rust_support/bindings.h
+++ b/lib/rust_support/bindings.h
@@ -1,7 +1,5 @@
+#include <arch/mmu.h>
 #include <err.h>
-
-#include <panic.h>
-
+#include <kernel/mutex.h>
 #include <kernel/vm.h>
-
-#include <arch/mmu.h>
+#include <panic.h>
diff --git a/lib/rust_support/lib.rs b/lib/rust_support/lib.rs
index c2fd111a..3e9a8793 100644
--- a/lib/rust_support/lib.rs
+++ b/lib/rust_support/lib.rs
@@ -34,10 +34,13 @@ use core::panic::PanicInfo;
 mod sys {
     #![allow(unused)]
     #![allow(non_camel_case_types)]
+    #![allow(non_upper_case_globals)]
     include!(env!("BINDGEN_INC_FILE"));
 }
+
 pub mod err;
 pub mod mmu;
+pub mod sync;
 pub mod vmm;
 
 pub use sys::paddr_t;
diff --git a/lib/rust_support/rules.mk b/lib/rust_support/rules.mk
index 7fc96fb8..acadcd24 100644
--- a/lib/rust_support/rules.mk
+++ b/lib/rust_support/rules.mk
@@ -36,17 +36,22 @@ MODULE_DEPS := \
 
 MODULE_BINDGEN_ALLOW_FUNCTIONS := \
 	_panic \
+	extern_is_mutex_held \
+	mutex_acquire_timeout \
+	mutex_destroy \
+	mutex_init \
+	mutex_release \
+	vaddr_to_paddr \
 	vmm_alloc_physical_etc \
 	vmm_alloc_contiguous \
 	vmm_free_region \
-	vaddr_to_paddr \
 
 MODULE_BINDGEN_ALLOW_VARS := \
+	_kernel_aspace \
 	ARCH_MMU_FLAG_.* \
 	ERR_.* \
 	PAGE_SIZE \
 	PAGE_SIZE_SHIFT \
-	_kernel_aspace \
 
 MODULE_BINDGEN_SRC_HEADER := $(LOCAL_DIR)/bindings.h
 
diff --git a/lib/rust_support/sync.rs b/lib/rust_support/sync.rs
new file mode 100644
index 00000000..47c74ba7
--- /dev/null
+++ b/lib/rust_support/sync.rs
@@ -0,0 +1,215 @@
+/*
+ * Copyright (c) 2024 Google Inc. All rights reserved
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files
+ * (the "Software"), to deal in the Software without restriction,
+ * including without limitation the rights to use, copy, modify, merge,
+ * publish, distribute, sublicense, and/or sell copies of the Software,
+ * and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+use core::cell::UnsafeCell;
+use core::fmt;
+use core::fmt::Debug;
+use core::fmt::Formatter;
+use core::mem;
+use core::ops::Deref;
+use core::ops::DerefMut;
+
+use alloc::boxed::Box;
+
+use crate::sys::extern_is_mutex_held;
+use crate::sys::mutex_acquire_timeout;
+use crate::sys::mutex_destroy;
+use crate::sys::mutex_init;
+use crate::sys::mutex_release;
+use crate::sys::mutex_t;
+use crate::sys::status_t;
+
+// TODO: `INFINITE_TIME` is defined in `lk/types.h` as `UINT32_MAX`,
+// which in turn is defined as `UINT_MAX`, which is not recognized
+// by bindgen according to the bug below so we use `u32::MAX`.
+// See <https://github.com/rust-lang/rust-bindgen/issues/1636>.
+const INFINITE_TIME: u32 = u32::MAX;
+
+/// Try to acquire the mutex with a timeout value.
+///
+/// # Safety
+///
+/// Same requirements as [`mutex_acquire_timeout`].
+#[inline]
+unsafe fn mutex_acquire(mutex: *mut mutex_t) -> status_t {
+    // SAFETY: Delegated.
+    unsafe { mutex_acquire_timeout(mutex, INFINITE_TIME) }
+}
+
+/// A mutex that does not encapsulate the data it protects.
+///
+/// This is done so that there is no manual `impl<T> Drop for Mutex<T>`,
+/// as that prevents moving out of it.
+struct LoneMutex {
+    /// We need to [`Box`] the C [`mutex_t`] because
+    /// it contains a [`list_node`], which cannot be moved.
+    /// `std` does a similar thing for immovable mutexes,
+    /// though they now use a `LazyBox` so that `fn new` can be a `const fn`.
+    /// We don't need that (yet), and `LazyBox` is tricky and not exposed by `std`,
+    /// so we just always allocate upfront.
+    ///
+    /// We need to put the [`mutex_t`] in an [`UnsafeCell`], too,
+    /// because it has interior mutability.
+    inner: Box<UnsafeCell<mutex_t>>,
+}
+
+impl LoneMutex {
+    fn new() -> Self {
+        // SAFETY: C types like `mutex_t` are zeroizable.
+        let mutex = unsafe { mem::zeroed() };
+        let mut inner = Box::new(UnsafeCell::new(mutex));
+        // SAFETY: `mutex_init` only writes to each field of a `mutex_t`
+        unsafe { mutex_init(inner.get_mut()) };
+        Self { inner }
+    }
+
+    fn get_raw(&self) -> *mut mutex_t {
+        self.inner.get()
+    }
+}
+
+impl Drop for LoneMutex {
+    fn drop(&mut self) {
+        // SAFETY: `mutex_destroy` is thread safe and it was `mutex_init`ialized.
+        unsafe {
+            mutex_destroy(self.get_raw());
+        }
+    }
+}
+
+impl Default for LoneMutex {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// A mutex wrapping the C [`mutex_t`].
+/// Its API is modeled after [`std::sync::Mutex`].
+///
+/// The main differences are that [`Mutex`]:
+///
+/// * doesn't support poisoning (it eagerly panics).
+///
+/// * doesn't lazily [`Box`] in [`Self::new`],
+///   so [`Self::new`] isn't a `const fn` either.
+#[derive(Default)]
+pub struct Mutex<T: ?Sized> {
+    mutex: LoneMutex,
+    value: UnsafeCell<T>,
+}
+
+impl<T> Mutex<T> {
+    pub fn new(value: T) -> Self {
+        Self { mutex: LoneMutex::new(), value: UnsafeCell::new(value) }
+    }
+
+    pub fn into_inner(self) -> T {
+        self.value.into_inner()
+    }
+}
+
+impl<T: ?Sized> Mutex<T> {
+    pub fn get_mut(&mut self) -> &mut T {
+        self.value.get_mut()
+    }
+}
+
+pub struct MutexGuard<'a, T: ?Sized> {
+    lock: &'a Mutex<T>,
+}
+
+impl<T: ?Sized> Mutex<T> {
+    pub fn lock(&self) -> MutexGuard<'_, T> {
+        // SAFETY: `mutex_acquire` is thread safe and it was `mutex_init`ialized.
+        let status = unsafe { mutex_acquire(self.mutex.get_raw()) };
+        assert_ne!(status, 0);
+        MutexGuard { lock: &self }
+    }
+
+    pub fn try_lock(&self) -> Option<MutexGuard<'_, T>> {
+        // SAFETY: `extern_is_mutex_held` is thread safe and it was `mutex_init`ialized.
+        let locked = unsafe { extern_is_mutex_held(self.mutex.get_raw()) };
+        if locked {
+            None
+        } else {
+            Some(self.lock())
+        }
+    }
+}
+
+impl<T: ?Sized> Drop for MutexGuard<'_, T> {
+    fn drop(&mut self) {
+        // SAFETY: `mutex_release` is thread safe and it was `mutex_init`ialized.
+        let status = unsafe { mutex_release(self.lock.mutex.get_raw()) };
+        assert_ne!(status, 0);
+    }
+}
+
+impl<T: ?Sized> Deref for MutexGuard<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: Interior mutability checked by `mutex_t`.
+        unsafe { &*self.lock.value.get() }
+    }
+}
+
+impl<T: ?Sized> DerefMut for MutexGuard<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        // SAFETY: Interior mutability checked by `mutex_t`.
+        unsafe { &mut *self.lock.value.get() }
+    }
+}
+
+/// SAFETY: This is a mutex wrapping [`mutex_t`].
+unsafe impl<T: ?Sized + Send> Send for Mutex<T> {}
+
+/// SAFETY: This is a mutex wrapping [`mutex_t`].
+unsafe impl<T: ?Sized + Send> Sync for Mutex<T> {}
+
+// Note: We don't impl `UnwindSafe` and `RefUnwindSafe`
+// because we don't poison our `Mutex` upon panicking
+// like `std::sync::Mutex` does.
+
+impl<T> From<T> for Mutex<T> {
+    fn from(value: T) -> Self {
+        Self::new(value)
+    }
+}
+
+/// Copied from [`std::sync::Mutex`],
+/// with minor changes as [`Mutex`] has no poisoning.
+impl<T: ?Sized + Debug> Debug for Mutex<T> {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        let mut d = f.debug_struct("Mutex");
+        match self.try_lock() {
+            Some(guard) => {
+                d.field("data", &&*guard);
+            }
+            None => {
+                d.field("data", &format_args!("<locked>"));
+            }
+        }
+        d.finish_non_exhaustive()
+    }
+}