ndk: Add way to manage ownership of linkToDeath cookie

Bug: 181971563
Bug: 197721058
Test: m, modified existing test
Change-Id: Ib9b4ab41fd7929f9ba56e112c8a5b58966040771
diff --git a/libs/binder/rust/src/proxy.rs b/libs/binder/rust/src/proxy.rs
index b03ed49..e1d01ef 100644
--- a/libs/binder/rust/src/proxy.rs
+++ b/libs/binder/rust/src/proxy.rs
@@ -31,8 +31,10 @@
 use std::convert::TryInto;
 use std::ffi::{c_void, CString};
 use std::fmt;
+use std::mem;
 use std::os::unix::io::AsRawFd;
 use std::ptr;
+use std::sync::Arc;
 
 /// A strong reference to a Binder remote object.
 ///
@@ -378,13 +380,17 @@
             // Safety: `SpIBinder` guarantees that `self` always contains a
             // valid pointer to an `AIBinder`. `recipient` can always be
             // converted into a valid pointer to an
-            // `AIBinder_DeathRecipient`. Any value is safe to pass as the
-            // cookie, although we depend on this value being set by
-            // `get_cookie` when the death recipient callback is called.
+            // `AIBinder_DeathRecipient`.
+            //
+            // The cookie is also the correct pointer, and by calling new_cookie,
+            // we have created a new ref-count to the cookie, which linkToDeath
+            // takes ownership of. Once the DeathRecipient is unlinked for any
+            // reason (including if this call fails), the onUnlinked callback
+            // will consume that ref-count.
             sys::AIBinder_linkToDeath(
                 self.as_native_mut(),
                 recipient.as_native_mut(),
-                recipient.get_cookie(),
+                recipient.new_cookie(),
             )
         })
     }
@@ -552,10 +558,20 @@
 }
 
 /// Rust wrapper around DeathRecipient objects.
+///
+/// The cookie in this struct represents an Arc<F> for the owned callback.
+/// This struct owns a ref-count of it, and so does every binder that we
+/// have been linked with.
 #[repr(C)]
 pub struct DeathRecipient {
     recipient: *mut sys::AIBinder_DeathRecipient,
-    callback: Box<dyn Fn() + Send + 'static>,
+    cookie: *mut c_void,
+    vtable: &'static DeathRecipientVtable,
+}
+
+struct DeathRecipientVtable {
+    cookie_incr_refcount: unsafe extern "C" fn(*mut c_void),
+    cookie_decr_refcount: unsafe extern "C" fn(*mut c_void),
 }
 
 impl DeathRecipient {
@@ -563,9 +579,9 @@
     /// associated object dies.
     pub fn new<F>(callback: F) -> DeathRecipient
     where
-        F: Fn() + Send + 'static,
+        F: Fn() + Send + Sync + 'static,
     {
-        let callback = Box::new(callback);
+        let callback: *const F = Arc::into_raw(Arc::new(callback));
         let recipient = unsafe {
             // Safety: The function pointer is a valid death recipient callback.
             //
@@ -574,34 +590,85 @@
             // no longer needed.
             sys::AIBinder_DeathRecipient_new(Some(Self::binder_died::<F>))
         };
+        unsafe {
+            // Safety: The function pointer is a valid onUnlinked callback.
+            //
+            // All uses of linkToDeath in this file correctly increment the
+            // ref-count that this onUnlinked callback will decrement.
+            sys::AIBinder_DeathRecipient_setOnUnlinked(recipient, Some(Self::cookie_decr_refcount::<F>));
+        }
         DeathRecipient {
             recipient,
-            callback,
+            cookie: callback as *mut c_void,
+            vtable: &DeathRecipientVtable {
+                cookie_incr_refcount: Self::cookie_incr_refcount::<F>,
+                cookie_decr_refcount: Self::cookie_decr_refcount::<F>,
+            },
         }
     }
 
+    /// Increment the ref-count for the cookie and return it.
+    ///
+    /// # Safety
+    ///
+    /// The caller must handle the returned ref-count correctly.
+    unsafe fn new_cookie(&self) -> *mut c_void {
+        (self.vtable.cookie_incr_refcount)(self.cookie);
+
+        // Return a raw pointer with ownership of a ref-count
+        self.cookie
+    }
+
     /// Get the opaque cookie that identifies this death recipient.
     ///
     /// This cookie will be used to link and unlink this death recipient to a
     /// binder object and will be passed to the `binder_died` callback as an
     /// opaque userdata pointer.
     fn get_cookie(&self) -> *mut c_void {
-        &*self.callback as *const _ as *mut c_void
+        self.cookie
     }
 
     /// Callback invoked from C++ when the binder object dies.
     ///
     /// # Safety
     ///
-    /// The `cookie` parameter must have been created with the `get_cookie`
-    /// method of this object.
+    /// The `cookie` parameter must be the cookie for an Arc<F> and
+    /// the caller must hold a ref-count to it.
     unsafe extern "C" fn binder_died<F>(cookie: *mut c_void)
     where
-        F: Fn() + Send + 'static,
+        F: Fn() + Send + Sync + 'static,
     {
-        let callback = (cookie as *mut F).as_ref().unwrap();
+        let callback = (cookie as *const F).as_ref().unwrap();
         callback();
     }
+
+    /// Callback that decrements the ref-count.
+    /// This is invoked from C++ when a binder is unlinked.
+    ///
+    /// # Safety
+    ///
+    /// The `cookie` parameter must be the cookie for an Arc<F> and
+    /// the owner must give up a ref-count to it.
+    unsafe extern "C" fn cookie_decr_refcount<F>(cookie: *mut c_void)
+    where
+        F: Fn() + Send + Sync + 'static,
+    {
+        drop(Arc::from_raw(cookie as *const F));
+    }
+
+    /// Callback that increments the ref-count.
+    ///
+    /// # Safety
+    ///
+    /// The `cookie` parameter must be the cookie for an Arc<F> and
+    /// the owner must handle the created ref-count properly.
+    unsafe extern "C" fn cookie_incr_refcount<F>(cookie: *mut c_void)
+    where
+        F: Fn() + Send + Sync + 'static,
+    {
+        let arc = mem::ManuallyDrop::new(Arc::from_raw(cookie as *const F));
+        mem::forget(Arc::clone(&arc));
+    }
 }
 
 /// # Safety
@@ -627,6 +694,12 @@
             // `AIBinder_DeathRecipient_new` when `self` was created. This
             // delete method can only be called once when `self` is dropped.
             sys::AIBinder_DeathRecipient_delete(self.recipient);
+
+            // Safety: We own a ref-count to the cookie, and so does every
+            // linked binder. This call gives up our ref-count. The linked
+            // binders should already have given up their ref-count, or should
+            // do so shortly.
+            (self.vtable.cookie_decr_refcount)(self.cookie)
         }
     }
 }
diff --git a/libs/binder/rust/tests/integration.rs b/libs/binder/rust/tests/integration.rs
index da8907d..ac9eed2 100644
--- a/libs/binder/rust/tests/integration.rs
+++ b/libs/binder/rust/tests/integration.rs
@@ -363,13 +363,58 @@
         );
     }
 
-    fn register_death_notification(binder: &mut SpIBinder) -> (Arc<AtomicBool>, DeathRecipient) {
+    struct Bools {
+        binder_died: Arc<AtomicBool>,
+        binder_dealloc: Arc<AtomicBool>,
+    }
+
+    impl Bools {
+        fn is_dead(&self) -> bool {
+            self.binder_died.load(Ordering::Relaxed)
+        }
+        fn assert_died(&self) {
+            assert!(
+                self.is_dead(),
+                "Did not receive death notification"
+            );
+        }
+        fn assert_dropped(&self) {
+            assert!(
+                self.binder_dealloc.load(Ordering::Relaxed),
+                "Did not dealloc death notification"
+            );
+        }
+        fn assert_not_dropped(&self) {
+            assert!(
+                !self.binder_dealloc.load(Ordering::Relaxed),
+                "Dealloc death notification too early"
+            );
+        }
+    }
+
+    fn register_death_notification(binder: &mut SpIBinder) -> (Bools, DeathRecipient) {
         let binder_died = Arc::new(AtomicBool::new(false));
+        let binder_dealloc = Arc::new(AtomicBool::new(false));
+
+        struct SetOnDrop {
+            binder_dealloc: Arc<AtomicBool>,
+        }
+        impl Drop for SetOnDrop {
+            fn drop(&mut self) {
+                self.binder_dealloc.store(true, Ordering::Relaxed);
+            }
+        }
 
         let mut death_recipient = {
             let flag = binder_died.clone();
+            let set_on_drop = SetOnDrop {
+                binder_dealloc: binder_dealloc.clone(),
+            };
             DeathRecipient::new(move || {
                 flag.store(true, Ordering::Relaxed);
+                // Force the closure to take ownership of set_on_drop. When the closure is
+                // dropped, the destructor of `set_on_drop` will run.
+                let _ = &set_on_drop;
             })
         };
 
@@ -377,7 +422,12 @@
             .link_to_death(&mut death_recipient)
             .expect("link_to_death failed");
 
-        (binder_died, death_recipient)
+        let bools = Bools {
+            binder_died,
+            binder_dealloc,
+        };
+
+        (bools, death_recipient)
     }
 
     /// Killing a remote service should unregister the service and trigger
@@ -390,7 +440,7 @@
         let service_process = ScopedServiceProcess::new(service_name);
         let mut remote = binder::get_service(service_name).expect("Could not retrieve service");
 
-        let (binder_died, _recipient) = register_death_notification(&mut remote);
+        let (bools, recipient) = register_death_notification(&mut remote);
 
         drop(service_process);
         remote
@@ -400,10 +450,12 @@
         // Pause to ensure any death notifications get delivered
         thread::sleep(Duration::from_secs(1));
 
-        assert!(
-            binder_died.load(Ordering::Relaxed),
-            "Did not receive death notification"
-        );
+        bools.assert_died();
+        bools.assert_not_dropped();
+
+        drop(recipient);
+
+        bools.assert_dropped();
     }
 
     /// Test unregistering death notifications.
@@ -415,7 +467,7 @@
         let service_process = ScopedServiceProcess::new(service_name);
         let mut remote = binder::get_service(service_name).expect("Could not retrieve service");
 
-        let (binder_died, mut recipient) = register_death_notification(&mut remote);
+        let (bools, mut recipient) = register_death_notification(&mut remote);
 
         remote
             .unlink_to_death(&mut recipient)
@@ -430,9 +482,13 @@
         thread::sleep(Duration::from_secs(1));
 
         assert!(
-            !binder_died.load(Ordering::Relaxed),
+            !bools.is_dead(),
             "Received unexpected death notification after unlinking",
         );
+
+        bools.assert_not_dropped();
+        drop(recipient);
+        bools.assert_dropped();
     }
 
     /// Dropping a remote handle should unregister any death notifications.
@@ -444,7 +500,7 @@
         let service_process = ScopedServiceProcess::new(service_name);
         let mut remote = binder::get_service(service_name).expect("Could not retrieve service");
 
-        let (binder_died, _recipient) = register_death_notification(&mut remote);
+        let (bools, recipient) = register_death_notification(&mut remote);
 
         // This should automatically unregister our death notification.
         drop(remote);
@@ -457,9 +513,13 @@
         // We dropped the remote handle, so we should not receive the death
         // notification when the remote process dies here.
         assert!(
-            !binder_died.load(Ordering::Relaxed),
+            !bools.is_dead(),
             "Received unexpected death notification after dropping remote handle"
         );
+
+        bools.assert_not_dropped();
+        drop(recipient);
+        bools.assert_dropped();
     }
 
     /// Test IBinder interface methods not exercised elsewhere.