Verify that an Ids' ownership is correct using associated objects

.github/workflows/ci.yml

@@ -153,7 +153,7 @@ jobs:
       ARGS: --no-default-features --features std --features ${{ matrix.runtime || 'apple' }} ${{ matrix.args }}
       # Use --no-fail-fast, except with dinghy
       TESTARGS: ${{ matrix.dinghy && ' ' || '--no-fail-fast' }} ${{ matrix.test-args }}
-      FEATURES: ${{ matrix.features || 'malloc,block,exception,catch_all,verify_message' }}
+      FEATURES: ${{ matrix.features || 'malloc,block,exception,catch_all,verify_message,unstable-verify-ownership' }}
       UNSTABLE_FEATURES: ${{ matrix.unstable-features || 'unstable-autoreleasesafe,unstable-c-unwind' }}
 
     runs-on: ${{ matrix.os }}

objc-sys/src/constants.rs

@@ -19,12 +19,32 @@ pub const nil: id = 0 as *mut _;
 /// A quick alias for a [`null_mut`][`core::ptr::null_mut`] class.
 pub const Nil: *mut objc_class = 0 as *mut _;
 
+/// Policies related to associative references.
+///
+/// These are options to [`objc_setAssociatedObject`].
+///
+/// [`objc_setAssociatedObject`]: crate::objc_setAssociatedObject
 pub type objc_AssociationPolicy = usize;
+/// Specifies a weak reference to the associated object.
+///
+/// This performs straight assignment, without message sends.
 pub const OBJC_ASSOCIATION_ASSIGN: objc_AssociationPolicy = 0;
+/// Specifies a strong reference to the associated object.
+///
+/// The association is not made atomically.
 pub const OBJC_ASSOCIATION_RETAIN_NONATOMIC: objc_AssociationPolicy = 1;
+/// Specifies that the associated object is copied.
+///
+/// The association is not made atomically.
 pub const OBJC_ASSOCIATION_COPY_NONATOMIC: objc_AssociationPolicy = 3;
-pub const OBJC_ASSOCIATION_RETAIN: objc_AssociationPolicy = 769;
-pub const OBJC_ASSOCIATION_COPY: objc_AssociationPolicy = 771;
+/// Specifies a strong reference to the associated object.
+///
+/// The association is made atomically.
+pub const OBJC_ASSOCIATION_RETAIN: objc_AssociationPolicy = 0o1401;
+/// Specifies that the associated object is copied.
+///
+/// The association is made atomically.
+pub const OBJC_ASSOCIATION_COPY: objc_AssociationPolicy = 0o1403;
 
 #[cfg(apple)]
 pub const OBJC_SYNC_SUCCESS: c_int = 0;
@@ -36,3 +56,17 @@ pub const OBJC_SYNC_TIMED_OUT: c_int = -2;
 /// Only relevant before macOS 10.13
 #[cfg(apple)]
 pub const OBJC_SYNC_NOT_INITIALIZED: c_int = -3;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_association_policy() {
+        assert_eq!(OBJC_ASSOCIATION_RETAIN, 769);
+        assert_eq!(OBJC_ASSOCIATION_COPY, 771);
+
+        // What the GNUStep headers define
+        assert_eq!(OBJC_ASSOCIATION_RETAIN, 0x301);
+        assert_eq!(OBJC_ASSOCIATION_COPY, 0x303);
+    }
+}

objc2/Cargo.toml

@@ -58,6 +58,10 @@ unstable-autoreleasesafe = []
 # Uses the nightly c_unwind feature to make throwing safe
 unstable-c-unwind = ["objc-sys/unstable-c-unwind"]
 
+# Uses associated types on every object to check whether Id's ownership rules
+# are being upheld.
+unstable-verify-ownership = []
+
 # Runtime selection. See `objc-sys` for details.
 apple = ["objc-sys/apple"]
 gnustep-1-7 = ["objc-sys/gnustep-1-7"]

objc2/src/__macro_helpers.rs

@@ -272,6 +272,8 @@ mod tests {
             // let copy: Id<RcTestObject, Shared> = unsafe { msg_send_id![&obj, copy].unwrap() };
             // let mutable_copy: Id<RcTestObject, Shared> = unsafe { msg_send_id![&obj, mutableCopy].unwrap() };
 
+            let obj: Id<RcTestObject, Shared> = obj.into();
+            expected.assert_current();
             let _self: Id<RcTestObject, Shared> = unsafe { msg_send_id![&obj, self].unwrap() };
             expected.retain += 1;
             expected.assert_current();

objc2/src/rc/id.rs

@@ -175,14 +175,15 @@ impl<T: Message + ?Sized, O: Ownership> Id<T, O> {
     #[inline]
     // Note: We don't take a reference as a parameter since it would be too
     // easy to accidentally create two aliasing mutable references.
-    pub unsafe fn new(ptr: *mut T) -> Option<Id<T, O>> {
+    pub unsafe fn new(ptr: *mut T) -> Option<Self> {
         // Should optimize down to nothing.
         // SAFETY: Upheld by the caller
-        NonNull::new(ptr).map(|ptr| unsafe { Id::new_nonnull(ptr) })
+        NonNull::new(ptr).map(|ptr| unsafe { Self::new_nonnull(ptr) })
     }
 
     #[inline]
-    unsafe fn new_nonnull(ptr: NonNull<T>) -> Id<T, O> {
+    unsafe fn new_nonnull(ptr: NonNull<T>) -> Self {
+        unsafe { O::__ensure_unique_if_owned(ptr.as_ptr().cast()) };
         Self {
             ptr,
             item: PhantomData,
@@ -199,12 +200,13 @@ impl<T: Message + ?Sized, O: Ownership> Id<T, O> {
     ///
     /// This is an associated method, and must be called as `Id::as_ptr(obj)`.
     #[inline]
-    pub fn as_ptr(this: &Id<T, O>) -> *const T {
+    pub fn as_ptr(this: &Self) -> *const T {
         this.ptr.as_ptr()
     }
 
     #[inline]
     pub(crate) fn consume_as_ptr(this: ManuallyDrop<Self>) -> *mut T {
+        unsafe { O::__relinquish_ownership(this.ptr.as_ptr().cast()) };
         this.ptr.as_ptr()
     }
 
@@ -214,7 +216,25 @@ impl<T: Message + ?Sized, O: Ownership> Id<T, O> {
         // So we just hack it with transmute!
 
         // SAFETY: Option<Id<T, _>> has the same size as *mut T
-        unsafe { mem::transmute::<ManuallyDrop<Option<Self>>, *mut T>(ManuallyDrop::new(obj)) }
+        let obj = ManuallyDrop::new(obj);
+        let ptr = unsafe { mem::transmute::<ManuallyDrop<Option<Self>>, *mut T>(obj) };
+        if !ptr.is_null() {
+            unsafe { O::__relinquish_ownership(ptr.cast()) };
+        }
+        ptr
+    }
+
+    // TODO: Pub?
+    #[inline]
+    fn forget(this: Self) -> *mut T {
+        unsafe { O::__relinquish_ownership(this.ptr.as_ptr().cast()) };
+        ManuallyDrop::new(this).ptr.as_ptr()
+    }
+
+    #[inline]
+    fn forget_nonnull(this: Self) -> NonNull<T> {
+        unsafe { O::__relinquish_ownership(this.ptr.as_ptr().cast()) };
+        ManuallyDrop::new(this).ptr
     }
 }
 
@@ -410,11 +430,11 @@ impl<T: Message, O: Ownership> Id<T, O> {
         // retained objects it is hard to imagine a case where the inner type
         // has a method with the same name.
 
-        let ptr = ManuallyDrop::new(self).ptr.as_ptr();
+        let ptr = Self::forget(self);
         // SAFETY: The `ptr` is guaranteed to be valid and have at least one
         // retain count.
-        // And because of the ManuallyDrop, we don't call the Drop
-        // implementation, so the object won't also be released there.
+        // And because of the `forget`, we don't call the Drop implementation,
+        // so the object won't also be released there.
         let res: *mut T = unsafe { ffi::objc_autorelease(ptr.cast()) }.cast();
         debug_assert_eq!(res, ptr, "objc_autorelease did not return the same pointer");
         res
@@ -467,7 +487,7 @@ impl<T: Message> Id<T, Owned> {
     #[inline]
     pub unsafe fn from_shared(obj: Id<T, Shared>) -> Self {
         // Note: We can't debug_assert retainCount because of autoreleases
-        let ptr = ManuallyDrop::new(obj).ptr;
+        let ptr = Id::forget_nonnull(obj);
         // SAFETY: The pointer is valid
         // Ownership rules are upheld by the caller
         unsafe { <Id<T, Owned>>::new_nonnull(ptr) }
@@ -496,7 +516,7 @@ impl<T: Message + ?Sized> From<Id<T, Owned>> for Id<T, Shared> {
     /// Downgrade from an owned to a shared [`Id`], allowing it to be cloned.
     #[inline]
     fn from(obj: Id<T, Owned>) -> Self {
-        let ptr = ManuallyDrop::new(obj).ptr;
+        let ptr = Id::forget_nonnull(obj);
         // SAFETY: The pointer is valid, and ownership is simply decreased
         unsafe { <Id<T, Shared>>::new_nonnull(ptr) }
     }
@@ -534,6 +554,8 @@ impl<T: ?Sized, O: Ownership> Drop for Id<T, O> {
     #[doc(alias = "release")]
     #[inline]
     fn drop(&mut self) {
+        unsafe { O::__relinquish_ownership(self.ptr.as_ptr().cast()) };
+
         // We could technically run the destructor for `T` when `O = Owned`,
         // and when `O = Shared` with (retainCount == 1), but that would be
         // confusing and inconsistent since we cannot guarantee that it's run.
@@ -704,4 +726,58 @@ mod tests {
         expected.retain += 1;
         expected.assert_current();
     }
+
+    #[test]
+    fn test_unique_owned() {
+        // Tests patterns that `unstable-verify-ownership` should allow
+        let obj = RcTestObject::new();
+
+        // To/from shared
+        let obj = obj.into();
+        let obj: Id<_, Owned> = unsafe { Id::from_shared(obj) };
+
+        // To/from autoreleased
+        let obj: Id<_, Owned> = autoreleasepool(|pool| {
+            let obj = obj.autorelease(pool);
+            unsafe { Id::retain(obj).unwrap() }
+        });
+
+        // Forget and bring back
+        let ptr = Id::forget(obj);
+        let _obj: Id<_, Owned> = unsafe { Id::new(ptr).unwrap() };
+    }
+
+    #[test]
+    #[cfg(feature = "unstable-verify-ownership")]
+    #[should_panic = "Another `Id<T, Owned>` has already been created from that object"]
+    fn test_double_owned() {
+        let mut obj = RcTestObject::new();
+        // Double-retain: This is unsound!
+        let _obj2: Id<_, Owned> = unsafe { Id::retain(&mut *obj).unwrap() };
+    }
+
+    #[test]
+    #[cfg(feature = "unstable-verify-ownership")]
+    #[should_panic = "Tried to give up ownership of `Id<T, Owned>` that wasn't owned"]
+    fn test_double_forgotten() {
+        let obj = RcTestObject::new();
+        let ptr = Id::forget(obj);
+
+        let _obj: Id<_, Owned> = Id {
+            ptr: NonNull::new(ptr).unwrap(),
+            item: PhantomData,
+            own: PhantomData,
+            notunwindsafe: PhantomData,
+        };
+        // Dropping fails
+    }
+
+    #[test]
+    #[cfg(feature = "unstable-verify-ownership")]
+    #[should_panic = "An `Id<T, Owned>` exists while trying to create `Id<T, Shared>`!"]
+    fn test_create_shared_when_owned_exists() {
+        let obj = RcTestObject::new();
+        let ptr: *const RcTestObject = &*obj;
+        let _obj: Id<_, Shared> = unsafe { Id::retain(ptr as *mut RcTestObject).unwrap() };
+    }
 }

objc2/src/rc/ownership.rs

@@ -1,3 +1,7 @@
+use core::ffi::c_void;
+
+use crate::ffi;
+
 /// A type used to mark that a struct owns the object(s) it contains,
 /// so it has the sole references to them.
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
@@ -20,7 +24,104 @@ mod private {
 ///
 /// This trait is sealed and not meant to be implemented outside of the this
 /// crate.
-pub trait Ownership: private::Sealed + 'static {}
+pub trait Ownership: private::Sealed + 'static {
+    #[doc(hidden)]
+    #[inline]
+    unsafe fn __ensure_unique_if_owned(_: *mut ffi::objc_object) {}
+    #[doc(hidden)]
+    #[inline]
+    unsafe fn __relinquish_ownership(_: *mut ffi::objc_object) {}
+}
+
+/// The value of this doesn't matter, it's only the address that we use.
+static OBJC2_ID_OWNED_UNIQUE_KEY: i32 = 0;
+
+#[inline]
+#[allow(unused)]
+fn key() -> *const c_void {
+    let ptr: *const i32 = &OBJC2_ID_OWNED_UNIQUE_KEY;
+    ptr.cast()
+}
+
+#[cfg(feature = "unstable-verify-ownership")]
+impl Ownership for Owned {
+    #[track_caller]
+    unsafe fn __ensure_unique_if_owned(obj: *mut ffi::objc_object) {
+        std::println!("\nOwned __ensure_unique_if_owned: {:?} / {:?}", obj, key());
+        if obj.is_null() {
+            panic!("Tried to take ownership of nil object!");
+        }
+
+        let associated = unsafe { ffi::objc_getAssociatedObject(obj, key()) };
+        std::println!("Owned associated: {:?}", associated);
+        if associated.is_null() {
+            // Set the associated object to something (it can be anything, so
+            // we just set it to the current object).
+            //
+            // We use `assign` because we don't want to retain or copy the
+            // object, since we just use it as a marker.
+            unsafe { ffi::objc_setAssociatedObject(obj, key(), obj, ffi::OBJC_ASSOCIATION_ASSIGN) };
+        } else {
+            panic!("Another `Id<T, Owned>` has already been created from that object!");
+        }
+        let associated = unsafe { ffi::objc_getAssociatedObject(obj, key()) };
+        std::println!("Owned associated: {:?}", associated);
+    }
 
+    #[track_caller]
+    unsafe fn __relinquish_ownership(obj: *mut ffi::objc_object) {
+        std::println!("\nOwned __relinquish_ownership: {:?} / {:?}", obj, key());
+        if obj.is_null() {
+            panic!("Tried to release ownership of nil object!");
+        }
+
+        let associated = unsafe { ffi::objc_getAssociatedObject(obj, key()) };
+        std::println!("Owned associated: {:?}", associated);
+        if associated.is_null() {
+            panic!("Tried to give up ownership of `Id<T, Owned>` that wasn't owned!");
+        } else {
+            // Clear the associated object
+            unsafe {
+                ffi::objc_setAssociatedObject(obj, key(), ffi::nil, ffi::OBJC_ASSOCIATION_ASSIGN)
+            };
+        }
+        let associated = unsafe { ffi::objc_getAssociatedObject(obj, key()) };
+        std::println!("Owned associated: {:?}", associated);
+    }
+}
+
+#[cfg(not(feature = "unstable-verify-ownership"))]
 impl Ownership for Owned {}
+
+#[cfg(feature = "unstable-verify-ownership")]
+impl Ownership for Shared {
+    #[track_caller]
+    unsafe fn __ensure_unique_if_owned(obj: *mut ffi::objc_object) {
+        std::println!("\nShared __ensure_unique_if_owned: {:?} / {:?}", obj, key());
+        if obj.is_null() {
+            panic!("Tried to take shared ownership of nil object!");
+        }
+
+        let associated = unsafe { ffi::objc_getAssociatedObject(obj, key()) };
+        std::println!("Shared associated: {:?}", associated);
+        if !associated.is_null() {
+            panic!("An `Id<T, Owned>` exists while trying to create `Id<T, Shared>`!");
+        }
+    }
+    #[track_caller]
+    unsafe fn __relinquish_ownership(obj: *mut ffi::objc_object) {
+        std::println!("\nShared __relinquish_ownership: {:?} / {:?}", obj, key());
+        if obj.is_null() {
+            panic!("Tried to release shared ownership of nil object!");
+        }
+
+        let associated = unsafe { ffi::objc_getAssociatedObject(obj, key()) };
+        std::println!("Shared associated: {:?}", associated);
+        if !associated.is_null() {
+            panic!("Tried to give up ownership of `Id<T, Shared>`!");
+        }
+    }
+}
+
+#[cfg(not(feature = "unstable-verify-ownership"))]
 impl Ownership for Shared {}

objc2/src/runtime.rs

@@ -662,10 +662,19 @@ impl Object {
         // SAFETY: Invariants upheld by caller
         unsafe { *self.ivar_mut::<T>(name) = value };
     }
+}
 
+/// Associated objects.
+///
+/// TODO: Make a public API for these.
+///
+/// See [Apple's documentation](https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ObjectiveC/Chapters/ocAssociativeReferences.html).
+impl Object {
     // objc_setAssociatedObject
     // objc_getAssociatedObject
     // objc_removeAssociatedObjects
+
+    // https://nshipster.com/associated-objects/
 }
 
 unsafe impl RefEncode for Object {