binder: Tighten up simultaneous auth handling fixing #10669

Include uncached methods in the deduping logic too.

Author: jdcormie

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

@@ -16,11 +16,14 @@
 
 package io.grpc.binder.internal;
 
+import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.common.util.concurrent.SettableFuture;
 import com.google.errorprone.annotations.CheckReturnValue;
 import io.grpc.Attributes;
 import io.grpc.Internal;
@@ -167,6 +170,7 @@ private static Status statusFromFailedAuthorizationFuture(Throwable cause) {
   static final class TransportAuthorizationState {
     private final int uid;
     private final ServerPolicyChecker serverPolicyChecker;
+    // Holds *all* pending policy check futures and *certain* complete ones that we want to cache.
     private final ConcurrentHashMap<String, ListenableFuture<Status>> serviceAuthorization;
     private final Executor executor;
 
@@ -185,44 +189,53 @@ static final class TransportAuthorizationState {
     @CheckReturnValue
     ListenableFuture<Status> checkAuthorization(MethodDescriptor<?, ?> method) {
       String serviceName = method.getServiceName();
-      // Only cache decisions if the method can be sampled for tracing,
-      // which is true for all generated methods. Otherwise, programmatically
-      // created methods could cause this cache to grow unbounded.
-      boolean useCache = method.isSampledToLocalTracing();
-      if (useCache) {
-        @Nullable ListenableFuture<Status> authorization = serviceAuthorization.get(serviceName);
-        if (authorization != null) {
-          // Authorization check exists and is a pending or successful future (even if for a
-          // failed authorization).
-          return authorization;
-        }
+      @Nullable
+      ListenableFuture<Status> pendingOrCachedAuthResult = serviceAuthorization.get(serviceName);
+      if (pendingOrCachedAuthResult != null) {
+        return nonCancellationPropagating(pendingOrCachedAuthResult);
       }
-      // Under high load, this may trigger a large number of concurrent authorization checks that
-      // perform essentially the same work and have the potential of exhausting the resources they
-      // depend on. This was a non-issue in the past with synchronous policy checks due to the
-      // fixed-size nature of the thread pool this method runs under.
-      //
-      // TODO(10669): evaluate if there should be at most a single pending authorization check per
-      //  (uid, serviceName) pair at any given time.
-      ListenableFuture<Status> authorization =
-          serverPolicyChecker.checkAuthorizationForServiceAsync(uid, serviceName);
-      if (useCache) {
-        serviceAuthorization.putIfAbsent(serviceName, authorization);
-        Futures.addCallback(
-            authorization,
-            new FutureCallback<Status>() {
-              @Override
-              public void onSuccess(Status result) {}
-
-              @Override
-              public void onFailure(Throwable t) {
-                serviceAuthorization.remove(serviceName, authorization);
-              }
-            },
-            MoreExecutors.directExecutor());
+
+      SettableFuture<Status> newPendingAuthResult = SettableFuture.create();
+      ListenableFuture<Status> checkThenActRaceWinner =
+          serviceAuthorization.putIfAbsent(serviceName, newPendingAuthResult);
+      if (checkThenActRaceWinner != null) {
+        // Another thread running this method must have also just saw no entry for serviceName, then
+        // beat us to calling putIfAbsent(). We can only track one check at a time so share theirs.
+        return nonCancellationPropagating(checkThenActRaceWinner);
+      }
+
+      try {
+        newPendingAuthResult.setFuture(
+            serverPolicyChecker.checkAuthorizationForServiceAsync(uid, serviceName));
+      } catch (Exception e) {  // Not just RuntimeException! Handle the "sneaky" checked case too.
+        newPendingAuthResult.setException(e);
       }
-      return authorization;
+
+      Futures.addCallback(
+          newPendingAuthResult,
+          new FutureCallback<Status>() {
+            @Override
+            public void onSuccess(Status result) {
+              // Auth checks can be expensive so we want to cache the results. But programmatically
+              // created service names could cause the cache to grow without bound. Conservatively,
+              // we only cache results for codegen service names as there can't be too many of them.
+              if (!method.isSampledToLocalTracing()) {
+                serviceAuthorization.remove(serviceName, newPendingAuthResult);
+              }
+            }
+
+            @Override
+            public void onFailure(Throwable t) {
+              // Not simply a non-OK auth result but a failure to return any decision at all. Never
+              // cache these so that if the caller retries, we'll retry the auth check as well.
+              serviceAuthorization.remove(serviceName, newPendingAuthResult);
+            }
+          },
+          MoreExecutors.directExecutor());
+
+      return nonCancellationPropagating(newPendingAuthResult);
     }
+
   }
 
   /**

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

@@ -82,38 +82,88 @@ public void tearDown() throws Exception {
   }
 
   @Test
-  public void checkAuthorization_doesNotCacheNonCodegenMethods() throws Exception {
+  public void checkAuthorization_deduplicatesSimultaneousNonCodegenMethods() throws Exception {
     ListenableFuture<Status> authResult1 = authState.checkAuthorization(NONCODEGEN_METHOD);
     assertThat(authResult1.isDone()).isFalse();
 
+    ListenableFuture<Status> authResult2 = authState.checkAuthorization(NONCODEGEN_METHOD);
+    assertThat(authResult2.isDone()).isFalse();
+
+    // The fake policy checker was only invoked ONCE thanks to deduplication.
+    // Completing that single underlying auth check satisfies both futures.
     fakePolicyChecker.takeNextAuthRequestOrDie().set(Status.OK);
+
     assertThat(authResult1.get()).isEqualTo(Status.OK);
+    assertThat(authResult2.get()).isEqualTo(Status.OK);
     assertThat(fakePolicyChecker.statusesToSet).isEmpty();
 
     // Because it's a non-codegen method, the auth result should not be cached.
-    ListenableFuture<Status> authResult2 = authState.checkAuthorization(NONCODEGEN_METHOD);
-    assertThat(authResult2.isDone()).isFalse();
+    ListenableFuture<Status> authResult3 = authState.checkAuthorization(NONCODEGEN_METHOD);
+    assertThat(authResult3.isDone()).isFalse();
 
     fakePolicyChecker.takeNextAuthRequestOrDie().set(Status.PERMISSION_DENIED);
-    assertThat(authResult2.get()).isEqualTo(Status.PERMISSION_DENIED);
+    assertThat(authResult3.get()).isEqualTo(Status.PERMISSION_DENIED);
   }
 
   @Test
-  public void checkAuthorization_cachesCodegenMethods() throws Exception {
+  public void checkAuthorization_cachesSimultaneousCodegenMethods() throws Exception {
     ListenableFuture<Status> authResult1 = authState.checkAuthorization(CODEGEN_METHOD);
     assertThat(authResult1.isDone()).isFalse();
 
+    ListenableFuture<Status> authResult2 = authState.checkAuthorization(CODEGEN_METHOD);
+    assertThat(authResult2.isDone()).isFalse();
+
+    // The fake policy checker was only invoked ONCE thanks to deduplication.
+    // Completing that single underlying auth check satisfies both futures.
     fakePolicyChecker.takeNextAuthRequestOrDie().set(Status.OK);
+
     assertThat(authResult1.get()).isEqualTo(Status.OK);
+    assertThat(authResult2.get()).isEqualTo(Status.OK);
     assertThat(fakePolicyChecker.statusesToSet).isEmpty();
 
     // Because it's a codegen method, the auth result should be cached for the life of the object.
+    ListenableFuture<Status> authResult3 = authState.checkAuthorization(CODEGEN_METHOD);
+    assertThat(authResult3.isDone()).isTrue();
+    assertThat(authResult3.get()).isEqualTo(Status.OK);
+    assertThat(fakePolicyChecker.statusesToSet).isEmpty();
+  }
+
+  @Test
+  public void checkAuthorization_cachesPermissionDeniedForCodegenMethods() throws Exception {
+    ListenableFuture<Status> authResult1 = authState.checkAuthorization(CODEGEN_METHOD);
+    assertThat(authResult1.isDone()).isFalse();
+
+    fakePolicyChecker.takeNextAuthRequestOrDie().set(Status.PERMISSION_DENIED);
+
+    assertThat(authResult1.get()).isEqualTo(Status.PERMISSION_DENIED);
+    assertThat(fakePolicyChecker.statusesToSet).isEmpty();
+
+    // Because it's a codegen method, the non-OK status should be cached for the life of the object.
     ListenableFuture<Status> authResult2 = authState.checkAuthorization(CODEGEN_METHOD);
     assertThat(authResult2.isDone()).isTrue();
-    assertThat(authResult2.get()).isEqualTo(Status.OK);
+    assertThat(authResult2.get()).isEqualTo(Status.PERMISSION_DENIED);
     assertThat(fakePolicyChecker.statusesToSet).isEmpty();
   }
 
+  @Test
+  public void checkAuthorization_cancellation_doesNotPropagateToUnderlyingCheck() throws Exception {
+    ListenableFuture<Status> authResult1 = authState.checkAuthorization(CODEGEN_METHOD);
+    ListenableFuture<Status> authResult2 = authState.checkAuthorization(CODEGEN_METHOD);
+
+    // Cancel the first future.
+    authResult1.cancel(true);
+
+    assertThat(authResult1.isCancelled()).isTrue();
+    // The second future should NOT be cancelled, because the cancellation shouldn't propagate
+    // to the underlying shared future.
+    assertThat(authResult2.isCancelled()).isFalse();
+
+    // Completing the underlying auth check satisfies the non-cancelled future.
+    fakePolicyChecker.takeNextAuthRequestOrDie().set(Status.OK);
+
+    assertThat(authResult2.get()).isEqualTo(Status.OK);
+  }
+
   @Test
   public void checkAuthorization_failedFuture_notCached() throws Exception {
     ListenableFuture<Status> authResult1 = authState.checkAuthorization(CODEGEN_METHOD);
@@ -133,11 +183,37 @@ public void checkAuthorization_failedFuture_notCached() throws Exception {
     assertThat(authResult2.get()).isEqualTo(Status.OK);
   }
 
+
+
+  @Test
+  public void checkAuthorization_synchronousException_doesNotLeaveStrandedFuture()
+      throws Exception {
+    IllegalStateException syncException = new IllegalStateException("ouch");
+    fakePolicyChecker.syncExceptionsToThrow.add(syncException);
+
+    // The synchronous exception is safely returned as a failed future.
+    ListenableFuture<Status> authResult1 = authState.checkAuthorization(CODEGEN_METHOD);
+    ExecutionException ee = assertThrows(ExecutionException.class, authResult1::get);
+    assertThat(ee).hasCauseThat().isSameInstanceAs(syncException);
+
+    // Ensure the failed check can be retried.
+    ListenableFuture<Status> authResult2 = authState.checkAuthorization(CODEGEN_METHOD);
+    assertThat(authResult2.isDone()).isFalse();
+
+    fakePolicyChecker.takeNextAuthRequestOrDie().set(Status.OK);
+    assertThat(authResult2.get()).isEqualTo(Status.OK);
+  }
+
   private static final class FakeServerPolicyChecker implements ServerPolicyChecker {
+    final LinkedBlockingQueue<RuntimeException> syncExceptionsToThrow = new LinkedBlockingQueue<>();
     final LinkedBlockingQueue<SettableFuture<Status>> statusesToSet = new LinkedBlockingQueue<>();
 
     @Override
     public ListenableFuture<Status> checkAuthorizationForServiceAsync(int uid, String serviceName) {
+      RuntimeException syncException = syncExceptionsToThrow.poll();
+      if (syncException != null) {
+        throw syncException;
+      }
       SettableFuture<Status> pendingResult = SettableFuture.create();
       statusesToSet.add(pendingResult);
       return pendingResult;