[KEYCLOAK-14147] - Request filter refactoring

Co-authored-by: Stian Thorgersen <stian@redhat.com>
Co-authored-by: Martin Kanis <mkanis@redhat.com>

distribution/feature-packs/server-feature-pack/src/main/resources/modules/system/layers/keycloak/org/keycloak/keycloak-server-subsystem/main/server-war/WEB-INF/web.xml

@@ -48,13 +48,13 @@
     </listener>
 
     <filter>
-        <filter-name>Keycloak Session Management</filter-name>
+        <filter-name>Client Connection Filter</filter-name>
-        <filter-class>org.keycloak.provider.wildfly.KeycloakSessionServletFilter</filter-class>
+        <filter-class>org.keycloak.provider.wildfly.WildFlyClientConnectionServletFilter</filter-class>
         <async-supported>true</async-supported>
     </filter>
 
     <filter-mapping>
-        <filter-name>Keycloak Session Management</filter-name>
+        <filter-name>Client Connection Filter</filter-name>
         <url-pattern>/*</url-pattern>
     </filter-mapping>
 

distribution/feature-packs/server-feature-pack/src/main/resources/modules/system/layers/keycloak/org/keycloak/keycloak-services/main/module.xml

@@ -59,6 +59,7 @@
         <module name="com.fasterxml.jackson.core.jackson-core"/>
         <module name="com.fasterxml.jackson.core.jackson-annotations"/>
         <module name="com.fasterxml.jackson.core.jackson-databind"/>
+        <module name="com.fasterxml.jackson.datatype.jackson-datatype-jdk8"/>
         <module name="com.fasterxml.jackson.jaxrs.jackson-jaxrs-json-provider"/>
         <module name="com.google.zxing.core"/>
         <module name="com.google.zxing.javase"/>

pom.xml

@@ -420,6 +420,11 @@
                 <artifactId>jackson-module-jaxb-annotations</artifactId>
                 <version>${jackson.version}</version>
             </dependency>
+            <dependency>
+                <groupId>com.fasterxml.jackson.datatype</groupId>
+                <artifactId>jackson-datatype-jdk8</artifactId>
+                <version>${jackson.version}</version>
+            </dependency>
             <dependency>
                 <groupId>com.fasterxml.jackson.jaxrs</groupId>
                 <artifactId>jackson-jaxrs-json-provider</artifactId>

quarkus/deployment/pom.xml

@@ -59,6 +59,11 @@
             <artifactId>quarkus-jdbc-mariadb-deployment</artifactId>
             <version>${quarkus.version}</version>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-vertx-web-deployment</artifactId>
+            <version>${quarkus.version}</version>
+        </dependency>
         <dependency>
             <groupId>io.quarkus</groupId>
             <artifactId>quarkus-bootstrap-core</artifactId>

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

@@ -20,6 +20,7 @@ import org.keycloak.provider.KeycloakDeploymentInfo;
 import org.keycloak.provider.ProviderFactory;
 import org.keycloak.provider.ProviderManager;
 import org.keycloak.provider.Spi;
+import org.keycloak.provider.quarkus.QuarkusClientConnectionFilter;
 import org.keycloak.runtime.KeycloakRecorder;
 import org.keycloak.transaction.JBossJtaTransactionManagerLookup;
 
@@ -30,6 +31,7 @@ import io.quarkus.deployment.annotations.ExecutionTime;
 import io.quarkus.deployment.annotations.Record;
 import io.quarkus.deployment.builditem.FeatureBuildItem;
 import io.quarkus.hibernate.orm.deployment.PersistenceUnitDescriptorBuildItem;
+import io.quarkus.vertx.http.deployment.FilterBuildItem;
 
 class KeycloakProcessor {
 
@@ -57,9 +59,12 @@ class KeycloakProcessor {
     }
 
     /**
-     * <p>Load the built-in provider factories during build time so we don't spend time looking up them at runtime.
+     * <p>
+     * Load the built-in provider factories during build time so we don't spend time looking up them at runtime.
      * 
-     * <p>User-defined providers are going to be loaded at startup</p>
+     * <p>
+     * User-defined providers are going to be loaded at startup
+     * </p>
      */
     @Record(ExecutionTime.STATIC_INIT)
     @BuildStep
@@ -69,7 +74,8 @@ class KeycloakProcessor {
 
     private Map<Spi, Set<Class<? extends ProviderFactory>>> loadBuiltInFactories() {
         ProviderManager pm = new ProviderManager(
-                KeycloakDeploymentInfo.create().services(), Thread.currentThread().getContextClassLoader(), Config.scope().getArray("providers"));
+                KeycloakDeploymentInfo.create().services(), Thread.currentThread().getContextClassLoader(),
+                Config.scope().getArray("providers"));
         Map<Spi, Set<Class<? extends ProviderFactory>>> result = new HashMap<>();
 
         for (Spi spi : pm.loadSpis()) {
@@ -92,4 +98,9 @@ class KeycloakProcessor {
 
         return result;
     }
+
+    @BuildStep
+    void initializeRouter(BuildProducer<FilterBuildItem> routes) {
+        routes.produce(new FilterBuildItem(new QuarkusClientConnectionFilter(), FilterBuildItem.AUTHORIZATION - 10));
+    }
 }

quarkus/pom.xml

@@ -31,11 +31,11 @@
     <packaging>pom</packaging>
 
     <properties>
-        <quarkus.version>1.5.1.Final</quarkus.version>
+        <quarkus.version>1.6.0.CR1</quarkus.version>
-        <resteasy.version>4.5.3.Final</resteasy.version>
+        <resteasy.version>4.5.5.Final</resteasy.version>
-        <jackson.version>2.10.2</jackson.version>
+        <jackson.version>2.10.4</jackson.version>
         <jackson.databind.version>${jackson.version}</jackson.databind.version>
-        <hibernate.version>5.4.16.Final</hibernate.version>
+        <hibernate.version>5.4.18.Final</hibernate.version>
         <snakeyaml.version>1.20</snakeyaml.version>
 
         <surefire-plugin.version>2.22.0</surefire-plugin.version>

quarkus/runtime/pom.xml

@@ -410,10 +410,6 @@
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-lang3</artifactId>
         </dependency>
-        <dependency>
-            <groupId>io.quarkus</groupId>
-            <artifactId>quarkus-vertx-web</artifactId>
-        </dependency>
         <dependency>
             <groupId>org.liquibase</groupId>
             <artifactId>liquibase-core</artifactId>

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

@@ -10,6 +10,9 @@ import javax.inject.Inject;
 import javax.persistence.EntityManagerFactory;
 import javax.ws.rs.ApplicationPath;
 
+import org.jboss.resteasy.plugins.server.servlet.ResteasyContextParameters;
+import org.jboss.resteasy.spi.ResteasyDeployment;
+import org.keycloak.common.util.Resteasy;
 import org.keycloak.models.utils.PostMigrationEvent;
 import org.keycloak.services.resources.KeycloakApplication;
 import org.keycloak.services.resources.QuarkusWelcomeResource;
@@ -30,6 +33,10 @@ public class QuarkusKeycloakApplication extends KeycloakApplication {
 
     @Override
     public Set<Object> getSingletons() {
+        //TODO: a temporary hack for https://github.com/quarkusio/quarkus/issues/9647, we need to disable the sanitizer to avoid
+        // escaping text/html responses from the server
+        Resteasy.getContextData(ResteasyDeployment.class).setProperty(ResteasyContextParameters.RESTEASY_DISABLE_HTML_SANITIZER, Boolean.TRUE);
+
         HashSet<Object> singletons = new HashSet<>(super.getSingletons().stream().filter(new Predicate<Object>() {
             @Override
             public boolean test(Object o) {

quarkus/runtime/src/main/java/org/keycloak/provider/quarkus/KeycloakConfigSourceProvider.java

@@ -22,9 +22,6 @@ import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 
-import io.quarkus.runtime.configuration.DeploymentProfileConfigSource;
-import io.quarkus.runtime.configuration.ExpandingConfigSource;
-
 import org.eclipse.microprofile.config.spi.ConfigSource;
 import org.eclipse.microprofile.config.spi.ConfigSourceProvider;
 import org.jboss.logging.Logger;
@@ -58,23 +55,16 @@ public class KeycloakConfigSourceProvider implements ConfigSourceProvider {
         }
 
         if (filePath != null) {
-            sources.add(wrap(new KeycloakPropertiesConfigSource.InFileSystem(filePath)));
+            sources.add(new KeycloakPropertiesConfigSource.InFileSystem(filePath));
         }
 
         // fall back to the default configuration within the server classpath
         if (sources.isEmpty()) {
             log.debug("Loading the default server configuration");
-            sources.add(wrap(new KeycloakPropertiesConfigSource.InJar()));
+            sources.add(new KeycloakPropertiesConfigSource.InJar());
         }
 
         return sources;
 
     }
-
-    private ConfigSource wrap(ConfigSource source) {
-        return ExpandingConfigSource.wrapper(new ExpandingConfigSource.Cache())
-            .compose(DeploymentProfileConfigSource.wrapper())
-            .apply(source);
-    }
-
 }

quarkus/runtime/src/main/java/org/keycloak/provider/quarkus/QuarkusClientConnectionFilter.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2019 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.provider.quarkus;
+
+import org.jboss.resteasy.plugins.server.servlet.ResteasyContextParameters;
+import org.jboss.resteasy.spi.ResteasyDeployment;
+import org.keycloak.common.ClientConnection;
+import org.keycloak.common.util.Resteasy;
+import org.keycloak.services.filters.AbstractClientConnectionFilter;
+
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Handler;
+import io.vertx.core.http.HttpServerRequest;
+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.
+ * 
+ * <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 QuarkusClientConnectionFilter extends AbstractClientConnectionFilter 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
+    };
+
+    @Override
+    public void handle(RoutingContext context) {
+        ClientConnection clientConnection = createClientConnection(context.request());
+
+        // 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 -> filter(clientConnection, (session) -> {
+            try {
+                context.next();
+                promise.complete();
+            } catch (Exception cause) {
+                promise.fail(cause);
+                // re-throw so that the any exception is handled from parent
+                throw new RuntimeException(cause);
+            }
+        }), EMPTY_RESULT);
+    }
+
+    private ClientConnection createClientConnection(HttpServerRequest request) {
+        return new ClientConnection() {
+            @Override
+            public String getRemoteAddr() {
+                return request.remoteAddress().host();
+            }
+
+            @Override
+            public String getRemoteHost() {
+                return request.remoteAddress().host();
+            }
+
+            @Override
+            public int getRemotePort() {
+                return request.remoteAddress().port();
+            }
+
+            @Override
+            public String getLocalAddr() {
+                return request.localAddress().host();
+            }
+
+            @Override
+            public int getLocalPort() {
+                return request.localAddress().port();
+            }
+        };
+    }
+}

quarkus/runtime/src/main/java/org/keycloak/provider/quarkus/QuarkusFilter.java

@@ -1,118 +0,0 @@
-/*
- * Copyright 2019 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.provider.quarkus;
-
-import io.vertx.core.http.HttpServerRequest;
-import io.vertx.ext.web.RoutingContext;
-import org.jboss.resteasy.plugins.server.servlet.ResteasyContextParameters;
-import org.jboss.resteasy.spi.ResteasyDeployment;
-import org.keycloak.common.ClientConnection;
-import org.keycloak.common.util.Resteasy;
-import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.KeycloakSessionFactory;
-import org.keycloak.models.KeycloakTransaction;
-import org.keycloak.services.resources.KeycloakApplication;
-
-import javax.annotation.Priority;
-import javax.inject.Inject;
-import javax.ws.rs.container.ContainerRequestContext;
-import javax.ws.rs.container.ContainerResponseContext;
-import javax.ws.rs.container.PreMatching;
-import javax.ws.rs.ext.Provider;
-import java.io.IOException;
-
-
-@PreMatching
-@Provider
-@Priority(1)
-public class QuarkusFilter implements javax.ws.rs.container.ContainerRequestFilter,
-        javax.ws.rs.container.ContainerResponseFilter  {
-
-    @Inject
-    KeycloakApplication keycloakApplication;
-    
-    @Inject
-    RoutingContext routingContext;
-
-    public QuarkusFilter() {
-        //TODO: a temporary hack for https://github.com/quarkusio/quarkus/issues/9647, we need to disable the sanitizer to avoid
-        // escaping text/html responses from the server
-        Resteasy.getContextData(ResteasyDeployment.class).setProperty(ResteasyContextParameters.RESTEASY_DISABLE_HTML_SANITIZER, Boolean.TRUE);
-    }
-
-    @Override
-    public void filter(ContainerRequestContext containerRequestContext) {
-        KeycloakSessionFactory sessionFactory = keycloakApplication.getSessionFactory();
-        KeycloakSession session = sessionFactory.create();
-
-        Resteasy.pushContext(KeycloakSession.class, session);
-        HttpServerRequest request = routingContext.request();
-
-        session.getContext().setConnection(createConnection(request));
-        Resteasy.pushContext(ClientConnection.class, session.getContext().getConnection());
-
-        KeycloakTransaction tx = session.getTransactionManager();
-        Resteasy.pushContext(KeycloakTransaction.class, tx);
-
-        tx.begin();
-    }
-
-    @Override
-    public void filter(ContainerRequestContext containerRequestContext, ContainerResponseContext containerResponseContext) throws IOException {
-        //End the session and clear context
-        KeycloakSession session = Resteasy.getContextData(KeycloakSession.class);
-
-        // KeycloakTransactionCommitter is responsible for committing the transaction, but if an exception is thrown it's not invoked and transaction
-        // should be rolled back
-        if (session.getTransactionManager() != null && session.getTransactionManager().isActive()) {
-            session.getTransactionManager().rollback();
-        }
-
-        session.close();
-        Resteasy.clearContextData();
-    }
-
-    private ClientConnection createConnection(HttpServerRequest request) {
-        return new ClientConnection() {
-            @Override
-            public String getRemoteAddr() {
-                return request.remoteAddress().host();
-            }
-
-            @Override
-            public String getRemoteHost() {
-                return request.remoteAddress().host();
-            }
-
-            @Override
-            public int getRemotePort() {
-                return request.remoteAddress().port();
-            }
-
-            @Override
-            public String getLocalAddr() {
-                return request.localAddress().host();
-            }
-
-            @Override
-            public int getLocalPort() {
-                return request.localAddress().port();
-            }
-        };
-    }
-}

quarkus/runtime/src/main/java/org/keycloak/provider/quarkus/QuarkusLifecycleObserver.java

@@ -21,9 +21,11 @@ import io.quarkus.runtime.ShutdownEvent;
 import io.quarkus.runtime.StartupEvent;
 import javax.enterprise.context.ApplicationScoped;
 import javax.enterprise.event.Observes;
-import javax.inject.Inject;
 
+import org.jboss.resteasy.plugins.server.servlet.ResteasyContextParameters;
+import org.jboss.resteasy.spi.ResteasyDeployment;
 import org.keycloak.Config;
+import org.keycloak.common.util.Resteasy;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
 import org.keycloak.models.KeycloakTransactionManager;
@@ -38,18 +40,13 @@ public class QuarkusLifecycleObserver {
     private static final String KEYCLOAK_ADMIN_ENV_VAR = "KEYCLOAK_ADMIN";
     private static final String KEYCLOAK_ADMIN_PASSWORD_ENV_VAR = "KEYCLOAK_ADMIN_PASSWORD";
 
-    @Inject
-    KeycloakApplication application;
-
     void onStartupEvent(@Observes StartupEvent event) {
-
         Runnable startupHook = ((QuarkusPlatform) Platform.getPlatform()).startupHook;
 
         if (startupHook != null) {
             startupHook.run();
             createAdminUser();
         }
-
     }
 
     void onShutdownEvent(@Observes ShutdownEvent event) {
@@ -70,7 +67,7 @@ public class QuarkusLifecycleObserver {
             return;
         }
 
-        KeycloakSessionFactory sessionFactory = application.getSessionFactory();
+        KeycloakSessionFactory sessionFactory = KeycloakApplication.getSessionFactory();
         KeycloakSession session = sessionFactory.create();
         KeycloakTransactionManager transaction = session.getTransactionManager();
 

server-spi/src/main/java/org/keycloak/models/KeycloakContext.java

@@ -65,8 +65,6 @@ public interface KeycloakContext {
 
     ClientConnection getConnection();
 
-    void setConnection(ClientConnection connection);
-
     Locale resolveLocale(UserModel user);
     
     /**

services/pom.xml

@@ -132,6 +132,10 @@
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-annotations</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.datatype</groupId>
+            <artifactId>jackson-datatype-jdk8</artifactId>
+        </dependency>
         <dependency>
             <groupId>com.fasterxml.woodstox</groupId>
             <artifactId>woodstox-core</artifactId>

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

@@ -118,12 +118,7 @@ public class DefaultKeycloakContext implements KeycloakContext {
 
     @Override
     public ClientConnection getConnection() {
-        return connection;
+        return getContextObject(ClientConnection.class);
-    }
-
-    @Override
-    public void setConnection(ClientConnection connection) {
-        this.connection = connection;
     }
 
     @Override
@@ -140,4 +135,5 @@ public class DefaultKeycloakContext implements KeycloakContext {
     public void setAuthenticationSession(AuthenticationSessionModel authenticationSession) {
         this.authenticationSession = authenticationSession;
     }
+
 }

services/src/main/java/org/keycloak/services/error/KeycloakErrorHandler.java

@@ -46,9 +46,6 @@ public class KeycloakErrorHandler implements ExceptionMapper<Throwable> {
 
     public static final String UNCAUGHT_SERVER_ERROR_TEXT = "Uncaught server error";
 
-    @Context
-    private KeycloakSession session;
-
     @Context
     private HttpHeaders headers;
 
@@ -57,7 +54,8 @@ public class KeycloakErrorHandler implements ExceptionMapper<Throwable> {
 
     @Override
     public Response toResponse(Throwable throwable) {
-        KeycloakTransaction tx = Resteasy.getContextData(KeycloakTransaction.class);
+        KeycloakSession session = Resteasy.getContextData(KeycloakSession.class);
+        KeycloakTransaction tx = session.getTransactionManager();
         tx.setRollbackOnly();
 
         int statusCode = getStatusCode(throwable);
@@ -78,14 +76,14 @@ public class KeycloakErrorHandler implements ExceptionMapper<Throwable> {
         }
 
         try {
-            RealmModel realm = resolveRealm();
+            RealmModel realm = resolveRealm(session);
 
             Theme theme = session.theme().getTheme(Theme.Type.LOGIN);
 
             Locale locale = session.getContext().resolveLocale(null);
 
             FreeMarkerUtil freeMarker = new FreeMarkerUtil();
-            Map<String, Object> attributes = initAttributes(realm, theme, locale, statusCode);
+            Map<String, Object> attributes = initAttributes(session, realm, theme, locale, statusCode);
 
             String templateName = "error.ftl";
 
@@ -121,7 +119,7 @@ public class KeycloakErrorHandler implements ExceptionMapper<Throwable> {
         return "unknown_error";
     }
 
-    private RealmModel resolveRealm() {
+    private RealmModel resolveRealm(KeycloakSession session) {
         String path = session.getContext().getUri().getPath();
         Matcher m = realmNamePattern.matcher(path);
         String realmName;
@@ -142,7 +140,7 @@ public class KeycloakErrorHandler implements ExceptionMapper<Throwable> {
         return realm;
     }
 
-    private Map<String, Object> initAttributes(RealmModel realm, Theme theme, Locale locale, int statusCode) throws IOException {
+    private Map<String, Object> initAttributes(KeycloakSession session, RealmModel realm, Theme theme, Locale locale, int statusCode) throws IOException {
         Map<String, Object> attributes = new HashMap<>();
         Properties messagesBundle = theme.getMessages(locale);
 

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

@@ -0,0 +1,66 @@
+package org.keycloak.services.filters;
+
+/*
+ * Copyright 2020 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.
+ */
+
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+import org.keycloak.common.ClientConnection;
+import org.keycloak.common.util.Resteasy;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.KeycloakTransactionManager;
+import org.keycloak.services.resources.KeycloakApplication;
+
+
+public abstract class AbstractClientConnectionFilter {
+
+    public void filter(ClientConnection clientConnection, Consumer<KeycloakSession> next) {
+        KeycloakSessionFactory sessionFactory = KeycloakApplication.getSessionFactory();
+        KeycloakSession session = sessionFactory.create();
+
+        KeycloakTransactionManager tx = session.getTransactionManager();
+        tx.begin();
+
+        try {
+            Resteasy.pushContext(ClientConnection.class, clientConnection);
+            Resteasy.pushContext(KeycloakSession.class, session);
+
+            next.accept(session);
+        } catch (Exception e) {
+            tx.setRollbackOnly();
+            throw new RuntimeException(e);
+        } finally {
+            close(session);
+        }
+    }
+
+    public void close(KeycloakSession session) {
+        KeycloakTransactionManager tx = session.getTransactionManager();
+        if (tx.isActive()) {
+            if (tx.getRollbackOnly()) {
+                tx.rollback();
+            } else {
+                tx.commit();
+            }
+        }
+
+        session.close();
+    }
+}

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

@@ -23,17 +23,19 @@ import org.keycloak.models.KeycloakSession;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerResponseContext;
 import javax.ws.rs.container.ContainerResponseFilter;
+import javax.ws.rs.container.PreMatching;
 
 /**
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
  */
+@PreMatching
 public class KeycloakSecurityHeadersFilter implements ContainerResponseFilter {
 
     @Override
-    public void filter(ContainerRequestContext requestContext, ContainerResponseContext responseContext) {
+    public void filter(ContainerRequestContext containerRequestContext, ContainerResponseContext containerResponseContext) {
         KeycloakSession session = Resteasy.getContextData(KeycloakSession.class);
-        SecurityHeadersProvider securityHeadersProvider = session.getProvider(SecurityHeadersProvider.class);
-        securityHeadersProvider.addHeaders(requestContext, responseContext);
-    }
 
+        SecurityHeadersProvider securityHeadersProvider = session.getProvider(SecurityHeadersProvider.class);
+        securityHeadersProvider.addHeaders(containerRequestContext, containerResponseContext);
+    }
 }

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

@@ -1,50 +0,0 @@
-/*
- * Copyright 2016 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.services.filters;
-
-import org.keycloak.common.util.Resteasy;
-import org.keycloak.models.KeycloakTransaction;
-
-import javax.annotation.Priority;
-import javax.ws.rs.container.ContainerRequestContext;
-import javax.ws.rs.container.ContainerResponseContext;
-import javax.ws.rs.container.ContainerResponseFilter;
-import java.io.IOException;
-
-/**
- * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
- */
-@Priority(2)
-public class KeycloakTransactionCommitter implements ContainerResponseFilter {
-
-    @Override
-    public void filter(ContainerRequestContext containerRequestContext, ContainerResponseContext containerResponseContext) throws IOException {
-        KeycloakTransaction tx = Resteasy.getContextData(KeycloakTransaction.class);
-        if (tx != null && tx.isActive()) {
-            if (tx.getRollbackOnly()) {
-                tx.rollback();
-            } else {
-                tx.commit();
-            }
-        }
-    }
-
-}

services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java

@@ -43,7 +43,6 @@ import org.keycloak.services.DefaultKeycloakSessionFactory;
 import org.keycloak.services.ServicesLogger;
 import org.keycloak.services.error.KeycloakErrorHandler;
 import org.keycloak.services.filters.KeycloakSecurityHeadersFilter;
-import org.keycloak.services.filters.KeycloakTransactionCommitter;
 import org.keycloak.services.managers.ApplianceBootstrap;
 import org.keycloak.services.managers.RealmManager;
 import org.keycloak.services.managers.UserStorageSyncManager;
@@ -89,7 +88,7 @@ public class KeycloakApplication extends Application {
     protected Set<Object> singletons = new HashSet<Object>();
     protected Set<Class<?>> classes = new HashSet<Class<?>>();
 
-    protected KeycloakSessionFactory sessionFactory;
+    protected static KeycloakSessionFactory sessionFactory;
 
     public KeycloakApplication() {
 
@@ -100,9 +99,6 @@ public class KeycloakApplication extends Application {
 
             loadConfig();
 
-            Resteasy.pushDefaultContextObject(KeycloakApplication.class, this);
-            Resteasy.pushContext(KeycloakApplication.class, this); // for injection
-
             singletons.add(new RobotsResource());
             singletons.add(new RealmsResource());
             singletons.add(new AdminRoot());
@@ -110,10 +106,9 @@ public class KeycloakApplication extends Application {
             classes.add(JsResource.class);
 
             classes.add(KeycloakSecurityHeadersFilter.class);
-            classes.add(KeycloakTransactionCommitter.class);
             classes.add(KeycloakErrorHandler.class);
 
-            singletons.add(new ObjectMapperResolver(Boolean.parseBoolean(System.getProperty("keycloak.jsonPrettyPrint", "false"))));
+            singletons.add(new ObjectMapperResolver());
             singletons.add(new WelcomeResource());
 
             platform.onStartup(this::startup);
@@ -281,7 +276,7 @@ public class KeycloakApplication extends Application {
         }
     }
 
-    public KeycloakSessionFactory getSessionFactory() {
+    public static KeycloakSessionFactory getSessionFactory() {
         return sessionFactory;
     }
 

services/src/main/java/org/keycloak/services/resources/admin/ClientResource.java

@@ -56,7 +56,6 @@ import org.keycloak.services.clientregistration.policy.RegistrationAuth;
 import org.keycloak.services.managers.ClientManager;
 import org.keycloak.services.managers.RealmManager;
 import org.keycloak.services.managers.ResourceAdminManager;
-import org.keycloak.services.resources.KeycloakApplication;
 import org.keycloak.services.resources.admin.permissions.AdminPermissionEvaluator;
 import org.keycloak.services.resources.admin.permissions.AdminPermissionManagement;
 import org.keycloak.services.resources.admin.permissions.AdminPermissions;
@@ -105,16 +104,9 @@ public class ClientResource {
     protected ClientModel client;
     protected KeycloakSession session;
 
-    @Context
-    protected KeycloakApplication keycloak;
-
     @Context
     protected ClientConnection clientConnection;
 
-    protected KeycloakApplication getKeycloakApplication() {
-        return keycloak;
-    }
-
     public ClientResource(RealmModel realm, AdminPermissionEvaluator auth, ClientModel clientModel, KeycloakSession session, AdminEventBuilder adminEvent) {
         this.realm = realm;
         this.auth = auth;

services/src/main/java/org/keycloak/services/resources/admin/RealmsAdminResource.java

@@ -18,7 +18,6 @@ package org.keycloak.services.resources.admin;
 
 import org.jboss.logging.Logger;
 import org.jboss.resteasy.annotations.cache.NoCache;
-import javax.ws.rs.NotFoundException;
 import org.jboss.resteasy.spi.ResteasyProviderFactory;
 import org.keycloak.common.ClientConnection;
 import org.keycloak.models.AdminRoles;
@@ -34,12 +33,12 @@ import org.keycloak.representations.idm.RealmRepresentation;
 import org.keycloak.services.ErrorResponse;
 import org.keycloak.services.ForbiddenException;
 import org.keycloak.services.managers.RealmManager;
-import org.keycloak.services.resources.KeycloakApplication;
 import org.keycloak.services.resources.admin.permissions.AdminPermissionEvaluator;
 import org.keycloak.services.resources.admin.permissions.AdminPermissions;
 
 import javax.ws.rs.Consumes;
 import javax.ws.rs.GET;
+import javax.ws.rs.NotFoundException;
 import javax.ws.rs.POST;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
@@ -68,9 +67,6 @@ public class RealmsAdminResource {
     @Context
     protected KeycloakSession session;
 
-    @Context
-    protected KeycloakApplication keycloak;
-
     @Context
     protected ClientConnection clientConnection;
 

services/src/main/java/org/keycloak/services/resources/admin/RoleContainerResource.java

@@ -55,11 +55,12 @@ import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.UriInfo;
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * @resource Roles
@@ -94,13 +95,13 @@ public class RoleContainerResource extends RoleResource {
     @GET
     @NoCache
     @Produces(MediaType.APPLICATION_JSON)
-    public List<RoleRepresentation> getRoles(@QueryParam("search") @DefaultValue("") String search,
+    public Stream<RoleRepresentation> getRoles(@QueryParam("search") @DefaultValue("") String search,
                                                @QueryParam("first") Integer firstResult,
                                                @QueryParam("max") Integer maxResults,
                                                @QueryParam("briefRepresentation") @DefaultValue("true") boolean briefRepresentation) {
         auth.roles().requireList(roleContainer);
 
-        Set<RoleModel> roleModels = new HashSet<RoleModel>();
+        Set<RoleModel> roleModels;
 
         if(search != null && search.trim().length() > 0) {
             roleModels = roleContainer.searchForRoles(search, firstResult, maxResults);
@@ -110,15 +111,10 @@ public class RoleContainerResource extends RoleResource {
             roleModels = roleContainer.getRoles();
         }
 
-        List<RoleRepresentation> roles = new ArrayList<RoleRepresentation>();
+        Function<RoleModel, RoleRepresentation> toRoleRepresentation = briefRepresentation ?
-        for (RoleModel roleModel : roleModels) {
+                ModelToRepresentation::toBriefRepresentation :
-            if(briefRepresentation) {
+                ModelToRepresentation::toRepresentation;
-                roles.add(ModelToRepresentation.toBriefRepresentation(roleModel));  
+        return roleModels.stream().map(toRoleRepresentation);
-            } else {
-                roles.add(ModelToRepresentation.toRepresentation(roleModel));               
-            }
-        }
-        return roles;
     }
 
     /**

services/src/main/java/org/keycloak/services/util/ObjectMapperResolver.java

@@ -18,11 +18,16 @@
 package org.keycloak.services.util;
 
 import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.databind.JavaType;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.SerializationFeature;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.fasterxml.jackson.databind.type.TypeFactory;
+import com.fasterxml.jackson.datatype.jdk8.StreamSerializer;
 
 import javax.ws.rs.ext.ContextResolver;
 import javax.ws.rs.ext.Provider;
+import java.util.stream.Stream;
 
 /**
  * Any class with package org.jboss.resteasy.skeleton.key will use NON_DEFAULT inclusion
@@ -34,9 +39,16 @@ import javax.ws.rs.ext.Provider;
 public class ObjectMapperResolver implements ContextResolver<ObjectMapper> {
     protected ObjectMapper mapper = new ObjectMapper();
 
-    public ObjectMapperResolver(boolean indent) {
+    public ObjectMapperResolver() {
+        JavaType type = TypeFactory.unknownType();
+        JavaType streamType = mapper.getTypeFactory().constructParametricType(Stream.class, type);
+
+        SimpleModule module = new SimpleModule();
+        module.addSerializer(new StreamSerializer(streamType, type));
+        mapper.registerModule(module);
+
         mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
-        if (indent) {
+        if (Boolean.parseBoolean(System.getProperty("keycloak.jsonPrettyPrint", "false"))) {
             mapper.enable(SerializationFeature.INDENT_OUTPUT);
         }
     }

testsuite/integration-arquillian/servers/auth-server/undertow/src/main/java/org/keycloak/testsuite/arquillian/undertow/KeycloakOnUndertow.java

@@ -50,7 +50,7 @@ import org.keycloak.services.managers.ApplianceBootstrap;
 import org.keycloak.services.resources.KeycloakApplication;
 import org.keycloak.testsuite.JsonConfigProviderFactory;
 import org.keycloak.testsuite.KeycloakServer;
-import org.keycloak.testsuite.TestKeycloakSessionServletFilter;
+import org.keycloak.testsuite.UndertowClientConnectionServletFilter;
 import org.keycloak.testsuite.utils.tls.TLSUtils;
 import org.keycloak.testsuite.utils.undertow.UndertowDeployerHelper;
 import org.keycloak.testsuite.utils.undertow.UndertowWarClassLoader;
@@ -99,7 +99,7 @@ public class KeycloakOnUndertow implements DeployableContainer<KeycloakOnUnderto
         di.setDefaultServletConfig(new DefaultServletConfig(true));
         di.addWelcomePage("theme/keycloak/welcome/resources/index.html");
 
-        FilterInfo filter = Servlets.filter("SessionFilter", TestKeycloakSessionServletFilter.class);
+        FilterInfo filter = Servlets.filter("SessionFilter", UndertowClientConnectionServletFilter.class);
         di.addFilter(filter);
         di.addFilterUrlMapping("SessionFilter", "/*", DispatcherType.REQUEST);
         filter.setAsyncSupported(true);
@@ -203,8 +203,7 @@ public class KeycloakOnUndertow implements DeployableContainer<KeycloakOnUnderto
 
         DeploymentInfo di = createAuthServerDeploymentInfo();
         undertow.deploy(di);
-        ResteasyDeployment deployment = (ResteasyDeployment) di.getServletContextAttributes().get(ResteasyDeployment.class.getName());
+        sessionFactory = KeycloakApplication.getSessionFactory();
-        sessionFactory = ((KeycloakApplication) deployment.getApplication()).getSessionFactory();
 
         setupDevConfig();
 

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

@@ -19,6 +19,8 @@ package org.keycloak.testsuite;
 
 import org.junit.After;
 import org.keycloak.admin.client.resource.RealmResource;
+import org.keycloak.common.ClientConnection;
+import org.keycloak.common.util.Resteasy;
 import org.keycloak.common.util.reflections.Reflections;
 import org.keycloak.events.Details;
 import org.keycloak.models.KeycloakSession;
@@ -120,7 +122,7 @@ public abstract class AbstractTestRealmKeycloakTest extends AbstractKeycloakTest
 
     /** KEYCLOAK-12065 Inherit Client Connection from parent session **/
     public static KeycloakSession inheritClientConnection(KeycloakSession parentSession, KeycloakSession currentSession) {
-        currentSession.getContext().setConnection(parentSession.getContext().getConnection());
+        Resteasy.pushContext(ClientConnection.class, parentSession.getContext().getConnection());
         return currentSession;
     }
 }

testsuite/utils/src/main/java/org/keycloak/testsuite/KeycloakServer.java

@@ -405,7 +405,7 @@ public class KeycloakServer {
             // KEYCLOAK-14178
             deployment.setProperty(ResteasyContextParameters.RESTEASY_DISABLE_HTML_SANITIZER, true);
 
-            FilterInfo filter = Servlets.filter("SessionFilter", TestKeycloakSessionServletFilter.class);
+            FilterInfo filter = Servlets.filter("SessionFilter", UndertowClientConnectionServletFilter.class);
             filter.setAsyncSupported(true);
 
             di.addFilter(filter);
@@ -413,7 +413,7 @@ public class KeycloakServer {
 
             server.deploy(di);
 
-            sessionFactory = ((KeycloakApplication) deployment.getApplication()).getSessionFactory();
+            sessionFactory = KeycloakApplication.getSessionFactory();
 
             setupDevConfig();
 

testsuite/utils/src/main/java/org/keycloak/testsuite/KeycloakSessionServletFilter.java

@@ -1,138 +0,0 @@
-/*
- * Copyright 2016 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.testsuite;
-
-import java.io.IOException;
-
-import javax.servlet.AsyncEvent;
-import javax.servlet.AsyncListener;
-import javax.servlet.Filter;
-import javax.servlet.FilterChain;
-import javax.servlet.FilterConfig;
-import javax.servlet.ServletException;
-import javax.servlet.ServletRequest;
-import javax.servlet.ServletResponse;
-import javax.servlet.http.HttpServletRequest;
-
-import org.keycloak.common.ClientConnection;
-import org.keycloak.common.util.Resteasy;
-import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.KeycloakSessionFactory;
-import org.keycloak.models.KeycloakTransaction;
-
-/**
- * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
- * @version $Revision: 1 $
- */
-public class KeycloakSessionServletFilter implements Filter {
-
-    @Override
-    public void init(FilterConfig filterConfig) throws ServletException {
-    }
-
-    @Override
-    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
-        servletRequest.setCharacterEncoding("UTF-8");
-
-        final HttpServletRequest request = (HttpServletRequest)servletRequest;
-
-        KeycloakSessionFactory sessionFactory = (KeycloakSessionFactory) servletRequest.getServletContext().getAttribute(KeycloakSessionFactory.class.getName());
-        KeycloakSession session = sessionFactory.create();
-        Resteasy.pushContext(KeycloakSession.class, session);
-        ClientConnection connection = new ClientConnection() {
-            @Override
-            public String getRemoteAddr() {
-                return request.getRemoteAddr();
-            }
-
-            @Override
-            public String getRemoteHost() {
-                return request.getRemoteHost();
-            }
-
-            @Override
-            public int getRemotePort() {
-                return request.getRemotePort();
-            }
-
-            @Override
-            public String getLocalAddr() {
-                return request.getLocalAddr();
-            }
-
-            @Override
-            public int getLocalPort() {
-                return request.getLocalPort();
-            }
-        };
-        session.getContext().setConnection(connection);
-        Resteasy.pushContext(ClientConnection.class, connection);
-
-        KeycloakTransaction tx = session.getTransactionManager();
-        Resteasy.pushContext(KeycloakTransaction.class, tx);
-        tx.begin();
-
-        try {
-            filterChain.doFilter(servletRequest, servletResponse);
-        } finally {
-            if (servletRequest.isAsyncStarted()) {
-                servletRequest.getAsyncContext().addListener(createAsyncLifeCycleListener(session));
-            } else {
-                closeSession(session);
-            }
-        }
-    }
-
-    private AsyncListener createAsyncLifeCycleListener(final KeycloakSession session) {
-        return new AsyncListener() {
-            @Override
-            public void onComplete(AsyncEvent event) {
-                closeSession(session);
-            }
-
-            @Override
-            public void onTimeout(AsyncEvent event) {
-                closeSession(session);
-            }
-
-            @Override
-            public void onError(AsyncEvent event) {
-                closeSession(session);
-            }
-
-            @Override
-            public void onStartAsync(AsyncEvent event) {
-            }
-        };
-    }
-
-    private void closeSession(KeycloakSession session) {
-        // KeycloakTransactionCommitter is responsible for committing the transaction, but if an exception is thrown it's not invoked and transaction
-        // should be rolled back
-        if (session.getTransactionManager() != null && session.getTransactionManager().isActive()) {
-            session.getTransactionManager().rollback();
-        }
-
-        session.close();
-        Resteasy.clearContextData();
-    }
-
-    @Override
-    public void destroy() {
-    }
-}

testsuite/utils/src/main/java/org/keycloak/testsuite/TestKeycloakSessionServletFilter.java

@@ -1,49 +0,0 @@
-/*
- * Copyright 2019 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.testsuite;
-
-import javax.servlet.FilterChain;
-import javax.servlet.ServletException;
-import javax.servlet.ServletRequest;
-import javax.servlet.ServletResponse;
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletRequestWrapper;
-import java.io.IOException;
-
-/**
- * @author <a href="mailto:psilva@redhat.com">Pedro Igor</a>
- */
-public class TestKeycloakSessionServletFilter extends KeycloakSessionServletFilter {
-
-    @Override
-    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain)
-            throws IOException, ServletException {
-        super.doFilter(new HttpServletRequestWrapper((HttpServletRequest) servletRequest) {
-            @Override
-            public String getRemoteAddr() {
-                String forwardedFor = getHeader("X-Forwarded-For");
-
-                if (forwardedFor != null) {
-                    return forwardedFor;
-                }
-
-                return super.getRemoteAddr();
-            }
-        }, servletResponse, filterChain);
-    }
-}

testsuite/utils/src/main/java/org/keycloak/testsuite/TestPlatform.java

@@ -17,21 +17,13 @@
 
 package org.keycloak.testsuite;
 
-import org.keycloak.common.util.Resteasy;
-import org.keycloak.models.KeycloakSessionFactory;
 import org.keycloak.platform.PlatformProvider;
-import org.keycloak.services.resources.KeycloakApplication;
-
-import javax.servlet.ServletContext;
 
 public class TestPlatform implements PlatformProvider {
 
     @Override
     public void onStartup(Runnable startupHook) {
         startupHook.run();
-        KeycloakApplication keycloakApplication = Resteasy.getContextData(KeycloakApplication.class);
-        ServletContext context = Resteasy.getContextData(ServletContext.class);
-        context.setAttribute(KeycloakSessionFactory.class.getName(),  keycloakApplication.getSessionFactory());
     }
 
     @Override

testsuite/utils/src/main/java/org/keycloak/testsuite/UndertowClientConnectionServletFilter.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2016 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.testsuite;
+
+import java.io.UnsupportedEncodingException;
+
+import javax.servlet.Filter;
+import javax.servlet.FilterChain;
+import javax.servlet.FilterConfig;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
+
+import org.keycloak.common.ClientConnection;
+import org.keycloak.services.filters.AbstractClientConnectionFilter;
+
+public class UndertowClientConnectionServletFilter extends AbstractClientConnectionFilter implements Filter {
+
+    @Override
+    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain)
+            throws UnsupportedEncodingException {
+        servletRequest.setCharacterEncoding("UTF-8");
+        final HttpServletRequest request = (HttpServletRequest) servletRequest;
+
+        filter(createClientConnection(request), (session) -> {
+            try {
+                filterChain.doFilter(servletRequest, servletResponse);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    private ClientConnection createClientConnection(HttpServletRequest request) {
+        return new ClientConnection() {
+            @Override
+            public String getRemoteAddr() {
+                String forwardedFor = request.getHeader("X-Forwarded-For");
+
+                if (forwardedFor != null) {
+                    return forwardedFor;
+                }
+
+                return request.getRemoteAddr();
+            }
+
+            @Override
+            public String getRemoteHost() {
+                return request.getRemoteHost();
+            }
+
+            @Override
+            public int getRemotePort() {
+                return request.getRemotePort();
+            }
+
+            @Override
+            public String getLocalAddr() {
+                return request.getLocalAddr();
+            }
+
+            @Override
+            public int getLocalPort() {
+                return request.getLocalPort();
+            }
+        };
+    }
+
+    @Override
+    public void init(FilterConfig filterConfig) {
+    }
+
+    @Override
+    public void destroy() {
+    }
+}

wildfly/extensions/src/main/java/org/keycloak/provider/wildfly/KeycloakSessionServletFilter.java

@@ -1,138 +0,0 @@
-/*
- * Copyright 2016 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.provider.wildfly;
-
-import java.io.IOException;
-
-import javax.servlet.AsyncEvent;
-import javax.servlet.AsyncListener;
-import javax.servlet.Filter;
-import javax.servlet.FilterChain;
-import javax.servlet.FilterConfig;
-import javax.servlet.ServletException;
-import javax.servlet.ServletRequest;
-import javax.servlet.ServletResponse;
-import javax.servlet.http.HttpServletRequest;
-
-import org.keycloak.common.ClientConnection;
-import org.keycloak.common.util.Resteasy;
-import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.KeycloakSessionFactory;
-import org.keycloak.models.KeycloakTransaction;
-
-/**
- * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
- * @version $Revision: 1 $
- */
-public class KeycloakSessionServletFilter implements Filter {
-
-    @Override
-    public void init(FilterConfig filterConfig) throws ServletException {
-    }
-
-    @Override
-    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
-        servletRequest.setCharacterEncoding("UTF-8");
-
-        final HttpServletRequest request = (HttpServletRequest)servletRequest;
-
-        KeycloakSessionFactory sessionFactory = (KeycloakSessionFactory) servletRequest.getServletContext().getAttribute(KeycloakSessionFactory.class.getName());
-        KeycloakSession session = sessionFactory.create();
-        Resteasy.pushContext(KeycloakSession.class, session);
-        ClientConnection connection = new ClientConnection() {
-            @Override
-            public String getRemoteAddr() {
-                return request.getRemoteAddr();
-            }
-
-            @Override
-            public String getRemoteHost() {
-                return request.getRemoteHost();
-            }
-
-            @Override
-            public int getRemotePort() {
-                return request.getRemotePort();
-            }
-
-            @Override
-            public String getLocalAddr() {
-                return request.getLocalAddr();
-            }
-
-            @Override
-            public int getLocalPort() {
-                return request.getLocalPort();
-            }
-        };
-        session.getContext().setConnection(connection);
-        Resteasy.pushContext(ClientConnection.class, connection);
-
-        KeycloakTransaction tx = session.getTransactionManager();
-        Resteasy.pushContext(KeycloakTransaction.class, tx);
-        tx.begin();
-
-        try {
-            filterChain.doFilter(servletRequest, servletResponse);
-        } finally {
-            if (servletRequest.isAsyncStarted()) {
-                servletRequest.getAsyncContext().addListener(createAsyncLifeCycleListener(session));
-            } else {
-                closeSession(session);
-            }
-        }
-    }
-
-    private AsyncListener createAsyncLifeCycleListener(final KeycloakSession session) {
-        return new AsyncListener() {
-            @Override
-            public void onComplete(AsyncEvent event) {
-                closeSession(session);
-            }
-
-            @Override
-            public void onTimeout(AsyncEvent event) {
-                closeSession(session);
-            }
-
-            @Override
-            public void onError(AsyncEvent event) {
-                closeSession(session);
-            }
-
-            @Override
-            public void onStartAsync(AsyncEvent event) {
-            }
-        };
-    }
-
-    private void closeSession(KeycloakSession session) {
-        // KeycloakTransactionCommitter is responsible for committing the transaction, but if an exception is thrown it's not invoked and transaction
-        // should be rolled back
-        if (session.getTransactionManager() != null && session.getTransactionManager().isActive()) {
-            session.getTransactionManager().rollback();
-        }
-
-        session.close();
-        Resteasy.clearContextData();
-    }
-
-    @Override
-    public void destroy() {
-    }
-}

wildfly/extensions/src/main/java/org/keycloak/provider/wildfly/WildFlyClientConnectionServletFilter.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2016 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.provider.wildfly;
+
+import java.io.UnsupportedEncodingException;
+
+import javax.servlet.Filter;
+import javax.servlet.FilterChain;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
+
+import org.keycloak.common.ClientConnection;
+import org.keycloak.services.filters.AbstractClientConnectionFilter;
+
+public class WildFlyClientConnectionServletFilter extends AbstractClientConnectionFilter implements Filter {
+
+    @Override
+    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain)
+            throws UnsupportedEncodingException {
+        servletRequest.setCharacterEncoding("UTF-8");
+        ClientConnection clientConnection = createConnection((HttpServletRequest) servletRequest);
+
+        filter(clientConnection, (session) -> {
+            try {
+                filterChain.doFilter(servletRequest, servletResponse);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    private ClientConnection createConnection(HttpServletRequest request) {
+        return new ClientConnection() {
+            @Override
+            public String getRemoteAddr() {
+                return request.getRemoteAddr();
+            }
+
+            @Override
+            public String getRemoteHost() {
+                return request.getRemoteHost();
+            }
+
+            @Override
+            public int getRemotePort() {
+                return request.getRemotePort();
+            }
+
+            @Override
+            public String getLocalAddr() {
+                return request.getLocalAddr();
+            }
+
+            @Override
+            public int getLocalPort() {
+                return request.getLocalPort();
+            }
+        };
+    }
+}

wildfly/extensions/src/main/java/org/keycloak/provider/wildfly/WildflyPlatform.java

@@ -17,13 +17,8 @@
 
 package org.keycloak.provider.wildfly;
 
-import org.keycloak.common.util.Resteasy;
-import org.keycloak.models.KeycloakSessionFactory;
 import org.keycloak.platform.PlatformProvider;
 import org.keycloak.services.ServicesLogger;
-import org.keycloak.services.resources.KeycloakApplication;
-
-import javax.servlet.ServletContext;
 
 public class WildflyPlatform implements PlatformProvider {
 
@@ -32,9 +27,6 @@ public class WildflyPlatform implements PlatformProvider {
     @Override
     public void onStartup(Runnable startupHook) {
         startupHook.run();
-        KeycloakApplication keycloakApplication = Resteasy.getContextData(KeycloakApplication.class);
-        ServletContext context = Resteasy.getContextData(ServletContext.class);
-        context.setAttribute(KeycloakSessionFactory.class.getName(),  keycloakApplication.getSessionFactory());
     }
 
     @Override