binder: Wait for security policy checks before transport termination

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

@@ -16,6 +16,7 @@
 package io.grpc.binder.internal;
 
 import android.os.IBinder;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import io.grpc.Attributes;
 import io.grpc.Grpc;
@@ -128,8 +129,8 @@ void notifyShutdown(Status status) {
   }
 
   @Override
-  void notifyTerminatedUnlocked() {
-    BinderTransportSecurity.notifyTerminatedUnlocked(getAttributes());
+  ListenableFuture<?> notifyTerminatedUnlocked() {
+    return BinderTransportSecurity.notifyTerminatedUnlocked(getAttributes());
   }
 
   @Override

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

@@ -251,12 +251,21 @@ final boolean isReady() {
   @GuardedBy("this")
   abstract void notifyTerminated();
 
-  void notifyTerminatedUnlocked() {}
+  ListenableFuture<?> notifyTerminatedUnlocked() {
+    return immediateFuture(null);
+  }
 
   void releaseExecutors() {
     executorServicePool.returnObject(scheduledExecutorService);
   }
 
+  private void finishTermination() {
+    synchronized (this) {
+      notifyTerminated();
+    }
+    releaseExecutors();
+  }
+
   // Registers the specified future for eventual safe cancellation upon shutdown/terminate.
   @GuardedBy("this")
   protected final <T extends Future<?>> T register(T future) {
@@ -337,12 +346,12 @@ final void shutdownInternal(Status shutdownStatus, boolean forceTerminate) {
               future.cancel(false); // No effect if already isDone().
             }
 
-            notifyTerminatedUnlocked();
-
-            synchronized (this) {
-              notifyTerminated();
+            ListenableFuture<?> termination = notifyTerminatedUnlocked();
+            if (termination.isDone()) {
+              finishTermination();
+            } else {
+              termination.addListener(this::finishTermination, scheduledExecutorService);
             }
-            releaseExecutors();
           });
     }
   }

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

@@ -16,6 +16,7 @@
 
 package io.grpc.binder.internal;
 
+import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -25,6 +26,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.errorprone.annotations.CheckReturnValue;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
 import io.grpc.Attributes;
 import io.grpc.Internal;
 import io.grpc.Metadata;
@@ -96,13 +98,15 @@ public static void attachAuthAttrs(
    *
    * <p>This method completes futures and therefore may execute arbitrary listener code on a
    * potentially direct Executor. To avoid deadlock, callers must not hold any locks.
+   * The returned future completes after all policy invocations have returned.
    */
   @Internal
-  public static void notifyTerminatedUnlocked(Attributes attributes) {
+  public static ListenableFuture<?> notifyTerminatedUnlocked(Attributes attributes) {
     TransportAuthorizationState state = attributes.get(TRANSPORT_AUTHORIZATION_STATE);
     if (state != null) {
-      state.notifyTerminatedUnlocked();
+      return state.notifyTerminatedUnlocked();
     }
+    return immediateFuture(null);
   }
 
   /**
@@ -192,7 +196,10 @@ static final class TransportAuthorizationState {
     private final ConcurrentHashMap<String, ListenableFuture<Status>> serviceAuthorization;
     private final Executor executor;
 
-    private volatile boolean isTerminated;
+    @GuardedBy("this") private boolean isTerminated;
+    @GuardedBy("this") private boolean cancellationComplete;
+    @GuardedBy("this") private int policyChecksInProgress;
+    private final SettableFuture<Void> terminationComplete = SettableFuture.create();
 
     /**
      * @param executor used for calling into the application. Must outlive the transport.
@@ -224,17 +231,30 @@ ListenableFuture<Status> checkAuthorization(MethodDescriptor<?, ?> method) {
         return nonCancellationPropagating(checkThenActRaceWinner);
       }
 
-      // We only check isTerminated *after* the new future is visible to other threads in
-      // serviceAuthorization. In case of a race with a simultaneous call to notifyTerminated(),
-      // better to harmlessly cancel the new future twice rather than not cancel it at all.
-      if (!isTerminated) {
-        // If notifyTerminated() already cancelled newPendingAuthResult, setFuture() will forward
-        // that cancellation to its argument so it's safe to ignore the return value here.
+      // Reserve the policy call before termination can decide that all checks have finished.
+      boolean runPolicyCheck;
+      synchronized (this) {
+        runPolicyCheck = !isTerminated;
+        if (runPolicyCheck) {
+          policyChecksInProgress++;
+        }
+      }
+      if (runPolicyCheck) {
         try {
           newPendingAuthResult.setFuture(
               serverPolicyChecker.checkAuthorizationForServiceAsync(uid, serviceName));
         } catch (Exception e) {  // Not just RuntimeException! Handle the "sneaky" checked case too.
           newPendingAuthResult.setException(e);
+        } finally {
+          boolean completeTermination;
+          synchronized (this) {
+            policyChecksInProgress--;
+            completeTermination =
+                isTerminated && cancellationComplete && policyChecksInProgress == 0;
+          }
+          if (completeTermination) {
+            terminationComplete.set(null);
+          }
         }
       } else {
         newPendingAuthResult.cancel(false);
@@ -265,20 +285,25 @@ public void onFailure(Throwable t) {
       return nonCancellationPropagating(newPendingAuthResult);
     }
 
-    /**
-     * After this method returns, every future ever returned by a prior or concurrent call to
-     * checkAuthorization() is guaranteed to be complete. Every future returned by subsequent call
-     * to checkAuthorization() is also guaranteed to be complete.
-     */
-    void notifyTerminatedUnlocked() {
-      // Any entries added to serviceAuthorization *after* this assignment will be immediately
-      // canceled by the adding thread in checkAuthorization().
-      isTerminated = true;
+    /** Cancels pending checks and waits for policy invocations to return before termination. */
+    ListenableFuture<?> notifyTerminatedUnlocked() {
+      synchronized (this) {
+        isTerminated = true;
+      }
 
-      // Cancel any entries added to serviceAuthorization *before* the volatile assignment above.
+      // Cancel outside the lock because listeners may run directly on this thread.
       for (ListenableFuture<Status> authResult : serviceAuthorization.values()) {
         authResult.cancel(false); // No-op for cached results (already done).
       }
+      boolean completeTermination;
+      synchronized (this) {
+        cancellationComplete = true;
+        completeTermination = policyChecksInProgress == 0;
+      }
+      if (completeTermination) {
+        terminationComplete.set(null);
+      }
+      return terminationComplete;
     }
   }
 

binder/src/test/java/io/grpc/binder/internal/TransportAuthorizationStateTest.java

@@ -20,6 +20,7 @@
 import static java.util.concurrent.TimeUnit.*;
 import static org.junit.Assert.assertThrows;
 
+import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.protobuf.Empty;
@@ -29,9 +30,11 @@
 import io.grpc.binder.internal.BinderTransportSecurity.TransportAuthorizationState;
 import io.grpc.protobuf.lite.ProtoLiteUtils;
 import java.util.NoSuchElementException;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
@@ -193,6 +196,39 @@ public void notifyTerminatedUnlocked_cancelsPendingAuthFutures() {
     assertThat(authResult.isCancelled()).isTrue();
   }
 
+  @Test
+  public void notifyTerminatedUnlocked_waitsForPolicyInvocation() throws Exception {
+    CountDownLatch policyEntered = new CountDownLatch(1);
+    CountDownLatch releasePolicy = new CountDownLatch(1);
+    authState =
+        new TransportAuthorizationState(
+            UID,
+            (uid, serviceName) -> {
+              policyEntered.countDown();
+              try {
+                releasePolicy.await();
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                return Futures.immediateFailedFuture(e);
+              }
+              return Futures.immediateFuture(Status.OK);
+            },
+            executor);
+
+    Future<?> check = executor.submit(() -> authState.checkAuthorization(CODEGEN_METHOD));
+    try {
+      assertThat(policyEntered.await(10, SECONDS)).isTrue();
+      ListenableFuture<?> terminated = authState.notifyTerminatedUnlocked();
+      assertThat(terminated.isDone()).isFalse();
+
+      releasePolicy.countDown();
+      check.get(10, SECONDS);
+      terminated.get(10, SECONDS);
+    } finally {
+      releasePolicy.countDown();
+    }
+  }
+
   @Test
   public void checkAuthorization_afterTermination_returnsCancelledFuture() {
     authState.notifyTerminatedUnlocked();