keycloak-aplcache

diff --git a/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/IdMapperUpdaterSessionListener.java b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/IdMapperUpdaterSessionListener.java
new file mode 100644
index 0000000..692413e
--- /dev/null
+++ b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/IdMapperUpdaterSessionListener.java
@@ -0,0 +1,103 @@
+/*
+ * 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.undertow;
+
+import org.keycloak.adapters.saml.SamlSession;
+import org.keycloak.adapters.spi.SessionIdMapper;
+
+import io.undertow.server.HttpServerExchange;
+import io.undertow.server.session.Session;
+import io.undertow.server.session.SessionListener;
+import java.util.Objects;
+
+/**
+ *
+ * @author hmlnarik
+ */
+public class IdMapperUpdaterSessionListener implements SessionListener {
+
+    private final SessionIdMapper idMapper;
+
+    public IdMapperUpdaterSessionListener(SessionIdMapper idMapper) {
+        this.idMapper = idMapper;
+    }
+
+    @Override
+    public void sessionCreated(Session session, HttpServerExchange exchange) {
+        Object value = session.getAttribute(SamlSession.class.getName());
+        map(session.getId(), value);
+    }
+
+    @Override
+    public void sessionDestroyed(Session session, HttpServerExchange exchange, SessionDestroyedReason reason) {
+        if (reason != SessionDestroyedReason.UNDEPLOY) {
+            unmap(session.getId(), session.getAttribute(SamlSession.class.getName()));
+        }
+    }
+
+    @Override
+    public void attributeAdded(Session session, String name, Object value) {
+        if (Objects.equals(name, SamlSession.class.getName())) {
+            map(session.getId(), value);
+        }
+    }
+
+    @Override
+    public void attributeUpdated(Session session, String name, Object newValue, Object oldValue) {
+        if (Objects.equals(name, SamlSession.class.getName())) {
+            unmap(session.getId(), oldValue);
+            map(session.getId(), newValue);
+        }
+    }
+
+    @Override
+    public void attributeRemoved(Session session, String name, Object oldValue) {
+        if (Objects.equals(name, SamlSession.class.getName())) {
+            unmap(session.getId(), oldValue);
+        }
+    }
+
+    @Override
+    public void sessionIdChanged(Session session, String oldSessionId) {
+        Object value = session.getAttribute(SamlSession.class.getName());
+        if (value != null) {
+            unmap(oldSessionId, value);
+            map(session.getId(), value);
+        }
+    }
+
+    private void map(String sessionId, Object value) {
+        if (! (value instanceof SamlSession) || sessionId == null) {
+            return;
+        }
+        SamlSession account = (SamlSession) value;
+
+        idMapper.map(account.getSessionIndex(), account.getPrincipal().getSamlSubject(), sessionId);
+    }
+
+    private void unmap(String sessionId, Object value) {
+        if (! (value instanceof SamlSession) || sessionId == null) {
+            return;
+        }
+
+        SamlSession samlSession = (SamlSession) value;
+        if (samlSession.getSessionIndex() != null) {
+            idMapper.removeSession(sessionId);
+        }
+    }
+
+}

diff --git a/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/SamlServletExtension.java b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/SamlServletExtension.java
index 0e6a1a1..4985dfb 100755
--- a/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/SamlServletExtension.java
+++ b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/SamlServletExtension.java
@@ -154,7 +154,7 @@ public class SamlServletExtension implements ServletExtension {
         servletContext.setAttribute(SamlDeploymentContext.class.getName(), deploymentContext);
         UndertowUserSessionManagement userSessionManagement = new UndertowUserSessionManagement();
         final ServletSamlAuthMech mech = createAuthMech(deploymentInfo, deploymentContext, userSessionManagement);
-
+        mech.addTokenStoreUpdaters(deploymentInfo);
 
         // setup handlers
 

diff --git a/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlAuthMech.java b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlAuthMech.java
index 8818171..b0a7339 100755
--- a/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlAuthMech.java
+++ b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlAuthMech.java
@@ -21,35 +21,89 @@ import io.undertow.security.api.SecurityContext;
 import io.undertow.server.HttpServerExchange;
 import io.undertow.servlet.handlers.ServletRequestContext;
 import io.undertow.util.Headers;
+
 import org.keycloak.adapters.saml.SamlDeployment;
 import org.keycloak.adapters.saml.SamlDeploymentContext;
 import org.keycloak.adapters.saml.SamlSessionStore;
-import org.keycloak.adapters.spi.HttpFacade;
-import org.keycloak.adapters.spi.InMemorySessionIdMapper;
-import org.keycloak.adapters.spi.SessionIdMapper;
+import org.keycloak.adapters.spi.*;
 import org.keycloak.adapters.undertow.ServletHttpFacade;
 import org.keycloak.adapters.undertow.UndertowHttpFacade;
 import org.keycloak.adapters.undertow.UndertowUserSessionManagement;
 
+import io.undertow.servlet.api.DeploymentInfo;
 import javax.servlet.RequestDispatcher;
 import javax.servlet.ServletException;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
 import java.io.IOException;
+import java.lang.reflect.*;
+import java.util.Map;
+import org.jboss.logging.Logger;
 
 /**
  * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
  * @version $Revision: 1 $
  */
 public class ServletSamlAuthMech extends AbstractSamlAuthMech {
+
+    private static final Logger LOG = Logger.getLogger(ServletSamlAuthMech.class);
+
     protected SessionIdMapper idMapper = new InMemorySessionIdMapper();
+    protected SessionIdMapperUpdater idMapperUpdater = SessionIdMapperUpdater.DIRECT;
+
     public ServletSamlAuthMech(SamlDeploymentContext deploymentContext, UndertowUserSessionManagement sessionManagement, String errorPage) {
         super(deploymentContext, sessionManagement, errorPage);
     }
 
+    public void addTokenStoreUpdaters(DeploymentInfo deploymentInfo) {
+        deploymentInfo.addSessionListener(new IdMapperUpdaterSessionListener(idMapper));    // This takes care of HTTP sessions manipulated locally
+        SessionIdMapperUpdater updater = SessionIdMapperUpdater.EXTERNAL;
+
+        try {
+            Map<String, String> initParameters = deploymentInfo.getInitParameters();
+            String idMapperSessionUpdaterClasses = initParameters == null
+              ? null
+              : initParameters.get("keycloak.sessionIdMapperUpdater.classes");
+            if (idMapperSessionUpdaterClasses == null) {
+                return;
+            }
+
+            for (String clazz : idMapperSessionUpdaterClasses.split("\\s*,\\s*")) {
+                if (! clazz.isEmpty()) {
+                    updater = invokeAddTokenStoreUpdaterMethod(clazz, deploymentInfo, updater);
+                }
+            }
+        } finally {
+            setIdMapperUpdater(updater);
+        }
+    }
+
+    private SessionIdMapperUpdater invokeAddTokenStoreUpdaterMethod(String idMapperSessionUpdaterClass, DeploymentInfo deploymentInfo,
+      SessionIdMapperUpdater previousIdMapperUpdater) {
+        try {
+            Class<?> clazz = deploymentInfo.getClassLoader().loadClass(idMapperSessionUpdaterClass);
+            Method addTokenStoreUpdatersMethod = clazz.getMethod("addTokenStoreUpdaters", DeploymentInfo.class, SessionIdMapper.class, SessionIdMapperUpdater.class);
+            if (! Modifier.isStatic(addTokenStoreUpdatersMethod.getModifiers())
+              || ! Modifier.isPublic(addTokenStoreUpdatersMethod.getModifiers())
+              || ! SessionIdMapperUpdater.class.isAssignableFrom(addTokenStoreUpdatersMethod.getReturnType())) {
+                LOG.errorv("addTokenStoreUpdaters method in class {0} has to be public static. Ignoring class.", idMapperSessionUpdaterClass);
+                return previousIdMapperUpdater;
+            }
+
+            LOG.debugv("Initializing sessionIdMapperUpdater class {0}", idMapperSessionUpdaterClass);
+            return (SessionIdMapperUpdater) addTokenStoreUpdatersMethod.invoke(null, deploymentInfo, idMapper, previousIdMapperUpdater);
+        } catch (ClassNotFoundException | NoSuchMethodException | SecurityException ex) {
+            LOG.warnv(ex, "Cannot use sessionIdMapperUpdater class {0}", idMapperSessionUpdaterClass);
+            return previousIdMapperUpdater;
+        } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException ex) {
+            LOG.warnv(ex, "Cannot use {0}.addTokenStoreUpdaters(DeploymentInfo, SessionIdMapper) method", idMapperSessionUpdaterClass);
+            return previousIdMapperUpdater;
+        }
+    }
+
     @Override
     protected SamlSessionStore getTokenStore(HttpServerExchange exchange, HttpFacade facade, SamlDeployment deployment, SecurityContext securityContext) {
-        return new ServletSamlSessionStore(exchange, sessionManagement, securityContext, idMapper, deployment);
+        return new ServletSamlSessionStore(exchange, sessionManagement, securityContext, idMapper, idMapperUpdater, deployment);
     }
 
     @Override
@@ -84,5 +138,11 @@ public class ServletSamlAuthMech extends AbstractSamlAuthMech {
         return null;
     }
 
+    public SessionIdMapperUpdater getIdMapperUpdater() {
+        return idMapperUpdater;
+    }
 
+    protected void setIdMapperUpdater(SessionIdMapperUpdater idMapperUpdater) {
+        this.idMapperUpdater = idMapperUpdater;
+    }
 }

diff --git a/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlSessionStore.java b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlSessionStore.java
index 72acda5..2bf2369 100755
--- a/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlSessionStore.java
+++ b/adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlSessionStore.java
@@ -24,11 +24,13 @@ import io.undertow.server.session.SessionManager;
 import io.undertow.servlet.handlers.ServletRequestContext;
 import io.undertow.servlet.spec.HttpSessionImpl;
 import org.jboss.logging.Logger;
+
 import org.keycloak.adapters.saml.SamlDeployment;
 import org.keycloak.adapters.saml.SamlSession;
 import org.keycloak.adapters.saml.SamlSessionStore;
 import org.keycloak.adapters.saml.SamlUtil;
 import org.keycloak.adapters.spi.SessionIdMapper;
+import org.keycloak.adapters.spi.SessionIdMapperUpdater;
 import org.keycloak.adapters.undertow.ChangeSessionId;
 import org.keycloak.adapters.undertow.SavedRequest;
 import org.keycloak.adapters.undertow.ServletHttpFacade;
@@ -44,6 +46,8 @@ import java.util.List;
 import java.util.Set;
 
 /**
+ * Session store manipulation methods per single HTTP exchange.
+ *
  * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
  * @version $Revision: 1 $
  */
@@ -55,17 +59,20 @@ public class ServletSamlSessionStore implements SamlSessionStore {
     private final UndertowUserSessionManagement sessionManagement;
     private final SecurityContext securityContext;
     private final SessionIdMapper idMapper;
+    private final SessionIdMapperUpdater idMapperUpdater;
     protected final SamlDeployment deployment;
 
 
     public ServletSamlSessionStore(HttpServerExchange exchange, UndertowUserSessionManagement sessionManagement,
                                    SecurityContext securityContext,
-                                   SessionIdMapper idMapper, SamlDeployment deployment) {
+                                   SessionIdMapper idMapper, SessionIdMapperUpdater idMapperUpdater,
+                                   SamlDeployment deployment) {
         this.exchange = exchange;
         this.sessionManagement = sessionManagement;
         this.securityContext = securityContext;
         this.idMapper = idMapper;
         this.deployment = deployment;
+        this.idMapperUpdater = idMapperUpdater;
     }
 
     @Override
@@ -97,7 +104,7 @@ public class ServletSamlSessionStore implements SamlSessionStore {
             SamlSession samlSession = (SamlSession)session.getAttribute(SamlSession.class.getName());
             if (samlSession != null) {
                 if (samlSession.getSessionIndex() != null) {
-                    idMapper.removeSession(session.getId());
+                    idMapperUpdater.removeSession(idMapper, session.getId());
                 }
                 session.removeAttribute(SamlSession.class.getName());
             }
@@ -113,7 +120,7 @@ public class ServletSamlSessionStore implements SamlSessionStore {
             ids.addAll(sessions);
             logoutSessionIds(ids);
             for (String id : ids) {
-                idMapper.removeSession(id);
+                idMapperUpdater.removeSession(idMapper, id);
             }
         }
 
@@ -127,7 +134,7 @@ public class ServletSamlSessionStore implements SamlSessionStore {
              String sessionId = idMapper.getSessionFromSSO(id);
              if (sessionId != null) {
                  sessionIds.add(sessionId);
-                 idMapper.removeSession(sessionId);
+                 idMapperUpdater.removeSession(idMapper, sessionId);
              }
 
         }
@@ -177,7 +184,7 @@ public class ServletSamlSessionStore implements SamlSessionStore {
         session.setAttribute(SamlSession.class.getName(), account);
         sessionManagement.login(servletRequestContext.getDeployment().getSessionManager());
         String sessionId = changeSessionId(session);
-        idMapper.map(account.getSessionIndex(), account.getPrincipal().getSamlSubject(), sessionId);
+        idMapperUpdater.map(idMapper, account.getSessionIndex(), account.getPrincipal().getSamlSubject(), sessionId);
 
     }
 

diff --git a/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlAuthMech.java b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlAuthMech.java
index bdf0606..ae4e242 100755
--- a/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlAuthMech.java
+++ b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlAuthMech.java
@@ -37,6 +37,6 @@ public class WildflySamlAuthMech extends ServletSamlAuthMech {
 
     @Override
     protected SamlSessionStore getTokenStore(HttpServerExchange exchange, HttpFacade facade, SamlDeployment deployment, SecurityContext securityContext) {
-        return new WildflySamlSessionStore(exchange, sessionManagement, securityContext, idMapper, deployment);
+        return new WildflySamlSessionStore(exchange, sessionManagement, securityContext, idMapper, getIdMapperUpdater(), deployment);
     }
 }

diff --git a/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlSessionStore.java b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlSessionStore.java
index b4c213a..5f8d717 100755
--- a/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlSessionStore.java
+++ b/adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/WildflySamlSessionStore.java
@@ -23,6 +23,7 @@ import org.keycloak.adapters.saml.SamlDeployment;
 import org.keycloak.adapters.saml.SamlSession;
 import org.keycloak.adapters.saml.undertow.ServletSamlSessionStore;
 import org.keycloak.adapters.spi.SessionIdMapper;
+import org.keycloak.adapters.spi.SessionIdMapperUpdater;
 import org.keycloak.adapters.undertow.UndertowUserSessionManagement;
 
 /**
@@ -31,8 +32,10 @@ import org.keycloak.adapters.undertow.UndertowUserSessionManagement;
  */
 public class WildflySamlSessionStore extends ServletSamlSessionStore {
     public WildflySamlSessionStore(HttpServerExchange exchange, UndertowUserSessionManagement sessionManagement,
-                                   SecurityContext securityContext, SessionIdMapper idMapper, SamlDeployment resolvedDeployment) {
-        super(exchange, sessionManagement, securityContext, idMapper, resolvedDeployment);
+                                   SecurityContext securityContext,
+                                   SessionIdMapper idMapper, SessionIdMapperUpdater idMapperUpdater,
+                                   SamlDeployment resolvedDeployment) {
+        super(exchange, sessionManagement, securityContext, idMapper, idMapperUpdater, resolvedDeployment);
     }
 
     @Override

diff --git a/adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapper.java b/adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapper.java
index 3b467d7..b50e37b 100755
--- a/adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapper.java
+++ b/adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapper.java
@@ -24,15 +24,43 @@ import java.util.Set;
  * @version $Revision: 1 $
  */
 public interface SessionIdMapper {
+    /**
+     * Returns {@code true} if the mapper contains mapping for the given HTTP session ID.
+     * @param id
+     * @return
+     */
     boolean hasSession(String id);
 
+    /**
+     * Clears all mappings from this mapper.
+     */
     void clear();
 
+    /**
+     * Returns set of HTTP session IDs for the given principal.
+     * @param principal Principal
+     * @return
+     */
     Set<String> getUserSessions(String principal);
 
+    /**
+     * Returns HTTP session ID from the given user session ID.
+     * @param sso User session ID
+     * @return
+     */
     String getSessionFromSSO(String sso);
 
+    /**
+     * Establishes mapping between user session ID, principal and HTTP session ID.
+     * @param sso User session ID
+     * @param principal Principal
+     * @param session HTTP session ID
+     */
     void map(String sso, String principal, String session);
 
+    /**
+     * Removes mappings for the given HTTP session ID.
+     * @param session HTTP session ID.
+     */
     void removeSession(String session);
 }

diff --git a/adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapperUpdater.java b/adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapperUpdater.java
new file mode 100644
index 0000000..a0ea2ff
--- /dev/null
+++ b/adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapperUpdater.java
@@ -0,0 +1,73 @@
+/*
+ * 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.spi;
+
+/**
+ * Classes implementing this interface represent a mechanism for updating {@link SessionIdMapper} entries.
+ * @author hmlnarik
+ */
+public interface SessionIdMapperUpdater {
+    /**
+     * {@link SessionIdMapper} entries are updated directly.
+     */
+    public static final SessionIdMapperUpdater DIRECT = new SessionIdMapperUpdater() {
+        @Override public void clear(SessionIdMapper idMapper) {
+            idMapper.clear();
+        }
+
+        @Override public void map(SessionIdMapper idMapper, String sso, String principal, String httpSessionId) {
+            idMapper.map(sso, principal, httpSessionId);
+        }
+
+        @Override public void removeSession(SessionIdMapper idMapper, String httpSessionId) {
+            idMapper.removeSession(httpSessionId);
+        }
+    };
+
+    /**
+     * Only HTTP session is manipulated with, {@link SessionIdMapper} entries are not updated by this updater and
+     * they have to be updated by some other means, e.g. by some listener of HTTP session changes.
+     */
+    public static final SessionIdMapperUpdater EXTERNAL = new SessionIdMapperUpdater() {
+        @Override public void clear(SessionIdMapper idMapper) { }
+
+        @Override public void map(SessionIdMapper idMapper, String sso, String principal, String httpSessionId) { }
+
+        @Override public void removeSession(SessionIdMapper idMapper, String httpSessionId) { }
+    };
+
+    /**
+     * Delegates to {@link SessionIdMapper#clear} method..
+     */
+    public abstract void clear(SessionIdMapper idMapper);
+
+    /**
+     * Delegates to {@link SessionIdMapper#map} method.
+     * @param idMapper Mapper
+     * @param sso User session ID
+     * @param principal Principal
+     * @param session HTTP session ID
+     */
+    public abstract void map(SessionIdMapper idMapper, String sso, String principal, String session);
+
+    /**
+     * Delegates to {@link SessionIdMapper#removeSession} method.
+     * @param idMapper Mapper
+     * @param session HTTP session ID.
+     */
+    public abstract void removeSession(SessionIdMapper idMapper, String session);
+}