Merge "Add methods to create a new NativeHandle and get fds and ints." into main

1 files changed, 152 insertions, 1 deletions

diff --git a/libs/nativewindow/rust/src/handle.rs b/libs/nativewindow/rust/src/handle.rs
index a3a9dc6258..c41ab8d1b8 100644
--- a/libs/nativewindow/rust/src/handle.rs
+++ b/libs/nativewindow/rust/src/handle.rs
@@ -12,7 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{mem::forget, ptr::NonNull};
+use std::{
+    ffi::c_int,
+    mem::forget,
+    os::fd::{BorrowedFd, FromRawFd, IntoRawFd, OwnedFd},
+    ptr::NonNull,
+};
 
 /// Rust wrapper around `native_handle_t`.
 ///
@@ -22,6 +27,108 @@ use std::{mem::forget, ptr::NonNull};
 pub struct NativeHandle(NonNull<ffi::native_handle_t>);
 
 impl NativeHandle {
+    /// Creates a new `NativeHandle` with the given file descriptors and integer values.
+    ///
+    /// The `NativeHandle` will take ownership of the file descriptors and close them when it is
+    /// dropped.
+    pub fn new(fds: Vec<OwnedFd>, ints: &[c_int]) -> Option<Self> {
+        let fd_count = fds.len();
+        // SAFETY: native_handle_create doesn't have any safety requirements.
+        let handle = unsafe {
+            ffi::native_handle_create(fd_count.try_into().unwrap(), ints.len().try_into().unwrap())
+        };
+        let handle = NonNull::new(handle)?;
+        for (i, fd) in fds.into_iter().enumerate() {
+            // SAFETY: `handle` must be valid because it was just created, and the array offset is
+            // within the bounds of what we allocated above.
+            unsafe {
+                *(*handle.as_ptr()).data.as_mut_ptr().add(i) = fd.into_raw_fd();
+            }
+        }
+        for (i, value) in ints.iter().enumerate() {
+            // SAFETY: `handle` must be valid because it was just created, and the array offset is
+            // within the bounds of what we allocated above. Note that `data` is uninitialized
+            // until after this so we can't use `slice::from_raw_parts_mut` or similar to create a
+            // reference to it so we use raw pointers arithmetic instead.
+            unsafe {
+                *(*handle.as_ptr()).data.as_mut_ptr().add(fd_count + i) = *value;
+            }
+        }
+        // SAFETY: `handle` must be valid because it was just created.
+        unsafe {
+            ffi::native_handle_set_fdsan_tag(handle.as_ptr());
+        }
+        Some(Self(handle))
+    }
+
+    /// Returns a borrowed view of all the file descriptors in this native handle.
+    pub fn fds(&self) -> Vec<BorrowedFd> {
+        self.data()[..self.fd_count()]
+            .iter()
+            .map(|fd| {
+                // SAFETY: The `native_handle_t` maintains ownership of the file descriptor so it
+                // won't be closed until this `NativeHandle` is destroyed. The `BorrowedFd` will
+                // have a lifetime constrained to that of `&self`, so it can't outlive it.
+                unsafe { BorrowedFd::borrow_raw(*fd) }
+            })
+            .collect()
+    }
+
+    /// Returns the integer values in this native handle.
+    pub fn ints(&self) -> &[c_int] {
+        &self.data()[self.fd_count()..]
+    }
+
+    /// Destroys the `NativeHandle`, taking ownership of the file descriptors it contained.
+    pub fn into_fds(self) -> Vec<OwnedFd> {
+        let fds = self.data()[..self.fd_count()]
+            .iter()
+            .map(|fd| {
+                // SAFETY: The `native_handle_t` has ownership of the file descriptor, and
+                // after this we destroy it without closing the file descriptor so we can take over
+                // ownership of it.
+                unsafe { OwnedFd::from_raw_fd(*fd) }
+            })
+            .collect();
+
+        // SAFETY: Our wrapped `native_handle_t` pointer is always valid, and it won't be accessed
+        // after this because we own it and forget it.
+        unsafe {
+            assert_eq!(ffi::native_handle_delete(self.0.as_ptr()), 0);
+        }
+        // Don't drop self, as that would cause `native_handle_close` to be called and close the
+        // file descriptors.
+        forget(self);
+        fds
+    }
+
+    /// Returns a reference to the underlying `native_handle_t`.
+    fn as_ref(&self) -> &ffi::native_handle_t {
+        // SAFETY: All the ways of creating a `NativeHandle` ensure that the `native_handle_t` is
+        // valid and initialised, and lives as long as the `NativeHandle`. We enforce Rust's
+        // aliasing rules by giving the reference a lifetime matching that of `&self`.
+        unsafe { self.0.as_ref() }
+    }
+
+    /// Returns the number of file descriptors included in the native handle.
+    fn fd_count(&self) -> usize {
+        self.as_ref().numFds.try_into().unwrap()
+    }
+
+    /// Returns the number of integer values included in the native handle.
+    fn int_count(&self) -> usize {
+        self.as_ref().numInts.try_into().unwrap()
+    }
+
+    /// Returns a slice reference for all the used `data` field of the native handle, including both
+    /// file descriptors and integers.
+    fn data(&self) -> &[c_int] {
+        let total_count = self.fd_count() + self.int_count();
+        // SAFETY: The data must have been initialised with this number of elements when the
+        // `NativeHandle` was created.
+        unsafe { self.as_ref().data.as_slice(total_count) }
+    }
+
     /// Wraps a raw `native_handle_t` pointer, taking ownership of it.
     ///
     /// # Safety
@@ -90,3 +197,47 @@ unsafe impl Send for NativeHandle {}
 // SAFETY: A `NativeHandle` can be used from different threads simultaneously, as is is just
 // integers and file descriptors.
 unsafe impl Sync for NativeHandle {}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::fs::File;
+
+    #[test]
+    fn create_empty() {
+        let handle = NativeHandle::new(vec![], &[]).unwrap();
+        assert_eq!(handle.fds().len(), 0);
+        assert_eq!(handle.ints(), &[]);
+    }
+
+    #[test]
+    fn create_with_ints() {
+        let handle = NativeHandle::new(vec![], &[1, 2, 42]).unwrap();
+        assert_eq!(handle.fds().len(), 0);
+        assert_eq!(handle.ints(), &[1, 2, 42]);
+    }
+
+    #[test]
+    fn create_with_fd() {
+        let file = File::open("/dev/null").unwrap();
+        let handle = NativeHandle::new(vec![file.into()], &[]).unwrap();
+        assert_eq!(handle.fds().len(), 1);
+        assert_eq!(handle.ints(), &[]);
+    }
+
+    #[test]
+    fn clone() {
+        let file = File::open("/dev/null").unwrap();
+        let original = NativeHandle::new(vec![file.into()], &[42]).unwrap();
+        assert_eq!(original.ints(), &[42]);
+        assert_eq!(original.fds().len(), 1);
+
+        let cloned = original.clone();
+        drop(original);
+
+        assert_eq!(cloned.ints(), &[42]);
+        assert_eq!(cloned.fds().len(), 1);
+
+        drop(cloned);
+    }
+}