[KEYCLOAK-19859] - Patching request filter to properly end responses

quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/integration/web/QuarkusRequestFilter.java

@@ -17,12 +17,17 @@
 
 package org.keycloak.quarkus.runtime.integration.web;
 
+import static org.keycloak.services.resources.KeycloakApplication.getSessionFactory;
+
 import org.keycloak.common.ClientConnection;
+import org.keycloak.common.util.Resteasy;
 import org.keycloak.models.KeycloakSession;
-import org.keycloak.services.filters.AbstractRequestFilter;
+import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.KeycloakTransactionManager;
 
 import io.vertx.core.AsyncResult;
 import io.vertx.core.Handler;
+import io.vertx.core.Promise;
 import io.vertx.core.http.HttpServerRequest;
 import io.vertx.ext.web.RoutingContext;
 
@@ -33,7 +38,7 @@ import io.vertx.ext.web.RoutingContext;
  * <p>The filter itself runs in a event loop and should delegate to worker threads any blocking code (for now, all requests are handled
  * as blocking).
  */
-public class QuarkusRequestFilter extends AbstractRequestFilter implements Handler<RoutingContext> {
+public class QuarkusRequestFilter implements Handler<RoutingContext> {
 
     private static final Handler<AsyncResult<Object>> EMPTY_RESULT = result -> {
         // we don't really care about the result because any exception thrown should be handled by the parent class
@@ -43,49 +48,71 @@ public class QuarkusRequestFilter extends AbstractRequestFilter implements Handl
     public void handle(RoutingContext context) {
         // our code should always be run as blocking until we don't provide a better support for running non-blocking code
         // in the event loop
-        context.vertx().executeBlocking(promise -> {
+        context.vertx().executeBlocking(createBlockingHandler(context), false, EMPTY_RESULT);
-            ClientConnection connection = createClientConnection(context.request());
-
-            filter(connection, session -> {
-                try {
-                    configureContextualData(context, connection, session);
-
-                    // we need to close the session before response is sent to the client, otherwise subsequent requests could
-                    // not get the latest state because the session from the previous request is still being closed
-                    // other methods from Vert.x to add a handler to the response works asynchronously
-                    context.addHeadersEndHandler(createEndHandler(context, promise, session));
-                    context.next();
-                } catch (Throwable cause) {
-                    context.fail(cause);
-                    // re-throw so that the any exception is handled from parent
-                    throw new RuntimeException(cause);
-                }
-            });
-        }, false, EMPTY_RESULT);
     }
 
-    private Handler<Void> createEndHandler(RoutingContext context, io.vertx.core.Promise<Object> promise,
+    private Handler<Promise<Object>> createBlockingHandler(RoutingContext context) {
-            KeycloakSession session) {
+        return promise -> {
-        return event -> {
+            KeycloakSessionFactory sessionFactory = getSessionFactory();
+            KeycloakSession session = sessionFactory.create();
+
+            configureContextualData(context, createClientConnection(context.request()), session);
+            configureEndHandler(context, promise, session);
+
+            KeycloakTransactionManager tx = session.getTransactionManager();
+
             try {
-                close(session);
+                tx.begin();
+                context.next();
                 promise.complete();
             } catch (Throwable cause) {
-                context.fail(cause);
+                promise.fail(cause);
+                // re-throw so that the any exception is handled from parent
+                throw new RuntimeException(cause);
+            } finally {
+                if (!context.response().headWritten()) {
+                    // make sure the session is closed in case the handler is not called
+                    // it might happen that, for whatever reason, downstream handlers do not end the response or
+                    // no data was written to the response
+                    close(session);
+                }
             }
         };
     }
 
-    private void configureContextualData(RoutingContext context, ClientConnection connection, KeycloakSession session) {
+    /**
-        // quarkus-resteasy changed and clears the context map before dispatching
+     * Creates a handler to close the {@link KeycloakSession} before the response is written to response but after Resteasy
-        // need to push keycloak contextual objects into the routing context for retrieving it later
+     * is done with processing its output.
-        context.data().put(KeycloakSession.class.getName(), session);
+     */
-        context.data().put(ClientConnection.class.getName(), connection);
+    private void configureEndHandler(RoutingContext context, Promise<Object> promise, KeycloakSession session) {
+        context.addHeadersEndHandler(event -> {
+            try {
+                close(session);
+            } catch (Throwable cause) {
+                promise.fail(cause);
+            }
+        });
     }
 
-    @Override
+    private void configureContextualData(RoutingContext context, ClientConnection connection, KeycloakSession session) {
-    protected boolean isAutoClose() {
+        Resteasy.pushContext(ClientConnection.class, connection);
-        return false;
+        Resteasy.pushContext(KeycloakSession.class, session);
+        // quarkus-resteasy changed and clears the context map before dispatching
+        // need to push keycloak contextual objects into the routing context for retrieving it later
+        context.put(KeycloakSession.class.getName(), session);
+        context.put(ClientConnection.class.getName(), connection);
+    }
+
+    protected void close(KeycloakSession session) {
+        KeycloakTransactionManager tx = session.getTransactionManager();
+        if (tx.isActive()) {
+            if (tx.getRollbackOnly()) {
+                tx.rollback();
+            } else {
+                tx.commit();
+            }
+        }
+        session.close();
     }
 
     private ClientConnection createClientConnection(HttpServerRequest request) {

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/arquillian/annotation/SetDefaultProvider.java

@@ -33,4 +33,6 @@ public @interface SetDefaultProvider {
      * the default provider is only set after enabling a feature.
      */
     boolean beforeEnableFeature() default true;
+
+    String defaultProvider() default "";
 }

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/arquillian/containers/KeycloakQuarkusServerDeployableContainer.java

@@ -176,12 +176,13 @@ public class KeycloakQuarkusServerDeployableContainer implements DeployableConta
         commands.add("--auto-build");
         commands.add("--http-enabled=true");
 
-        if (configuration.getDebugPort() > 0) {
+        if (Boolean.parseBoolean(System.getProperty("auth.server.debug", "false"))) {
             commands.add("--debug");
-            commands.add(Integer.toString(configuration.getDebugPort()));
+            if (configuration.getDebugPort() > 0) {
-        } else if (Boolean.parseBoolean(System.getProperty("auth.server.debug", "false"))) {
+                commands.add(Integer.toString(configuration.getDebugPort()));
-            commands.add("--debug");
+            } else {
-            commands.add(System.getProperty("auth.server.debug.port", "5005"));
+                commands.add(System.getProperty("auth.server.debug.port", "5005"));
+            }
         }
 
         commands.add("--http-port=" + configuration.getBindHttpPort());
@@ -194,18 +195,9 @@ public class KeycloakQuarkusServerDeployableContainer implements DeployableConta
         commands.add("--cluster=" + System.getProperty("auth.server.quarkus.cluster.config", "local"));
         commands.addAll(getAdditionalBuildArgs());
 
-        if (!containsUserProfileSpiConfiguration(commands)) {
-            // ensure that at least one user profile provider is configured
-            commands.add("--spi-user-profile-provider=declarative-user-profile");
-        }
-
         return commands.toArray(new String[0]);
     }
 
-    private boolean containsUserProfileSpiConfiguration(List<String> commands) {
-        return commands.stream().anyMatch(s -> s.startsWith("--spi-user-profile-provider="));
-    }
-
     private void waitForReadiness() throws MalformedURLException, LifecycleException {
         SuiteContext suiteContext = this.suiteContext.get();
         //TODO: not sure if the best endpoint but it makes sure that everything is properly initialized. Once we have

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/util/SpiProvidersSwitchingUtils.java

@@ -7,6 +7,7 @@ import org.keycloak.testsuite.arquillian.ContainerInfo;
 import org.keycloak.testsuite.arquillian.SuiteContext;
 import org.keycloak.testsuite.arquillian.annotation.SetDefaultProvider;
 import org.keycloak.testsuite.arquillian.containers.KeycloakQuarkusServerDeployableContainer;
+import org.keycloak.utils.StringUtil;
 import org.wildfly.extras.creaper.core.online.CliException;
 import org.wildfly.extras.creaper.core.online.ModelNodeResult;
 import org.wildfly.extras.creaper.core.online.OnlineManagementClient;
@@ -27,7 +28,8 @@ public class SpiProvidersSwitchingUtils {
     private enum SpiSwitcher {
         UNDERTOW {
             @Override
-            public Optional<String> getCurrentDefaultProvider(Container container, String spiName) {
+            public Optional<String> getCurrentDefaultProvider(Container container, String spiName,
+                    SetDefaultProvider annotation) {
                 return Optional.ofNullable(System.getProperty(getProviderPropertyName(spiName)));
             }
 
@@ -48,7 +50,8 @@ public class SpiProvidersSwitchingUtils {
         WILDFLY {
             
             @Override
-            public Optional<String> getCurrentDefaultProvider(Container container, String spiName) {
+            public Optional<String> getCurrentDefaultProvider(Container container, String spiName,
+                    SetDefaultProvider annotation) {
                 String cliCmd = SUBSYSTEM_KEYCLOAK_SERVER_SPI + spiName + ":read-attribute(name=default-provider)";
                 return runInCli(cliCmd).filter(ModelNodeResult::isSuccess)
                         .map(n -> n.get("result").asString());
@@ -88,12 +91,6 @@ public class SpiProvidersSwitchingUtils {
         },
         QUARKUS {
 
-            @Override
-            public Optional<String> getCurrentDefaultProvider(Container container, String spiName) {
-                return Optional.ofNullable(
-                        getQuarkusContainer(container).getCurrentlyConfiguredSpiProviderFor(toDashCase(spiName)));
-            }
-
             @Override
             public void setDefaultProvider(Container container, String spiName, String providerId) {
                 getQuarkusContainer(container).setAdditionalBuildArgs(Collections
@@ -102,8 +99,7 @@ public class SpiProvidersSwitchingUtils {
 
             @Override
             public void removeProviderConfig(Container container, String spiName) {
-                getQuarkusContainer(container).setAdditionalBuildArgs(Collections
+                getQuarkusContainer(container).setAdditionalBuildArgs(Collections.emptyList());
-                        .singletonList(KEYCLOAKX_ARG_SPI_PREFIX + toDashCase(spiName) + "-provider=default"));
             }
 
             private KeycloakQuarkusServerDeployableContainer getQuarkusContainer(Container container) {
@@ -136,7 +132,14 @@ public class SpiProvidersSwitchingUtils {
             }
         };
 
-        public abstract Optional<String> getCurrentDefaultProvider(Container container, String spiName);
+        public Optional<String> getCurrentDefaultProvider(Container container, String spiName,
+                SetDefaultProvider annotation) {
+            String defaultProvider = annotation.defaultProvider();
+            if (StringUtil.isNotBlank(defaultProvider)) {
+                return Optional.of(defaultProvider);
+            }
+            return Optional.empty();
+        }
 
         public abstract void setDefaultProvider(Container container, String spiName, String providerId);
 
@@ -172,7 +175,7 @@ public class SpiProvidersSwitchingUtils {
         log.infof("Setting default provider for %s to %s", spi, annotation.providerId());
 
         if (annotation.onlyUpdateDefault()) {
-            spiSwitcher.getCurrentDefaultProvider(container, spi).ifPresent(v -> originalSettingsBackup.put(spi, v));
+            spiSwitcher.getCurrentDefaultProvider(container, spi, annotation).ifPresent(v -> originalSettingsBackup.put(spi, v));
             spiSwitcher.updateDefaultProvider(container, spi, annotation.providerId());
         } else {
             spiSwitcher.setDefaultProvider(container, spi, annotation.providerId());

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/user/profile/CustomUserProfileTest.java

@@ -43,7 +43,7 @@ import java.util.Map;
 /**
  * @author <a href="mailto:joerg.matysiak@bosch.io">Jörg Matysiak</a>
  */
-@SetDefaultProvider(spi="userProfile", providerId="custom-user-profile", onlyUpdateDefault = true)
+@SetDefaultProvider(spi="userProfile", providerId="custom-user-profile", defaultProvider="declarative-user-profile", onlyUpdateDefault = true)
 public class CustomUserProfileTest extends AbstractUserProfileTest {
     
     @Test