Close the session within resteasy boundaries (#15193)

Closes #15192

quarkus/deployment/src/main/java/org/keycloak/quarkus/deployment/KeycloakProcessor.java

@@ -64,6 +64,7 @@ import io.quarkus.deployment.IsDevelopment;
 import io.quarkus.deployment.annotations.Consume;
 import io.quarkus.deployment.builditem.BootstrapConfigSetupCompleteBuildItem;
 import io.quarkus.deployment.builditem.CombinedIndexBuildItem;
+import io.quarkus.deployment.builditem.ExecutorBuildItem;
 import io.quarkus.deployment.builditem.GeneratedResourceBuildItem;
 import io.quarkus.deployment.builditem.HotDeploymentWatchedFileBuildItem;
 import io.quarkus.deployment.builditem.IndexDependencyBuildItem;
@@ -515,9 +516,10 @@ class KeycloakProcessor {
         indexDependencyBuildItemBuildProducer.produce(new IndexDependencyBuildItem("org.keycloak", "keycloak-model-jpa"));
     }
 
-    @Record(ExecutionTime.STATIC_INIT)
+    @Record(ExecutionTime.RUNTIME_INIT)
     @BuildStep
-    void initializeFilter(BuildProducer<FilterBuildItem> filters, KeycloakRecorder recorder, HttpBuildTimeConfig httpBuildConfig) {
+    void initializeFilter(BuildProducer<FilterBuildItem> filters, KeycloakRecorder recorder, HttpBuildTimeConfig httpBuildConfig,
+            ExecutorBuildItem executor) {
         String rootPath = httpBuildConfig.rootPath;
         List<String> ignoredPaths = new ArrayList<>();
 
@@ -529,7 +531,7 @@ class KeycloakProcessor {
             ignoredPaths.add(rootPath + "metrics");
         }
 
-        filters.produce(new FilterBuildItem(recorder.createRequestFilter(ignoredPaths),FilterBuildItem.AUTHORIZATION - 10));
+        filters.produce(new FilterBuildItem(recorder.createRequestFilter(ignoredPaths, executor.getExecutorProxy()),FilterBuildItem.AUTHORIZATION - 10));
     }
 
     @BuildStep

quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/KeycloakRecorder.java

@@ -20,6 +20,7 @@ package org.keycloak.quarkus.runtime;
 import java.lang.annotation.Annotation;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ExecutorService;
 import java.util.function.BiConsumer;
 import java.util.function.Predicate;
 
@@ -143,8 +144,8 @@ public class KeycloakRecorder {
         };
     }
 
-    public QuarkusRequestFilter createRequestFilter(List<String> ignoredPaths) {
+    public QuarkusRequestFilter createRequestFilter(List<String> ignoredPaths, ExecutorService executor) {
-        return new QuarkusRequestFilter(createIgnoredHttpPathsPredicate(ignoredPaths));
+        return new QuarkusRequestFilter(createIgnoredHttpPathsPredicate(ignoredPaths), executor);
     }
 
     private Predicate<RoutingContext> createIgnoredHttpPathsPredicate(List<String> ignoredPaths) {

quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/integration/jaxrs/TransactionalResponseFilter.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2022 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.quarkus.runtime.integration.jaxrs;
+
+import java.io.IOException;
+import java.util.stream.Stream;
+import javax.annotation.Priority;
+import javax.ws.rs.container.ContainerRequestContext;
+import javax.ws.rs.container.ContainerResponseContext;
+import javax.ws.rs.container.ContainerResponseFilter;
+import javax.ws.rs.container.PreMatching;
+import javax.ws.rs.core.StreamingOutput;
+import javax.ws.rs.ext.Provider;
+import org.keycloak.common.util.Resteasy;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.quarkus.runtime.transaction.TransactionalSessionHandler;
+
+@Provider
+@PreMatching
+@Priority(1)
+public class TransactionalResponseFilter implements ContainerResponseFilter, TransactionalSessionHandler {
+
+    @Override
+    public void filter(ContainerRequestContext requestContext, ContainerResponseContext responseContext)
+            throws IOException {
+        Object entity = responseContext.getEntity();
+
+        if (shouldDelaySessionClose(entity)) {
+            return;
+        }
+
+        close(Resteasy.getContextData(KeycloakSession.class));
+    }
+
+    private static boolean shouldDelaySessionClose(Object entity) {
+        // do not close the session if the response entity is a stream
+        // that is because we need the session open until the stream is transformed as it might require access to the database
+        return entity instanceof Stream || entity instanceof StreamingOutput;
+    }
+}

quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/integration/jaxrs/TransactionalResponseInterceptor.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2022 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.quarkus.runtime.integration.jaxrs;
+
+import java.io.IOException;
+import javax.annotation.Priority;
+import javax.ws.rs.ConstrainedTo;
+import javax.ws.rs.RuntimeType;
+import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.ext.Provider;
+import javax.ws.rs.ext.WriterInterceptor;
+import javax.ws.rs.ext.WriterInterceptorContext;
+import org.keycloak.common.util.Resteasy;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.quarkus.runtime.transaction.TransactionalSessionHandler;
+
+@Provider
+@ConstrainedTo(RuntimeType.SERVER)
+@Priority(10000)
+public class TransactionalResponseInterceptor implements WriterInterceptor, TransactionalSessionHandler {
+    @Override
+    public void aroundWriteTo(WriterInterceptorContext context) throws IOException, WebApplicationException {
+        KeycloakSession session = Resteasy.getContextData(KeycloakSession.class);
+
+        try {
+            context.proceed();
+        } finally {
+            // make sure response is closed after writing to the response output stream
+            // this is needed in order to support streams from endpoints as they need access to underlying resources like database
+            close(session);
+        }
+    }
+}

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

@@ -17,47 +17,46 @@
 
 package org.keycloak.quarkus.runtime.integration.web;
 
-import static org.keycloak.services.resources.KeycloakApplication.getSessionFactory;
+import java.util.concurrent.ExecutorService;
-
 import java.util.function.Predicate;
 import org.keycloak.common.ClientConnection;
 import org.keycloak.common.util.Resteasy;
 import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.quarkus.runtime.transaction.TransactionalSessionHandler;
-import org.keycloak.models.KeycloakTransactionManager;
 
-import io.netty.handler.codec.http.HttpHeaderNames;
-import io.netty.handler.codec.http.HttpHeaderValues;
-import io.netty.handler.codec.http.HttpResponseStatus;
-import io.vertx.core.AsyncResult;
 import io.vertx.core.Handler;
-import io.vertx.core.Promise;
-import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.HttpServerRequest;
-import io.vertx.core.http.HttpServerResponse;
 import io.vertx.ext.web.RoutingContext;
 
 /**
  * <p>This filter is responsible for managing the request lifecycle as well as setting up the necessary context to process incoming
- * requests.
+ * requests. We need this filter running on the top of the chain in order to push contextual objects before executing Resteasy. It is not
+ * possible to use a {@link javax.ws.rs.container.ContainerRequestFilter} for this purpose because some mechanisms like error handling
+ * will not be able to access these contextual objects.
  * 
- * <p>The filter itself runs in a event loop and should delegate to worker threads any blocking code (for now, all requests are handled
+ * <p>The filter itself runs in an event loop and should delegate to worker threads any blocking code (for now, all requests are handled
  * as blocking).
+ *
+ * <p>Note that this filter is only responsible to close the {@link KeycloakSession} if not already closed when running Resteasy code. The reason is that closing it should be done at the
+ * Resteasy level so that we don't block event loop threads even if they execute in a worker thread. Vert.x handlers and their
+ * callbacks are not designed to run blocking code. If the session is eventually closed here is because Resteasy was not executed.
+ *
+ * @see org.keycloak.quarkus.runtime.integration.jaxrs.TransactionalResponseInterceptor
+ * @see org.keycloak.quarkus.runtime.integration.jaxrs.TransactionalResponseFilter
  */
-public class QuarkusRequestFilter implements Handler<RoutingContext> {
+public class QuarkusRequestFilter implements Handler<RoutingContext>, TransactionalSessionHandler {
 
-    private static final Handler<AsyncResult<Object>> EMPTY_RESULT = result -> {
+    private final ExecutorService executor;
-        // we don't really care about the result because any exception thrown should be handled by the parent class
-    };
 
     private Predicate<RoutingContext> contextFilter;
 
     public QuarkusRequestFilter() {
-        this(null);
+        this(null, null);
     }
 
-    public QuarkusRequestFilter(Predicate<RoutingContext> contextFilter) {
+    public QuarkusRequestFilter(Predicate<RoutingContext> contextFilter, ExecutorService executor) {
         this.contextFilter = contextFilter;
+        this.executor = executor;
     }
 
     @Override
@@ -68,89 +67,43 @@ public class QuarkusRequestFilter implements Handler<RoutingContext> {
         }
         // 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(createBlockingHandler(context), false, EMPTY_RESULT);
+        executor.execute(createBlockingHandler(context));
     }
 
     private boolean ignoreContext(RoutingContext context) {
         return contextFilter != null && contextFilter.test(context);
     }
 
-    private Handler<Promise<Object>> createBlockingHandler(RoutingContext context) {
+    private Runnable createBlockingHandler(RoutingContext context) {
-        return promise -> {
+        return () -> {
-            KeycloakSessionFactory sessionFactory = getSessionFactory();
+            KeycloakSession session = configureContextualData(context);
-            KeycloakSession session = sessionFactory.create();
-
-            configureContextualData(context, createClientConnection(context.request()), session);
-            configureEndHandler(context, session);
-
-            KeycloakTransactionManager tx = session.getTransactionManager();
 
             try {
-                tx.begin();
                 context.next();
-                promise.tryComplete();
             } catch (Throwable cause) {
-                promise.fail(cause);
                 // re-throw so that the any exception is handled from parent
                 throw new RuntimeException(cause);
             } finally {
-                if (!context.response().headWritten()) {
+                // force closing the session if not already closed
-                    // make sure the session is closed in case the handler is not called
+                // under some circumstances resteasy might not be invoked like when no route is found for a particular path
-                    // it might happen that, for whatever reason, downstream handlers do not end the response or
+                // in this case context is set with status code 404, and we need to close the session
-                    // no data was written to the response
                 close(session);
             }
-            }
         };
     }
 
-    /**
+    private KeycloakSession configureContextualData(RoutingContext context) {
-     * Creates a handler to close the {@link KeycloakSession} before the response is written to response but after Resteasy
+        KeycloakSession session = create();
-     * is done with processing its output.
-     */
-    private void configureEndHandler(RoutingContext context, KeycloakSession session) {
-        context.addHeadersEndHandler(event -> {
-            try {
-                close(session);
-            } catch (Throwable cause) {
-                unexpectedErrorResponse(context.response());
-            }
-        });
-    }
 
-    private void unexpectedErrorResponse(HttpServerResponse response) {
-        response.headers().clear();
-        response.putHeader(HttpHeaderNames.CONTENT_TYPE, HttpHeaderValues.TEXT_PLAIN);
-        response.putHeader(HttpHeaderNames.CONTENT_LENGTH, "0");
-        response.setStatusCode(HttpResponseStatus.INTERNAL_SERVER_ERROR.code());
-        response.putHeader(HttpHeaderNames.CONNECTION, HttpHeaderValues.CLOSE);
-        // writes an empty buffer to replace any data previously written
-        response.write(Buffer.buffer(""));
-    }
-
-    private void configureContextualData(RoutingContext context, ClientConnection connection, KeycloakSession session) {
-        Resteasy.pushContext(ClientConnection.class, connection);
         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);
+
+        ClientConnection connection = createClientConnection(context.request());
+
+        Resteasy.pushContext(ClientConnection.class, connection);
         context.put(ClientConnection.class.getName(), connection);
-    }
 
-    protected void close(KeycloakSession session) {
+        return session;
-        KeycloakTransactionManager tx = session.getTransactionManager();
-
-        try {
-            if (tx.isActive()) {
-                if (tx.getRollbackOnly()) {
-                    tx.rollback();
-                } else {
-                    tx.commit();
-                }
-            }
-        } finally {
-            session.close();
-        }
     }
 
     private ClientConnection createClientConnection(HttpServerRequest request) {

quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/transaction/TransactionalSessionHandler.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2022 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.quarkus.runtime.transaction;
+
+import static org.keycloak.services.resources.KeycloakApplication.getSessionFactory;
+
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.KeycloakTransactionManager;
+import org.keycloak.services.DefaultKeycloakSession;
+
+/**
+ * <p>A {@link TransactionalSessionHandler} is responsible for managing transaction sessions and its lifecycle. Its subtypes
+ * are usually related to components available from the underlying stack that runs on top of the request processing chain
+ * as well as at the end in order to create transaction sessions and close them accordingly, respectively.
+ */
+public interface TransactionalSessionHandler {
+
+    /**
+     * Creates a transactional {@link KeycloakSession}.
+     *
+     * @return a transactional keycloak session
+     */
+    default KeycloakSession create() {
+        KeycloakSessionFactory sessionFactory = getSessionFactory();
+        KeycloakSession session = sessionFactory.create();
+        KeycloakTransactionManager tx = session.getTransactionManager();
+        tx.begin();
+        return session;
+    }
+
+    /**
+     * Closes a transactional {@link KeycloakSession}.
+     *
+     * @param session a transactional session
+     */
+    default void close(KeycloakSession session) {
+        if (DefaultKeycloakSession.class.cast(session).isClosed()) {
+            return;
+        }
+
+        KeycloakTransactionManager tx = session.getTransactionManager();
+
+        try {
+            if (tx.isActive()) {
+                if (tx.getRollbackOnly()) {
+                    tx.rollback();
+                } else {
+                    tx.commit();
+                }
+            }
+        } finally {
+            session.close();
+        }
+    }
+}

services/src/main/java/org/keycloak/services/DefaultKeycloakSession.java

@@ -69,7 +69,6 @@ import java.util.stream.Collectors;
  */
 public class DefaultKeycloakSession implements KeycloakSession {
 
-    private final static Logger log = Logger.getLogger(DefaultKeycloakSession.class);
     private final DefaultKeycloakSessionFactory factory;
     private final Map<Integer, Provider> providers = new HashMap<>();
     private final List<Provider> closable = new LinkedList<>();
@@ -89,6 +88,8 @@ public class DefaultKeycloakSession implements KeycloakSession {
     private VaultTranscriber vaultTranscriber;
     private ClientPolicyManager clientPolicyManager;
 
+    private boolean closed;
+
     public DefaultKeycloakSession(DefaultKeycloakSessionFactory factory) {
         this.factory = factory;
         this.transactionManager = new DefaultKeycloakTransactionManager(this);
@@ -445,7 +446,11 @@ public class DefaultKeycloakSession implements KeycloakSession {
         return clientPolicyManager;
     }
 
+    @Override
     public void close() {
+        if (closed) {
+            throw new IllegalStateException("Illegal call to #close() on already closed KeycloakSession");
+        }
         Consumer<? super Provider> safeClose = p -> {
             try {
                 p.close();
@@ -458,6 +463,10 @@ public class DefaultKeycloakSession implements KeycloakSession {
         for (Entry<InvalidableObjectType, Set<Object>> me : invalidationMap.entrySet()) {
             factory.invalidate(this, me.getKey(), me.getValue().toArray());
         }
+        closed = true;
     }
 
+    public boolean isClosed() {
+        return closed;
+    }
 }

services/src/main/java/org/keycloak/services/filters/KeycloakSecurityHeadersFilter.java

@@ -20,6 +20,7 @@ import org.keycloak.common.util.Resteasy;
 import org.keycloak.headers.SecurityHeadersProvider;
 import org.keycloak.models.KeycloakSession;
 
+import javax.annotation.Priority;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerResponseContext;
 import javax.ws.rs.container.ContainerResponseFilter;
@@ -29,6 +30,7 @@ import javax.ws.rs.container.PreMatching;
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
  */
 @PreMatching
+@Priority(10)
 public class KeycloakSecurityHeadersFilter implements ContainerResponseFilter {
 
     @Override

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/admin/authentication/FlowTest.java

@@ -20,12 +20,15 @@ package org.keycloak.testsuite.admin.authentication;
 import org.junit.Assert;
 import org.junit.Assume;
 import org.junit.Test;
+import org.keycloak.common.Profile;
+import org.keycloak.common.util.StreamUtil;
 import org.keycloak.events.admin.OperationType;
 import org.keycloak.events.admin.ResourceType;
 import org.keycloak.representations.idm.AuthenticationExecutionExportRepresentation;
 import org.keycloak.representations.idm.AuthenticationExecutionInfoRepresentation;
 import org.keycloak.representations.idm.AuthenticationFlowRepresentation;
 import org.keycloak.representations.idm.ComponentRepresentation;
+import org.keycloak.testsuite.ProfileAssume;
 import org.keycloak.testsuite.util.AdminEventPaths;
 import org.keycloak.testsuite.util.ContainerAssume;
 
@@ -36,6 +39,9 @@ import javax.ws.rs.NotFoundException;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.Response.Status;
 import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.Charset;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -46,6 +52,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.keycloak.testsuite.util.Matchers.body;
 import static org.keycloak.testsuite.util.Matchers.statusCodeIs;
@@ -440,7 +447,8 @@ public class FlowTest extends AbstractAuthenticationTest {
     }
 
     @Test
-    public void failWithLongDescription() {
+    public void failWithLongDescription() throws IOException {
+        ProfileAssume.assumeFeatureDisabled(Profile.Feature.MAP_STORAGE);
         ContainerAssume.assumeAuthServerQuarkus();
         AuthenticationFlowRepresentation rep = authMgmtResource.getFlows().stream()
                 .filter(new Predicate<AuthenticationFlowRepresentation>() {
@@ -462,11 +470,11 @@ public class FlowTest extends AbstractAuthenticationTest {
 
         try {
             authMgmtResource.updateFlow(rep.getId(), rep);
+            fail("Should fail because the description is too long");
         } catch (InternalServerErrorException isee) {
             try (Response response = isee.getResponse()) {
                 assertEquals(500, response.getStatus());
-                assertEquals(0, response.getLength());
+                assertFalse(StreamUtil.readString((InputStream) response.getEntity(), Charset.forName("UTF-8")).toLowerCase().contains("exception"));
-                assertEquals(0, ByteArrayInputStream.class.cast(response.getEntity()).available());
             }
         } catch (Exception e) {
             fail("Unexpected exception");

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java

@@ -335,6 +335,8 @@ public class UserSessionProviderTest extends AbstractTestRealmKeycloakTest {
 
         Map<String, Set<String>> clientSessionsKept = new HashMap<>();
         for (UserSessionModel s : sessions) {
+            // session associated with the model was closed, load it by id into a new session
+            s = session.sessions().getUserSession(realm, s.getId());
             Set<String> clientUUIDS = new HashSet<>(s.getAuthenticatedClientSessions().keySet());
             clientUUIDS.remove(thirdPartyClientUUID); // This client will be later removed, hence his clientSessions too
             clientSessionsKept.put(s.getId(), clientUUIDS);