Security annotations don't work with CompletionStage

extensions/security/deployment/src/test/java/io/quarkus/security/test/cdi/CDIAccessDefaultTest.java

@@ -6,13 +6,16 @@
 import static io.quarkus.security.test.utils.SecurityTestUtils.assertFailureFor;
 import static io.quarkus.security.test.utils.SecurityTestUtils.assertSuccess;
 
+import java.util.concurrent.ExecutionException;
+
 import javax.inject.Inject;
 import javax.inject.Named;
 
 import org.jboss.shrinkwrap.api.ShrinkWrap;
 import org.jboss.shrinkwrap.api.spec.JavaArchive;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.function.Executable;
 
 import io.quarkus.security.ForbiddenException;
 import io.quarkus.security.UnauthorizedException;
@@ -73,6 +76,33 @@ public void shouldRestrictAccessToSpecificRole() {
         assertSuccess(() -> bean.securedMethod(), "accessibleForAdminOnly", ADMIN);
     }
 
+    @Test
+    public void shouldRestrictAccessToSpecificRoleUni() {
+        assertFailureFor(() -> bean.securedMethodUni().await().indefinitely(), UnauthorizedException.class, ANONYMOUS);
+        assertFailureFor(() -> bean.securedMethodUni().await().indefinitely(), ForbiddenException.class, USER);
+        assertSuccess(() -> bean.securedMethodUni().await().indefinitely(), "accessibleForAdminOnly", ADMIN);
+    }
+
+    @Test
+    public void shouldRestrictAccessToSpecificRoleCompletionState() {
+        Executable executable = () -> {
+            try {
+                bean.securedMethodCompletionStage().toCompletableFuture().get();
+            } catch (ExecutionException e) {
+                throw e.getCause();
+            }
+        };
+        assertFailureFor(executable, UnauthorizedException.class, ANONYMOUS);
+        assertFailureFor(executable, ForbiddenException.class, USER);
+        assertSuccess(() -> {
+            try {
+                return bean.securedMethodCompletionStage().toCompletableFuture().get();
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            }
+        }, "accessibleForAdminOnly", ADMIN);
+    }
+
     @Test
     public void shouldFailToAccessForbiddenOnClass() {
         assertFailureFor(() -> denyAllBean.noAdditionalConstraints(), UnauthorizedException.class, ANONYMOUS);

extensions/security/deployment/src/test/java/io/quarkus/security/test/cdi/app/BeanWithSecuredMethods.java

@@ -1,10 +1,14 @@
 package io.quarkus.security.test.cdi.app;
 
+import java.util.concurrent.CompletionStage;
+
 import javax.annotation.security.DenyAll;
 import javax.annotation.security.RolesAllowed;
 import javax.enterprise.context.ApplicationScoped;
 import javax.inject.Named;
 
+import io.smallrye.mutiny.Uni;
+
 /**
  * @author Michal Szynkiewicz, [email protected]
  */
@@ -23,6 +27,16 @@ public String securedMethod() {
         return "accessibleForAdminOnly";
     }
 
+    @RolesAllowed("admin")
+    public Uni<String> securedMethodUni() {
+        return Uni.createFrom().item("accessibleForAdminOnly");
+    }
+
+    @RolesAllowed("admin")
+    public CompletionStage<String> securedMethodCompletionStage() {
+        return Uni.createFrom().item("accessibleForAdminOnly").subscribeAsCompletionStage();
+    }
+
     public String unsecuredMethod() {
         return "accessibleForAll";
     }

extensions/security/runtime/src/main/java/io/quarkus/security/runtime/interceptor/SecurityHandler.java

@@ -1,6 +1,5 @@
 package io.quarkus.security.runtime.interceptor;
 
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.function.Function;
 
@@ -33,8 +32,13 @@ public Object handle(InvocationContext ic) throws Exception {
                     .onItem().transformToUni(new UniContinuation(ic));
         } else if (CompletionStage.class.isAssignableFrom(returnType)) {
             return constrainer.nonBlockingCheck(ic.getMethod(), ic.getParameters())
-                    .subscribeAsCompletionStage()
-                    .thenApply(new CompletionStageContinuation(ic));
+                    .onItem().transformToUni((s) -> {
+                        try {
+                            return Uni.createFrom().completionStage((CompletionStage<?>) ic.proceed());
+                        } catch (Exception e) {
+                            throw new RuntimeException(e);
+                        }
+                    }).subscribeAsCompletionStage();
         } else if (Multi.class.isAssignableFrom(returnType)) {
             return constrainer.nonBlockingCheck(ic.getMethod(), ic.getParameters())
                     .onItem().transformToMulti(new MultiContinuation(ic));
@@ -65,23 +69,6 @@ public Uni<?> apply(Object o) {
         }
     }
 
-    private static class CompletionStageContinuation implements Function<Object, CompletionStage<?>> {
-        private final InvocationContext ic;
-
-        CompletionStageContinuation(InvocationContext invocationContext) {
-            ic = invocationContext;
-        }
-
-        @Override
-        public CompletionStage<?> apply(Object o) {
-            try {
-                return (CompletionStage<?>) ic.proceed();
-            } catch (Exception e) {
-                return CompletableFuture.failedFuture(e);
-            }
-        }
-    }
-
     private static class MultiContinuation implements Function<Object, Multi<?>> {
         private final InvocationContext ic;