Add $destroyed flag to wrapper types to prevent use-after-free bugs

Author: ktoso

SwiftKit/src/main/java/org/swift/swiftkit/SwiftInstance.java

@@ -16,6 +16,7 @@
 
 import java.lang.foreign.GroupLayout;
 import java.lang.foreign.MemorySegment;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public interface SwiftInstance {
 
@@ -39,4 +40,13 @@ public interface SwiftInstance {
     default boolean isReferenceType() {
         return this instanceof SwiftHeapObject;
     }
+
+    /**
+     * Exposes a boolean value which can be used to indicate if the object was destroyed.
+     * <p/>
+     * This is exposing the object, rather than performing the action because we don't want to accidentally
+     * form a strong reference to the {@code SwiftInstance} which could prevent the cleanup from running,
+     * if using an GC managed instance (e.g. using an {@link AutoSwiftMemorySession}.
+     */
+    AtomicBoolean $statusDestroyedFlag();
 }

SwiftKit/src/test/java/org/swift/swiftkit/AutoArenaTest.java

@@ -18,12 +18,11 @@
 
 import java.lang.foreign.GroupLayout;
 import java.lang.foreign.MemorySegment;
-import java.lang.ref.Cleaner;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class AutoArenaTest {
 
-
     @Test
     @SuppressWarnings("removal") // System.runFinalization() will be removed
     public void cleaner_releases_native_resource() {
@@ -34,16 +33,21 @@ public void cleaner_releases_native_resource() {
 
         // we're retaining the `object`, register it with the arena:
         AutoSwiftMemorySession arena = (AutoSwiftMemorySession) SwiftArena.ofAuto();
-        arena.register(object, new SwiftHeapObjectCleanup(object.$memorySegment(), object.$swiftType()) {
-            @Override
-            public void run() throws UnexpectedRetainCountException {
-                cleanupLatch.countDown();
-            }
-        });
+
+        var statusDestroyedFlag = object.$statusDestroyedFlag();
+        Runnable markAsDestroyed = () -> statusDestroyedFlag.set(true);
+
+        arena.register(object,
+                new SwiftHeapObjectCleanup(object.$memorySegment(), object.$swiftType(), markAsDestroyed) {
+                    @Override
+                    public void run() throws UnexpectedRetainCountException {
+                        cleanupLatch.countDown();
+                    }
+                });
 
         // Release the object and hope it gets GC-ed soon
 
-        //noinspection UnusedAssignment
+        // noinspection UnusedAssignment
         object = null;
 
         var i = 1_000;
@@ -76,5 +80,10 @@ public FakeSwiftHeapObject() {
         public SwiftAnyType $swiftType() {
             return null;
         }
+
+        @Override
+        public AtomicBoolean $statusDestroyedFlag() {
+            return new AtomicBoolean();
+        }
     }
 }