KEYCLOAK-15907 Fix new host in SAML adapter cannot restore session

adapters/saml/as7-eap6/adapter/src/main/java/org/keycloak/adapters/saml/jbossweb/infinispan/SsoCacheSessionIdMapperUpdater.java

@@ -46,6 +46,16 @@ public class SsoCacheSessionIdMapperUpdater implements SessionIdMapperUpdater {
         this.delegate.clear(idMapper);
     }
 
+    @Override
+    public boolean refreshMapping(SessionIdMapper idMapper, String httpSessionId) {
+        String[] ssoAndPrincipal = httpSessionToSsoCache.get(httpSessionId);
+        if (ssoAndPrincipal != null) {
+            this.delegate.map(idMapper, ssoAndPrincipal[0], ssoAndPrincipal[1], httpSessionId);
+            return true;
+        }
+        return false;
+    }
+
     @Override
     public void map(SessionIdMapper idMapper, String sso, String principal, String httpSessionId) {
         httpSessionToSsoCache.put(httpSessionId, new String[] {sso, principal});

adapters/saml/undertow/src/main/java/org/keycloak/adapters/saml/undertow/ServletSamlSessionStore.java

@@ -159,7 +159,7 @@ public class ServletSamlSessionStore implements SamlSessionStore {
             return false;
         }
 
-        if (! idMapper.hasSession(session.getId())) {
+        if (! idMapper.hasSession(session.getId()) && ! idMapperUpdater.refreshMapping(idMapper, session.getId())) {
             log.debugf("Session %s has expired on some other node", session.getId());
             session.removeAttribute(SamlSession.class.getName());
             return false;

adapters/saml/wildfly-elytron/src/main/java/org/keycloak/adapters/saml/elytron/ElytronSamlSessionStore.java

@@ -148,7 +148,7 @@ public class ElytronSamlSessionStore implements SamlSessionStore, ElytronTokeSto
             return false;
         }
 
-        if (! idMapper.hasSession(session.getID())) {
+        if (! idMapper.hasSession(session.getID()) && ! idMapperUpdater.refreshMapping(idMapper, session.getID())) {
             log.debugf("Session %s has expired on some other node", session.getID());
             session.setAttachment(SamlSession.class.getName(), null);
             return false;

adapters/saml/wildfly-elytron/src/main/java/org/keycloak/adapters/saml/elytron/infinispan/SsoCacheSessionIdMapperUpdater.java

@@ -57,6 +57,18 @@ public abstract class SsoCacheSessionIdMapperUpdater implements SessionIdMapperU
         this.delegate.map(idMapper, sso, principal, httpSessionId);
     }
 
+    @Override
+    public boolean refreshMapping(SessionIdMapper idMapper, String httpSessionId) {
+        LOG.debugf("Refreshing session %s", httpSessionId);
+
+        String[] ssoAndPrincipal = httpSessionToSsoCache.get(httpSessionId);
+        if (ssoAndPrincipal != null) {
+            this.delegate.map(idMapper, ssoAndPrincipal[0], ssoAndPrincipal[1], httpSessionId);
+            return true;
+        }
+        return false;
+    }
+
     @Override
     public void removeSession(SessionIdMapper idMapper, String httpSessionId) {
         LOG.debugf("Removing session %s", httpSessionId);

adapters/saml/wildfly-elytron/src/main/java/org/keycloak/adapters/saml/elytron/infinispan/SsoSessionCacheListener.java

@@ -26,6 +26,7 @@ import org.infinispan.client.hotrod.annotation.ClientCacheEntryRemoved;
 import org.infinispan.client.hotrod.annotation.ClientListener;
 import org.infinispan.client.hotrod.event.ClientCacheEntryCreatedEvent;
 import org.infinispan.client.hotrod.event.ClientCacheEntryRemovedEvent;
+import org.infinispan.context.Flag;
 import org.infinispan.notifications.Listener;
 import org.infinispan.notifications.cachelistener.annotation.*;
 import org.infinispan.notifications.cachelistener.event.*;
@@ -206,6 +207,7 @@ public class SsoSessionCacheListener {
             @Override
             public void run() {
                 idMapper.removeSession((String) event.getKey());
+                ssoCache.getAdvancedCache().withFlags(Flag.SKIP_CACHE_STORE).remove((String) event.getKey());
             }
         });
     }

adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoCacheSessionIdMapperUpdater.java

@@ -24,6 +24,7 @@ import io.undertow.server.HttpServerExchange;
 import io.undertow.server.session.Session;
 import io.undertow.server.session.SessionListener;
 import org.infinispan.Cache;
+import org.jboss.logging.Logger;
 
 /**
  *
@@ -31,6 +32,8 @@ import org.infinispan.Cache;
  */
 public class SsoCacheSessionIdMapperUpdater implements SessionIdMapperUpdater, SessionListener {
 
+    private static final Logger LOG = Logger.getLogger(SsoCacheSessionIdMapperUpdater.class.getName());
+
     private final SessionIdMapperUpdater delegate;
     /**
      * Cache where key is a HTTP session ID, and value is a pair (user session ID, principal name) of Strings.
@@ -52,12 +55,28 @@ public class SsoCacheSessionIdMapperUpdater implements SessionIdMapperUpdater, S
 
     @Override
     public void map(SessionIdMapper idMapper, String sso, String principal, String httpSessionId) {
+        LOG.debugf("Adding mapping (%s, %s, %s)", sso, principal, httpSessionId);
+
         httpSessionToSsoCache.put(httpSessionId, new String[] {sso, principal});
         this.delegate.map(idMapper, sso, principal, httpSessionId);
     }
 
+    @Override
+    public boolean refreshMapping(SessionIdMapper idMapper, String httpSessionId) {
+        LOG.debugf("Refreshing session %s", httpSessionId);
+
+        String[] ssoAndPrincipal = httpSessionToSsoCache.get(httpSessionId);
+        if (ssoAndPrincipal != null) {
+            this.delegate.map(idMapper, ssoAndPrincipal[0], ssoAndPrincipal[1], httpSessionId);
+            return true;
+        }
+        return false;
+    }
+
     @Override
     public void removeSession(SessionIdMapper idMapper, String httpSessionId) {
+        LOG.debugf("Removing session %s", httpSessionId);
+
         httpSessionToSsoCache.remove(httpSessionId);
         this.delegate.removeSession(idMapper, httpSessionId);
     }

adapters/saml/wildfly/wildfly-adapter/src/main/java/org/keycloak/adapters/saml/wildfly/infinispan/SsoSessionCacheListener.java

@@ -26,6 +26,7 @@ import org.infinispan.client.hotrod.annotation.ClientCacheEntryRemoved;
 import org.infinispan.client.hotrod.annotation.ClientListener;
 import org.infinispan.client.hotrod.event.ClientCacheEntryCreatedEvent;
 import org.infinispan.client.hotrod.event.ClientCacheEntryRemovedEvent;
+import org.infinispan.context.Flag;
 import org.infinispan.notifications.Listener;
 import org.infinispan.notifications.cachelistener.annotation.*;
 import org.infinispan.notifications.cachelistener.event.*;
@@ -197,5 +198,6 @@ public class SsoSessionCacheListener {
         LOG.tracev("remoteCacheEntryRemoved {0}", event.getKey());
 
         this.idMapper.removeSession((String) event.getKey());
+        ssoCache.getAdvancedCache().withFlags(Flag.SKIP_CACHE_STORE).remove((String) event.getKey());
     }
 }

adapters/spi/adapter-spi/src/main/java/org/keycloak/adapters/spi/SessionIdMapperUpdater.java

@@ -36,6 +36,10 @@ public interface SessionIdMapperUpdater {
         @Override public void removeSession(SessionIdMapper idMapper, String httpSessionId) {
             idMapper.removeSession(httpSessionId);
         }
+
+        @Override public boolean refreshMapping(SessionIdMapper idMapper, String httpSessionId) {
+            return false;
+        }
     };
 
     /**
@@ -48,12 +52,14 @@ public interface SessionIdMapperUpdater {
         @Override public void map(SessionIdMapper idMapper, String sso, String principal, String httpSessionId) { }
 
         @Override public void removeSession(SessionIdMapper idMapper, String httpSessionId) { }
+
+        @Override public boolean refreshMapping(SessionIdMapper idMapper, String httpSessionId) { return false; }
     };
 
     /**
      * Delegates to {@link SessionIdMapper#clear} method..
      */
-    public abstract void clear(SessionIdMapper idMapper);
+    void clear(SessionIdMapper idMapper);
 
     /**
      * Delegates to {@link SessionIdMapper#map} method.
@@ -62,12 +68,22 @@ public interface SessionIdMapperUpdater {
      * @param principal Principal
      * @param session HTTP session ID
      */
-    public abstract void map(SessionIdMapper idMapper, String sso, String principal, String session);
+    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);
+    void removeSession(SessionIdMapper idMapper, String session);
+
+    /**
+     * Refreshes the mapping in the {@code idMapper} from the internal source of this mapped updater
+     * and maps it via {@link SessionIdMapper#map} method.
+     * @param idMapper Mapper
+     * @param session HTTP session ID.
+     * @return {@code true} if the mapping existed in the internal source of this mapped updater
+     * and has been refreshed, {@code false} otherwise
+     */
+    boolean refreshMapping(SessionIdMapper idMapper, String httpSessionId);
 }

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

@@ -39,6 +39,7 @@ import org.keycloak.representations.idm.*;
 import org.keycloak.common.util.Retry;
 import org.keycloak.testsuite.adapter.page.EmployeeServletDistributable;
 import org.keycloak.testsuite.admin.ApiUtil;
+import org.keycloak.testsuite.arquillian.ContainerInfo;
 import org.keycloak.testsuite.util.Matchers;
 import org.keycloak.testsuite.util.SamlClient;
 import org.keycloak.testsuite.util.SamlClient.Binding;
@@ -191,4 +192,43 @@ public abstract class AbstractSAMLAdapterClusteredTest extends AbstractAdapterCl
             ;
         });
     }
+
+    @Test
+    public void testNodeRestartResiliency(@ArquillianResource
+      @OperateOnDeployment(value = EmployeeServletDistributable.DEPLOYMENT_NAME) URL employeeUrl) throws Exception {
+        ContainerInfo containerInfo = testContext.getAppServerBackendsInfo().get(0);
+
+        setPasswordFor(bburkeUser, CredentialRepresentation.PASSWORD);
+
+        String employeeUrlString = getProxiedUrl(employeeUrl);
+        SamlClient samlClient = new SamlClientBuilder()
+          // Go to employee URL at reverse proxy which is set to forward to first node
+          .navigateTo(employeeUrlString)
+
+          // process redirection to login page
+          .processSamlResponse(Binding.POST).build()
+          .login().user(bburkeUser).build()
+          .processSamlResponse(Binding.POST).build()
+
+          // Returned to the page
+          .assertResponse(Matchers.bodyHC(containsString("principal=bburke")))
+
+          .execute();
+
+        controller.stop(containerInfo.getQualifier());
+        updateProxy(NODE_2_NAME, NODE_2_URI, NODE_1_URI);   // Update the proxy to forward to the second node.
+        samlClient.execute(new SamlClientBuilder()
+          .navigateTo(employeeUrlString)
+          .doNotFollowRedirects()
+          .assertResponse(Matchers.bodyHC(containsString("principal=bburke")))
+          .getSteps());
+
+        controller.start(containerInfo.getQualifier());
+        updateProxy(NODE_1_NAME, NODE_1_URI, NODE_2_URI);   // Update the proxy to forward to the first node.
+        samlClient.execute(new SamlClientBuilder()
+          .navigateTo(employeeUrlString)
+          .doNotFollowRedirects()
+          .assertResponse(Matchers.bodyHC(containsString("principal=bburke")))
+          .getSteps());
+    }
 }