binder: Use register() pattern to cancel futures safely without holding locks

Signed-off-by: Hyunsang Han <[email protected]>

Author: HyunSangHan

binder/src/main/java/io/grpc/binder/internal/BinderClientTransport.java

@@ -87,13 +87,6 @@ public final class BinderClientTransport extends BinderTransport
   @GuardedBy("this")
   private ScheduledFuture<?> readyTimeoutFuture; // != null iff timeout scheduled.
 
-  @GuardedBy("this")
-  @Nullable
-  private ListenableFuture<Status> authResultFuture; // null before we check auth.
-
-  @GuardedBy("this")
-  @Nullable
-  private ListenableFuture<Status> preAuthResultFuture; // null before we pre-auth.
 
   /**
    * Constructs a new transport instance.
@@ -195,7 +188,8 @@ private void preAuthorize(ServiceInfo serviceInfo) {
     // unauthorized server a chance to run, but the connection will still fail by SecurityPolicy
     // check later in handshake. Pre-auth remains effective at mitigating abuse because malware
     // can't typically control the exact timing of its installation.
-    preAuthResultFuture = checkServerAuthorizationAsync(serviceInfo.applicationInfo.uid);
+    ListenableFuture<Status> preAuthResultFuture =
+        register(checkServerAuthorizationAsync(serviceInfo.applicationInfo.uid));
     Futures.addCallback(
         preAuthResultFuture,
         new FutureCallback<Status>() {
@@ -316,9 +310,6 @@ void notifyTerminated() {
       readyTimeoutFuture.cancel(false);
       readyTimeoutFuture = null;
     }
-    cancelAsyncIfNeeded(preAuthResultFuture);
-    cancelAsyncIfNeeded(authResultFuture);
-
     serviceBinding.unbind();
     clientTransportListener.transportTerminated();
   }
@@ -337,7 +328,8 @@ protected void handleSetupTransport(Parcel parcel) {
         shutdownInternal(
             Status.UNAVAILABLE.withDescription("Malformed SETUP_TRANSPORT data"), true);
       } else {
-        authResultFuture = checkServerAuthorizationAsync(remoteUid);
+        ListenableFuture<Status> authResultFuture =
+            register(checkServerAuthorizationAsync(remoteUid));
         Futures.addCallback(
             authResultFuture,
             new FutureCallback<Status>() {
@@ -396,16 +388,6 @@ protected void handlePingResponse(Parcel parcel) {
     pingTracker.onPingResponse(parcel.readInt());
   }
 
-  /**
-   * Enqueues a future for cancellation later, on another thread.
-   * Useful when the caller wants to cancel while holding locks but the future is visible to
-   * user code which might have added listeners to run on directExecutor().
-   */
-  private void cancelAsyncIfNeeded(@Nullable ListenableFuture<?> future) {
-    if (future != null && !future.isDone()) {
-      offloadExecutor.execute(() -> future.cancel(false));
-    }
-  }
 
   private static ClientStream newFailingClientStream(
       Status failure, Attributes attributes, Metadata headers, ClientStreamTracer[] tracers) {

binder/src/main/java/io/grpc/binder/internal/BinderTransport.java

@@ -41,9 +41,11 @@
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -161,6 +163,9 @@ protected enum TransportState {
   @GuardedBy("this")
   private final LinkedHashSet<Integer> callIdsToNotifyWhenReady = new LinkedHashSet<>();
 
+  @GuardedBy("this")
+  private final List<Future<?>> ownedFutures = new ArrayList<>(); // To cancel upon terminate.
+
   @GuardedBy("this")
   protected Attributes attributes;
 
@@ -236,6 +241,13 @@ void releaseExecutors() {
     executorServicePool.returnObject(scheduledExecutorService);
   }
 
+  // Registers the specified future for eventual safe cancellation upon shutdown/terminate.
+  @GuardedBy("this")
+  protected final <T extends Future<?>> T register(T future) {
+    ownedFutures.add(future);
+    return future;
+  }
+
   @GuardedBy("this")
   boolean inState(TransportState transportState) {
     return this.transportState == transportState;
@@ -286,6 +298,8 @@ final void shutdownInternal(Status shutdownStatus, boolean forceTerminate) {
       sendShutdownTransaction();
       ArrayList<Inbound<?>> calls = new ArrayList<>(ongoingCalls.values());
       ongoingCalls.clear();
+      ArrayList<Future<?>> futuresToCancel = new ArrayList<>(ownedFutures);
+      ownedFutures.clear();
       scheduledExecutorService.execute(
           () -> {
             for (Inbound<?> inbound : calls) {
@@ -297,6 +311,11 @@ final void shutdownInternal(Status shutdownStatus, boolean forceTerminate) {
               notifyTerminated();
             }
             releaseExecutors();
+            
+            for (Future<?> future : futuresToCancel) {
+              // Not holding any locks here just in case some listener runs on a direct Executor.
+              future.cancel(false); // No effect if already isDone().
+            }
           });
     }
   }