Simplify blocking operation in identity providers

quarkusio · Oct 29, 2019 · 9b960e3 · 9b960e3
1 parent 0bd8185
commit 9b960e3
Showing 1 changed file with 26 additions and 36 deletions.
diff --git a/...runtime/src/main/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImpl.java b/...runtime/src/main/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImpl.java
@@ -13,6 +13,7 @@
 
 import org.jboss.logging.Logger;
 
+import io.quarkus.runtime.BlockingOperationControl;
 import io.quarkus.security.AuthenticationFailedException;
 import io.quarkus.security.identity.AuthenticationRequestContext;
 import io.quarkus.security.identity.IdentityProvider;
@@ -32,17 +33,33 @@ public class QuarkusIdentityProviderManagerImpl implements IdentityProviderManag
     private final List<SecurityIdentityAugmentor> augmenters;
     private final Executor blockingExecutor;
 
-    private static final AuthenticationRequestContext blockingRequestContext = new AuthenticationRequestContext() {
+    private final AuthenticationRequestContext blockingRequestContext = new AuthenticationRequestContext() {
         @Override
         public CompletionStage<SecurityIdentity> runBlocking(Supplier<SecurityIdentity> function) {
-            CompletableFuture<SecurityIdentity> ret = new CompletableFuture<>();
-            try {
-                SecurityIdentity result = function.get();
-                ret.complete(result);
-            } catch (Throwable t) {
-                ret.completeExceptionally(t);
+
+            if (BlockingOperationControl.isBlockingAllowed()) {
+                CompletableFuture<SecurityIdentity> ret = new CompletableFuture<>();
+                try {
+                    SecurityIdentity result = function.get();
+                    ret.complete(result);
+                } catch (Throwable t) {
+                    ret.completeExceptionally(t);
+                }
+                return ret;
+            } else {
+                CompletableFuture<SecurityIdentity> cf = new CompletableFuture<>();
+                blockingExecutor.execute(new Runnable() {
+                    @Override
+                    public void run() {
+                        try {
+                            cf.complete(function.get());
+                        } catch (Throwable t) {
+                            cf.completeExceptionally(t);
+                        }
+                    }
+                });
+                return cf;
             }
-            return ret;
         }
     };
 
@@ -69,7 +86,7 @@ public CompletionStage<SecurityIdentity> authenticate(AuthenticationRequest requ
                     "No IdentityProviders were registered to handle AuthenticationRequest " + request));
             return cf;
         }
-        return handleProvider(0, (List) providers, request, new AsyncAuthenticationRequestContext());
+        return handleProvider(0, (List) providers, request, blockingRequestContext);
     }
 
     /**
@@ -211,31 +228,4 @@ public int compare(SecurityIdentityAugmentor o1, SecurityIdentityAugmentor o2) {
         }
     }
 
-    private class AsyncAuthenticationRequestContext implements AuthenticationRequestContext {
-
-        private boolean inBlocking = false;
-
-        @Override
-        public CompletionStage<SecurityIdentity> runBlocking(Supplier<SecurityIdentity> function) {
-            if (inBlocking) {
-                return blockingRequestContext.runBlocking(function);
-            }
-            CompletableFuture<SecurityIdentity> cf = new CompletableFuture<>();
-            blockingExecutor.execute(new Runnable() {
-                @Override
-                public void run() {
-                    try {
-                        inBlocking = true;
-                        cf.complete(function.get());
-                    } catch (Throwable t) {
-                        cf.completeExceptionally(t);
-                    } finally {
-                        inBlocking = false;
-                    }
-                }
-            });
-
-            return cf;
-        }
-    }
 }