Fix a potential use-after-free

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.5.2
+
+- fix a narrow race in async cancellation. If `cancel()` on a `CompletableFuture` returned by an async function landed in the few instructions between the pipeline's `isCancelled()` check and its own `rust_future_free`, or between a wake and the re-poll it triggers, the Rust future was freed twice or polled after being freed. Seen once in CI as a glibc `tcache_thread_shutdown()` abort. Every use of the handle now runs under the future's monitor and stops once it has been freed. No measurable change to async call overhead.
+
 ## 0.5.1
 
 - adjust dependency requirements to allow consumption of `uniffi` patch bumps

src/templates/Async.java

@@ -44,29 +44,68 @@ interface AsyncCompleteVoidFunction {
         void apply(java.lang.foreign.SegmentAllocator allocator, long rustFuture, java.lang.foreign.MemorySegment status);
     }
 
-    static class UniffiFreeingFuture<T> extends java.util.concurrent.CompletableFuture<T> {
-        private java.util.function.Consumer<java.lang.Long> freeFunc;
-        private long rustFuture;
-
-        public UniffiFreeingFuture(long rustFuture, java.util.function.Consumer<java.lang.Long> freeFunc) {
-            this.freeFunc = freeFunc;
+    // Every use of the handle holds the monitor and checks `freed` first. cancel() frees on the
+    // calling thread while the poll/complete pipeline may be mid-call on an executor thread.
+    static final class UniffiFreeingFuture<T> extends java.util.concurrent.CompletableFuture<T> {
+        private final long rustFuture;
+        private final java.util.function.Consumer<java.lang.Long> freeFunc;
+        private boolean freed;
+
+        UniffiFreeingFuture(long rustFuture, java.util.function.Consumer<java.lang.Long> freeFunc) {
             this.rustFuture = rustFuture;
+            this.freeFunc = freeFunc;
         }
 
-        // Cancellation calls freeFunc immediately, which frees the underlying Rust future.
-        // This races with the thenApplyAsync pipeline stage that calls completeFunc on
-        // the same handle. The pipeline guards against this with isCancelled() checks
-        // before touching the Rust future, but a narrow race window remains if cancel()
-        // fires after the check passes. This is safe in practice because uniffi's
-        // rust_future_free is idempotent (see the double-free and poll-after-free tests).
         @Override
-        public boolean cancel(boolean ignored) {
-            boolean cancelled = super.cancel(ignored);
+        public boolean cancel(boolean mayInterruptIfRunning) {
+            boolean cancelled = super.cancel(mayInterruptIfRunning);
             if (cancelled) {
-                freeFunc.accept(rustFuture);
+                free();
             }
             return cancelled;
         }
+
+        synchronized void free() {
+            if (freed) {
+                return;
+            }
+            freed = true;
+            freeFunc.accept(rustFuture);
+        }
+
+        // Returns false, without polling, once freed.
+        synchronized boolean poll(PollingFunction pollFunc, long continuationHandle) {
+            if (freed) {
+                return false;
+            }
+            pollFunc.apply(rustFuture, CONTINUATION_CALLBACK_STUB, continuationHandle);
+            return true;
+        }
+
+        // Returns null once freed, which the pipeline only reaches after cancel().
+        synchronized <F, E extends java.lang.Exception> F completeRust(
+            AsyncCompleteFunction<F> completeFunc,
+            UniffiRustCallStatusErrorHandler<E> errorHandler
+        ) throws E {
+            if (freed) {
+                return null;
+            }
+            return UniffiHelpers.uniffiRustCallWithError(errorHandler, (_allocator, status) -> {
+                return completeFunc.apply(_allocator, rustFuture, status);
+            });
+        }
+
+        synchronized <E extends java.lang.Exception> void completeRustVoid(
+            AsyncCompleteVoidFunction completeFunc,
+            UniffiRustCallStatusErrorHandler<E> errorHandler
+        ) throws E {
+            if (freed) {
+                return;
+            }
+            UniffiHelpers.uniffiRustCallWithError(errorHandler, (_allocator, status) -> {
+                completeFunc.apply(_allocator, rustFuture, status);
+            });
+        }
     }
 
     // Helper so both the Java completable future and the job that handles it finishing and
@@ -97,14 +136,14 @@ static <T, F, E extends java.lang.Exception> java.util.concurrent.CompletableFut
         java.util.function.Function<F, T> liftFunc,
         UniffiRustCallStatusErrorHandler<E> errorHandler
     ){
-        java.util.concurrent.CompletableFuture<T> future = new UniffiFreeingFuture<>(rustFuture, freeFunc);
+        UniffiFreeingFuture<T> future = new UniffiFreeingFuture<>(rustFuture, freeFunc);
 
         java.util.concurrent.CompletableFuture<java.lang.Void> pollChain;
         try {
-            pollChain = pollUntilReady(rustFuture, pollFunc, uniffiExecutor);
+            pollChain = pollUntilReady(future, pollFunc, uniffiExecutor);
         } catch (java.lang.Exception e) {
-            freeFunc.accept(rustFuture);
             future.completeExceptionally(e);
+            future.free();
             return future;
         }
 
@@ -113,19 +152,13 @@ static <T, F, E extends java.lang.Exception> java.util.concurrent.CompletableFut
                 return null;
             }
             try {
-                F result = UniffiHelpers.uniffiRustCallWithError(errorHandler, (_allocator, status) -> {
-                    return completeFunc.apply(_allocator, rustFuture, status);
-                });
-                return liftFunc.apply(result);
+                F result = future.completeRust(completeFunc, errorHandler);
+                return result == null ? null : liftFunc.apply(result);
             } catch (java.lang.Exception e) {
                 throw new java.util.concurrent.CompletionException(e);
             }
         }, uniffiExecutor).whenComplete((result, throwable) -> {
-            if (future.isCancelled()) {
-                return;
-            }
             try {
-                // If we failed in the chain somewhere, now complete the future with the failure
                 if (throwable != null) {
                     java.lang.Throwable cause = throwable;
                     if (cause instanceof java.util.concurrent.CompletionException && cause.getCause() != null) {
@@ -136,7 +169,7 @@ static <T, F, E extends java.lang.Exception> java.util.concurrent.CompletableFut
                     future.complete(result);
                 }
             } finally {
-                freeFunc.accept(rustFuture);
+                future.free();
             }
         });
 
@@ -155,14 +188,14 @@ static <E extends java.lang.Exception> java.util.concurrent.CompletableFuture<ja
         java.lang.Runnable liftFunc,
         UniffiRustCallStatusErrorHandler<E> errorHandler
     ){
-        java.util.concurrent.CompletableFuture<java.lang.Void> future = new UniffiFreeingFuture<>(rustFuture, freeFunc);
+        UniffiFreeingFuture<java.lang.Void> future = new UniffiFreeingFuture<>(rustFuture, freeFunc);
 
         java.util.concurrent.CompletableFuture<java.lang.Void> pollChain;
         try {
-            pollChain = pollUntilReady(rustFuture, pollFunc, uniffiExecutor);
+            pollChain = pollUntilReady(future, pollFunc, uniffiExecutor);
         } catch (java.lang.Exception e) {
-            freeFunc.accept(rustFuture);
             future.completeExceptionally(e);
+            future.free();
             return future;
         }
 
@@ -171,19 +204,13 @@ static <E extends java.lang.Exception> java.util.concurrent.CompletableFuture<ja
                 return null;
             }
             try {
-                UniffiHelpers.uniffiRustCallWithError(errorHandler, (_allocator, status) -> {
-                    completeFunc.apply(_allocator, rustFuture, status);
-                });
+                future.completeRustVoid(completeFunc, errorHandler);
             } catch (java.lang.Exception e) {
                 throw new java.util.concurrent.CompletionException(e);
             }
             return null;
         }, uniffiExecutor).whenComplete((result, throwable) -> {
-            if (future.isCancelled()) {
-                return;
-            }
             try {
-                // If we failed in the chain somewhere, now complete the future with the failure
                 if (throwable != null) {
                     java.lang.Throwable cause = throwable;
                     if (cause instanceof java.util.concurrent.CompletionException && cause.getCause() != null) {
@@ -194,22 +221,25 @@ static <E extends java.lang.Exception> java.util.concurrent.CompletableFuture<ja
                     future.complete(null);
                 }
             } finally {
-                freeFunc.accept(rustFuture);
+                future.free();
             }
         });
 
         return future;
     }
 
-    private static java.util.concurrent.CompletableFuture<java.lang.Void> pollUntilReady(long rustFuture, PollingFunction pollFunc, java.util.concurrent.Executor uniffiExecutor) {
+    private static java.util.concurrent.CompletableFuture<java.lang.Void> pollUntilReady(UniffiFreeingFuture<?> future, PollingFunction pollFunc, java.util.concurrent.Executor uniffiExecutor) {
         java.util.concurrent.CompletableFuture<java.lang.Byte> pollFuture = new java.util.concurrent.CompletableFuture<>();
         var handle = uniffiContinuationHandleMap.insert(pollFuture);
-        pollFunc.apply(rustFuture, CONTINUATION_CALLBACK_STUB, handle);
+        if (!future.poll(pollFunc, handle)) {
+            uniffiContinuationHandleMap.remove(handle);
+            return java.util.concurrent.CompletableFuture.completedFuture(null);
+        }
         return pollFuture.thenComposeAsync(pollResult -> {
             if (pollResult == UNIFFI_RUST_FUTURE_POLL_READY) {
                 return java.util.concurrent.CompletableFuture.completedFuture(null);
             } else {
-                return pollUntilReady(rustFuture, pollFunc, uniffiExecutor);
+                return pollUntilReady(future, pollFunc, uniffiExecutor);
             }
         }, uniffiExecutor);
     }

tests/scripts/TestFixtureFutures/TestFixtureFutures.java

@@ -450,41 +450,49 @@ public CompletableFuture<java.lang.Void> tryDelay(String delayMs) {
         System.out.println("immediate cancellation (100 iterations) ... ok");
       }
 
-      // There is a theoretical race in our async code: if cancel() fires between the
-      // isCancelled() check and the freeFunc call in whenComplete, both paths call
-      // rust_future_free on the same handle. Currently this is safe because uniffi's
-      // rust_future_free is effectively idempotent:
-      //   - Handle::into_arc_borrowed increments the Arc refcount before creating the Arc,
-      //     so the RustFuture allocation stays alive across multiple free calls.
-      //   - RustFuture::free() just clears internal state (future=None, result=None) and
-      //     cancels the scheduler; the second call is a no-op.
-      //
-      // This test remains in place to catch any regression if uniffi changes its handle
-      // management to be less tolerant of double-free.
+      // cancel() from many threads racing the pipeline's complete/free and re-poll stages on
+      // the common pool. rust_future_free consumes the handle's Arc reference, so a second free
+      // or a poll after free corrupts the heap; glibc reports it as "tcache_thread_shutdown():
+      // unaligned tcache chunk detected" when the fixture's timer thread exits. Crashed the
+      // broken bindings 3/3 within 2s on an M4 where the 200-iteration single-thread loop passed.
       {
-        for (int i = 0; i < 200; i++) {
-            // 1ms sleep means the future may complete around the same time we cancel
-            var job = Futures.sayAfter((short)1, "race-" + i);
-            // Small random-ish delay to vary the race timing
-            if (i % 3 == 0) {
-                Thread.yield();
-            }
-            job.cancel(true);
+        int threads = 8, iters = 20_000;
+        var failure = new java.util.concurrent.atomic.AtomicReference<Throwable>();
+        Thread[] workers = new Thread[threads];
+        for (int t = 0; t < threads; t++) {
+            final int seed = t;
+            workers[t] = new Thread(() -> {
+                try {
+                    for (int i = 0; i < iters; i++) {
+                        CompletableFuture<?> job = switch ((i + seed) % 3) {
+                            case 0 -> Futures.sayAfter((short)1, "race");
+                            case 1 -> Futures.brokenSleep((short)1, (short)1);
+                            default -> Futures.alwaysReady();
+                        };
+                        if (i % 2 == 0) {
+                            Thread.yield();
+                        }
+                        job.cancel(true);
+                    }
+                } catch (Throwable e) {
+                    failure.compareAndSet(null, e);
+                }
+            });
+            workers[t].start();
+        }
+        for (Thread w : workers) {
+            w.join();
         }
-        // Verify the system is still healthy after many race attempts.
+        assert failure.get() == null : "cancel stress threw: " + failure.get();
+        // Lets the 1ms timer threads finish.
+        TestFixtureFutures.delay(50).get();
         var result = Futures.sayAfter((short)1, "post-race").get();
-        assert result.equals("Hello, post-race!") : "async broken after double-free race test";
-        System.out.println("double-free race (200 iterations) ... ok");
+        assert result.equals("Hello, post-race!") : "async broken after cancel race test";
+        System.out.println(MessageFormat.format("cancel race ({0} threads x {1} iterations) ... ok", threads, iters));
       }
 
-      // When a future is cancelled, our pollUntilReady chain may still have an in-flight
-      // poll when freeFunc is called. The orphaned chain can then call rust_future_poll on
-      // the freed handle. Currently this is safe because uniffi's Scheduler enters the
-      // Cancelled state on free, and any subsequent poll short-circuits to Ready via
-      // is_cancelled() without touching the inner future.
-      //
-      // This test remains in place to catch any regression if uniffi changes its
-      // post-free poll behavior.
+      // cancel() racing an in-flight re-poll. rust_future_poll on a freed handle increments a
+      // refcount in freed memory.
       {
         for (int i = 0; i < 50; i++) {
             // brokenSleep calls the waker multiple times, creating multiple polls.