keycloak-aplcache

diff --git a/adapters/saml/wildfly/wildfly-adapter/pom.xml b/adapters/saml/wildfly/wildfly-adapter/pom.xml
index 09c126e..2c073b4 100755
--- a/adapters/saml/wildfly/wildfly-adapter/pom.xml
+++ b/adapters/saml/wildfly/wildfly-adapter/pom.xml
@@ -67,6 +67,10 @@
             <artifactId>keycloak-jboss-adapter-core</artifactId>
         </dependency>
         <dependency>
+            <groupId>org.infinispan</groupId>
+            <artifactId>infinispan-core</artifactId>
+        </dependency>
+        <dependency>
             <groupId>org.picketbox</groupId>
             <artifactId>picketbox</artifactId>
             <version>4.0.20.Final</version>

diff --git a/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/InfinispanSessionCacheIdMapperUpdater.java b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/InfinispanSessionCacheIdMapperUpdater.java
new file mode 100644
index 0000000..489d1d5
--- /dev/null
+++ b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/InfinispanSessionCacheIdMapperUpdater.java
@@ -0,0 +1,107 @@
+/*
+ * Copyright 2017 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.adapters.saml.wildfly.infinispan;
+
+import org.keycloak.adapters.spi.SessionIdMapper;
+import org.keycloak.adapters.spi.SessionIdMapperUpdater;
+
+import io.undertow.servlet.api.DeploymentInfo;
+import java.util.*;
+import javax.naming.InitialContext;
+import javax.naming.NamingException;
+import org.infinispan.Cache;
+import org.infinispan.configuration.cache.CacheMode;
+import org.infinispan.configuration.cache.Configuration;
+import org.infinispan.manager.EmbeddedCacheManager;
+import org.jboss.logging.Logger;
+
+/**
+ *
+ * @author hmlnarik
+ */
+public class InfinispanSessionCacheIdMapperUpdater {
+
+    private static final Logger LOG = Logger.getLogger(InfinispanSessionCacheIdMapperUpdater.class);
+
+    public static final String DEFAULT_CACHE_CONTAINER_JNDI_NAME = "java:jboss/infinispan/container/web";
+
+    private static final String DEPLOYMENT_CACHE_CONTAINER_JNDI_NAME_PARAM_NAME = "keycloak.sessionIdMapperUpdater.infinispan.cacheContainerJndi";
+    private static final String DEPLOYMENT_CACHE_NAME_PARAM_NAME = "keycloak.sessionIdMapperUpdater.infinispan.deploymentCacheName";
+    private static final String SSO_CACHE_NAME_PARAM_NAME = "keycloak.sessionIdMapperUpdater.infinispan.cacheName";
+
+    public static SessionIdMapperUpdater addTokenStoreUpdaters(DeploymentInfo deploymentInfo, SessionIdMapper mapper, SessionIdMapperUpdater previousIdMapperUpdater) {
+       boolean distributable = Objects.equals(
+          deploymentInfo.getSessionManagerFactory().getClass().getName(),
+          "org.wildfly.clustering.web.undertow.session.DistributableSessionManagerFactory"
+        );
+
+        if (! distributable) {
+            LOG.warnv("Deployment {0} does not use supported distributed session cache mechanism", deploymentInfo.getDeploymentName());
+            return previousIdMapperUpdater;
+        }
+
+        Map<String, String> initParameters = deploymentInfo.getInitParameters();
+        String cacheContainerLookup = (initParameters != null && initParameters.get(DEPLOYMENT_CACHE_CONTAINER_JNDI_NAME_PARAM_NAME) != null)
+          ? initParameters.get(DEPLOYMENT_CACHE_CONTAINER_JNDI_NAME_PARAM_NAME)
+          : DEFAULT_CACHE_CONTAINER_JNDI_NAME;
+        boolean deploymentSessionCacheNamePreset = initParameters != null && initParameters.get(DEPLOYMENT_CACHE_NAME_PARAM_NAME) != null;
+        String deploymentSessionCacheName = deploymentSessionCacheNamePreset
+          ? initParameters.get(DEPLOYMENT_CACHE_NAME_PARAM_NAME)
+          : deploymentInfo.getDeploymentName();
+        boolean ssoCacheNamePreset = initParameters != null && initParameters.get(SSO_CACHE_NAME_PARAM_NAME) != null;
+        String ssoCacheName = ssoCacheNamePreset
+          ? initParameters.get(SSO_CACHE_NAME_PARAM_NAME)
+          : deploymentSessionCacheName + ".ssoCache";
+
+        try {
+            EmbeddedCacheManager cacheManager = (EmbeddedCacheManager) new InitialContext().lookup(cacheContainerLookup);
+
+            Configuration ssoCacheConfiguration = cacheManager.getCacheConfiguration(ssoCacheName);
+            if (ssoCacheConfiguration == null) {
+                Configuration cacheConfiguration = cacheManager.getCacheConfiguration(deploymentSessionCacheName);
+                if (cacheConfiguration == null) {
+                    LOG.debugv("Using default cache container configuration for SSO cache. lookup={0}, looked up configuration of cache={1}", cacheContainerLookup, deploymentSessionCacheName);
+                    ssoCacheConfiguration = cacheManager.getDefaultCacheConfiguration();
+                } else {
+                    LOG.debugv("Using distributed HTTP session cache configuration for SSO cache. lookup={0}, configuration taken from cache={1}", cacheContainerLookup, deploymentSessionCacheName);
+                    ssoCacheConfiguration = cacheConfiguration;
+                    cacheManager.defineConfiguration(ssoCacheName, ssoCacheConfiguration);
+                }
+            } else {
+                LOG.debugv("Using custom configuration of SSO cache. lookup={0}, cache name={1}", cacheContainerLookup, ssoCacheName);
+            }
+
+            CacheMode ssoCacheMode = ssoCacheConfiguration.clustering().cacheMode();
+            if (ssoCacheMode != CacheMode.REPL_ASYNC && ssoCacheMode != CacheMode.REPL_SYNC) {
+                LOG.warnv("SSO cache mode is {0}, it is recommended to use replicated mode instead", ssoCacheConfiguration.clustering().cacheModeString());
+            }
+
+            Cache<String, String[]> ssoCache = cacheManager.getCache(ssoCacheName, true);
+            ssoCache.addListener(new SsoSessionCacheListener(mapper));
+
+            LOG.debugv("Added distributed SSO session cache, lookup={0}, cache name={1}", cacheContainerLookup, deploymentSessionCacheName);
+
+            SsoCacheSessionIdMapperUpdater updater = new SsoCacheSessionIdMapperUpdater(ssoCache, previousIdMapperUpdater);
+            deploymentInfo.addSessionListener(updater);
+
+            return updater;
+        } catch (NamingException ex) {
+            LOG.warnv("Failed to obtain distributed session cache container, lookup={0}", cacheContainerLookup);
+            return previousIdMapperUpdater;
+        }
+    }
+}

diff --git a/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoCacheSessionIdMapperUpdater.java b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoCacheSessionIdMapperUpdater.java
new file mode 100644
index 0000000..e2b8ba2
--- /dev/null
+++ b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoCacheSessionIdMapperUpdater.java
@@ -0,0 +1,96 @@
+/*
+ * Copyright 2017 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.adapters.saml.wildfly.infinispan;
+
+import org.keycloak.adapters.saml.SamlSession;
+import org.keycloak.adapters.spi.SessionIdMapper;
+import org.keycloak.adapters.spi.SessionIdMapperUpdater;
+
+import io.undertow.server.HttpServerExchange;
+import io.undertow.server.session.Session;
+import io.undertow.server.session.SessionListener;
+import org.infinispan.Cache;
+
+/**
+ *
+ * @author hmlnarik
+ */
+public class SsoCacheSessionIdMapperUpdater implements SessionIdMapperUpdater, SessionListener {
+
+    private final SessionIdMapperUpdater delegate;
+    /**
+     * Cache where key is a HTTP session ID, and value is a pair (user session ID, principal name) of Strings.
+     */
+    private final Cache<String, String[]> httpSessionToSsoCache;
+
+    public SsoCacheSessionIdMapperUpdater(Cache<String, String[]> httpSessionToSsoCache, SessionIdMapperUpdater previousIdMapperUpdater) {
+        this.delegate = previousIdMapperUpdater;
+        this.httpSessionToSsoCache = httpSessionToSsoCache;
+    }
+
+    // SessionIdMapperUpdater methods
+
+    @Override
+    public void clear(SessionIdMapper idMapper) {
+        httpSessionToSsoCache.clear();
+        this.delegate.clear(idMapper);
+    }
+
+    @Override
+    public void map(SessionIdMapper idMapper, String sso, String principal, String httpSessionId) {
+        httpSessionToSsoCache.put(httpSessionId, new String[] {sso, principal});
+        this.delegate.map(idMapper, sso, principal, httpSessionId);
+    }
+
+    @Override
+    public void removeSession(SessionIdMapper idMapper, String httpSessionId) {
+        httpSessionToSsoCache.remove(httpSessionId);
+        this.delegate.removeSession(idMapper, httpSessionId);
+    }
+
+    // Undertow HTTP session listener methods
+
+    @Override
+    public void sessionCreated(Session session, HttpServerExchange exchange) {
+    }
+
+    @Override
+    public void sessionDestroyed(Session session, HttpServerExchange exchange, SessionDestroyedReason reason) {
+    }
+
+    @Override
+    public void attributeAdded(Session session, String name, Object value) {
+    }
+
+    @Override
+    public void attributeUpdated(Session session, String name, Object newValue, Object oldValue) {
+    }
+
+    @Override
+    public void attributeRemoved(Session session, String name, Object oldValue) {
+    }
+
+    @Override
+    public void sessionIdChanged(Session session, String oldSessionId) {
+        this.httpSessionToSsoCache.remove(oldSessionId);
+        Object value = session.getAttribute(SamlSession.class.getName());
+        if (value instanceof SamlSession) {
+            SamlSession sess = (SamlSession) value;
+            httpSessionToSsoCache.put(session.getId(), new String[] {sess.getSessionIndex(), sess.getPrincipal().getSamlSubject()});
+        }
+    }
+}

diff --git a/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoSessionCacheListener.java b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoSessionCacheListener.java
new file mode 100644
index 0000000..ccd102e
--- /dev/null
+++ b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoSessionCacheListener.java
@@ -0,0 +1,131 @@
+/*
+ * Copyright 2017 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.adapters.saml.wildfly.infinispan;
+
+import org.keycloak.adapters.spi.SessionIdMapper;
+
+import java.util.*;
+import java.util.concurrent.*;
+import org.infinispan.notifications.Listener;
+import org.infinispan.notifications.cachelistener.annotation.*;
+import org.infinispan.notifications.cachelistener.event.*;
+import org.infinispan.notifications.cachemanagerlistener.annotation.CacheStarted;
+import org.infinispan.notifications.cachemanagerlistener.annotation.CacheStopped;
+import org.infinispan.notifications.cachemanagerlistener.event.CacheStartedEvent;
+import org.infinispan.notifications.cachemanagerlistener.event.CacheStoppedEvent;
+import org.jboss.logging.Logger;
+
+/**
+ *
+ * @author hmlnarik
+ */
+@Listener
+public class SsoSessionCacheListener {
+
+    private static final Logger LOG = Logger.getLogger(SsoSessionCacheListener.class);
+
+    private final ConcurrentMap<String, Queue<Event>> map = new ConcurrentHashMap<>();
+
+    private final SessionIdMapper idMapper;
+
+    private ExecutorService executor = Executors.newSingleThreadExecutor();
+
+    public SsoSessionCacheListener(SessionIdMapper idMapper) {
+        this.idMapper = idMapper;
+    }
+
+    @TransactionRegistered
+    public void startTransaction(TransactionRegisteredEvent event) {
+        map.put(event.getGlobalTransaction().globalId(), new ConcurrentLinkedQueue<Event>());
+    }
+
+    @CacheStarted
+    public void cacheStarted(CacheStartedEvent event) {
+        this.executor = Executors.newSingleThreadExecutor();
+    }
+
+    @CacheStopped
+    public void cacheStopped(CacheStoppedEvent event) {
+        this.executor.shutdownNow();
+    }
+
+    @CacheEntryCreated
+    @CacheEntryRemoved
+    public void addEvent(TransactionalEvent event) {
+        if (event.isPre() == false) {
+            map.get(event.getGlobalTransaction().globalId()).add(event);
+        }
+    }
+
+    @TransactionCompleted
+    public void endTransaction(TransactionCompletedEvent event) {
+        Queue<Event> events = map.remove(event.getGlobalTransaction().globalId());
+
+        if (events == null || ! event.isTransactionSuccessful()) {
+            return;
+        }
+
+        if (event.isOriginLocal()) {
+            // Local events are processed by local HTTP session listener
+            return;
+        }
+
+        for (final Event e : events) {
+            switch (e.getType()) {
+                case CACHE_ENTRY_CREATED:
+                    this.executor.submit(new Runnable() {
+                        @Override public void run() {
+                            cacheEntryCreated((CacheEntryCreatedEvent) e);
+                        }
+                    });
+                    break;
+
+                case CACHE_ENTRY_REMOVED:
+                    this.executor.submit(new Runnable() {
+                        @Override public void run() {
+                            cacheEntryRemoved((CacheEntryRemovedEvent) e);
+                        }
+                    });
+                    break;
+            }
+        }
+    }
+
+    private void cacheEntryCreated(CacheEntryCreatedEvent event) {
+        if (! (event.getKey() instanceof String) || ! (event.getValue() instanceof String[])) {
+            return;
+        }
+        String httpSessionId = (String) event.getKey();
+        String[] value = (String[]) event.getValue();
+        String ssoId = value[0];
+        String principal = value[1];
+
+        LOG.tracev("cacheEntryCreated {0}:{1}", httpSessionId, ssoId);
+
+        this.idMapper.map(ssoId, principal, httpSessionId);
+    }
+
+    private void cacheEntryRemoved(CacheEntryRemovedEvent event) {
+        if (! (event.getKey() instanceof String)) {
+            return;
+        }
+
+        LOG.tracev("cacheEntryRemoved {0}", event.getKey());
+
+        this.idMapper.removeSession((String) event.getKey());
+    }
+}

diff --git a/distribution/saml-adapters/wildfly-adapter/wildfly-modules/src/main/resources/modules/org/keycloak/keycloak-saml-wildfly-adapter/main/module.xml b/distribution/saml-adapters/wildfly-adapter/wildfly-modules/src/main/resources/modules/org/keycloak/keycloak-saml-wildfly-adapter/main/module.xml
index ee00fcc..6cb3c73 100755
--- a/distribution/saml-adapters/wildfly-adapter/wildfly-modules/src/main/resources/modules/org/keycloak/keycloak-saml-wildfly-adapter/main/module.xml
+++ b/distribution/saml-adapters/wildfly-adapter/wildfly-modules/src/main/resources/modules/org/keycloak/keycloak-saml-wildfly-adapter/main/module.xml
@@ -42,6 +42,10 @@
         <module name="org.keycloak.keycloak-saml-adapter-core"/>
         <module name="org.keycloak.keycloak-common"/>
         <module name="org.apache.httpcomponents"/>
+        <module name="org.infinispan"/>
+        <module name="org.infinispan.commons"/>
+        <module name="org.infinispan.cachestore.remote"/>
+        <module name="org.infinispan.client.hotrod"/>
     </dependencies>
 
 </module>