From 373673826be5b4747a8c3a1729c46acd6e430027 Mon Sep 17 00:00:00 2001
From: Stuart Douglas <stuart.w.douglas@gmail.com>
Date: Tue, 29 Oct 2019 11:58:32 +1100
Subject: [PATCH] Simplify blocking operation in identity providers

---
 .../QuarkusIdentityProviderManagerImpl.java   | 62 ++++++++-----------
 1 file changed, 26 insertions(+), 36 deletions(-)

diff --git a/extensions/security/runtime/src/main/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImpl.java b/extensions/security/runtime/src/main/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImpl.java
index 4c49fd2f85550..29a08c94c8127 100644
--- a/extensions/security/runtime/src/main/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImpl.java
+++ b/extensions/security/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;
-        }
-    }
 }