LB check with failover and session expiration test (#647)

* LB check with failover and session expiration test Signed-off-by: Martin Kanis <mkanis@redhat.com> * Make redirectUri configurable for exchangeCode Without this the code cannot be exchanged using DC_1 and DC_2 because Keycloak expects redirectUri to be equal LOAD_BALANCER's URL Signed-off-by: Michal Hajas <mhajas@redhat.com> * Do not print site down logs with each lb-check execution Signed-off-by: Michal Hajas <mhajas@redhat.com> * Move wait functionality to KeycloakClient Signed-off-by: Michal Hajas <mhajas@redhat.com> * Report useful output when expected input is not set Signed-off-by: Michal Hajas <mhajas@redhat.com> * Make log more clear Signed-off-by: Michal Hajas <mhajas@redhat.com> * Remove unused import Signed-off-by: Michal Hajas <mhajas@redhat.com> --------- Signed-off-by: Martin Kanis <mkanis@redhat.com> Signed-off-by: Michal Hajas <mhajas@redhat.com> Co-authored-by: Michal Hajas <mhajas@redhat.com>
keycloak · Dec 21, 2023 · e6394ec · e6394ec
1 parent cee78bc
commit e6394ec
Show file tree

Hide file tree

Showing 12 changed files with 312 additions and 8 deletions.
diff --git a/dataset/src/main/java/org/keycloak/benchmark/dataset/DatasetResourceProvider.java b/dataset/src/main/java/org/keycloak/benchmark/dataset/DatasetResourceProvider.java
@@ -29,6 +29,7 @@
 import org.keycloak.benchmark.dataset.config.ConfigUtil;
 import org.keycloak.benchmark.dataset.config.DatasetConfig;
 import org.keycloak.benchmark.dataset.config.DatasetException;
+import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.events.Event;
 import org.keycloak.events.EventStoreProvider;
 import org.keycloak.events.EventType;
@@ -870,6 +871,28 @@ public AuthorizationProvisioner authz() {
         return new AuthorizationProvisioner(baseSession);
     }
 
+    @GET
+    @Path("/take-dc-down")
+    @NoCache
+    @Produces(MediaType.APPLICATION_JSON)
+    public Response takeDCDown() {
+        String siteName = baseSession.getProvider(InfinispanConnectionProvider.class).getTopologyInfo().getMySiteName();
+        baseSession.realms().getRealmByName("master").setAttribute("is-site-" + siteName + "-down", true);
+
+        return Response.ok(TaskResponse.statusMessage("Site " + siteName + " was marked as down.")).build();
+    }
+
+    @GET
+    @Path("/take-dc-up")
+    @NoCache
+    @Produces(MediaType.APPLICATION_JSON)
+    public Response takeDCUp() {
+        String siteName = baseSession.getProvider(InfinispanConnectionProvider.class).getTopologyInfo().getMySiteName();
+        baseSession.realms().getRealmByName("master").removeAttribute("is-site-" + siteName + "-down");
+
+        return Response.ok(TaskResponse.statusMessage("Site " + siteName + " was marked as up.")).build();
+    }
+
     @Override
     public void close() {
     }

diff --git a/dataset/src/main/java/org/keycloak/benchmark/lb/TestMultiSiteLoadBalancerCheckProvider.java b/dataset/src/main/java/org/keycloak/benchmark/lb/TestMultiSiteLoadBalancerCheckProvider.java
@@ -0,0 +1,31 @@
+package org.keycloak.benchmark.lb;
+
+import org.jboss.logging.Logger;
+import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
+import org.keycloak.health.LoadBalancerCheckProvider;
+import org.keycloak.models.KeycloakSession;
+
+public class TestMultiSiteLoadBalancerCheckProvider implements LoadBalancerCheckProvider {
+
+    protected static final Logger logger = Logger.getLogger(TestMultiSiteLoadBalancerCheckProvider.class);
+
+    private final KeycloakSession session;
+
+    public TestMultiSiteLoadBalancerCheckProvider(KeycloakSession session) {
+        this.session = session;
+    }
+
+    @Override
+    public boolean isDown() {
+        String siteName = session.getProvider(InfinispanConnectionProvider.class).getTopologyInfo().getMySiteName();
+        boolean isDown = session.realms().getRealmByName("master").getAttribute("is-site-" + siteName + "-down", false);
+
+        logger.debugf("Site %s is %s", siteName, isDown ? "DOWN" : "UP");
+        return isDown;
+    }
+
+    @Override
+    public void close() {
+
+    }
+}
diff --git a/...rc/main/java/org/keycloak/benchmark/lb/TestMultiSiteLoadBalancerCheckProviderFactory.java b/...rc/main/java/org/keycloak/benchmark/lb/TestMultiSiteLoadBalancerCheckProviderFactory.java
@@ -0,0 +1,35 @@
+package org.keycloak.benchmark.lb;
+
+import org.keycloak.Config;
+import org.keycloak.health.LoadBalancerCheckProvider;
+import org.keycloak.health.LoadBalancerCheckProviderFactory;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+
+public class TestMultiSiteLoadBalancerCheckProviderFactory implements LoadBalancerCheckProviderFactory {
+
+    @Override
+    public LoadBalancerCheckProvider create(KeycloakSession keycloakSession) {
+        return new TestMultiSiteLoadBalancerCheckProvider(keycloakSession);
+    }
+
+    @Override
+    public void init(Config.Scope scope) {
+
+    }
+
+    @Override
+    public void postInit(KeycloakSessionFactory keycloakSessionFactory) {
+
+    }
+
+    @Override
+    public void close() {
+
+    }
+
+    @Override
+    public String getId() {
+        return "test-multisite";
+    }
+}
diff --git a/...src/main/resources/META-INF/services/org.keycloak.health.LoadBalancerCheckProviderFactory b/...src/main/resources/META-INF/services/org.keycloak.health.LoadBalancerCheckProviderFactory
@@ -0,0 +1 @@
+org.keycloak.benchmark.lb.TestMultiSiteLoadBalancerCheckProviderFactory
diff --git a/...hmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/AbstractCrossDCTest.java b/...hmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/AbstractCrossDCTest.java
@@ -15,10 +15,15 @@
 import org.keycloak.representations.idm.RealmRepresentation;
 import org.keycloak.representations.idm.UserRepresentation;
 
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.net.UnknownHostException;
 import java.net.http.HttpClient;
 import java.util.Collections;
 import java.util.List;
 
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.keycloak.benchmark.crossdc.util.HttpClientUtils.MOCK_COOKIE_MANAGER;
 import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.DISTRIBUTED_CACHES;
 
@@ -35,18 +40,24 @@ public abstract class AbstractCrossDCTest {
     public static final String MAIN_PASSWORD = System.getProperty("main.password");
 
     static {
+        assertNotNull(MAIN_PASSWORD, "Main password must be set");
         DC_1 = new DatacenterInfo(HTTP_CLIENT, System.getProperty("keycloak.dc1.url"), System.getProperty("infinispan.dc1.url"));
         DC_2 = new DatacenterInfo(HTTP_CLIENT, System.getProperty("keycloak.dc2.url"), System.getProperty("infinispan.dc2.url"));
         LOAD_BALANCER_KEYCLOAK = new KeycloakClient(HTTP_CLIENT, System.getProperty("load-balancer.url"));
     }
 
     @BeforeEach
-    public void setUpTestEnvironment() {
+    public void setUpTestEnvironment() throws UnknownHostException {
+        assertTrue(DC_1.kc().isActive(LOAD_BALANCER_KEYCLOAK));
+
         Keycloak adminClient = DC_1.kc().adminClient();
         LOG.info("Setting up test environment");
         LOG.info("-------------------------------------------");
         LOG.info("Status of caches before test:");
-        DISTRIBUTED_CACHES.forEach(cache -> {
+        DISTRIBUTED_CACHES
+                .stream()
+                .filter(cache -> !cache.equals(InfinispanUtils.WORK))
+                .forEach(cache -> {
             LOG.infof("External cache %s " + cache + " in DC1: %d - entries [%s]", cache, DC_1.ispn().cache(cache).size(), DC_1.ispn().cache(cache).keys());
             LOG.infof("External cache %s " + cache + " in DC2: %d - entries [%s]", cache, DC_2.ispn().cache(cache).size(), DC_2.ispn().cache(cache).keys());
         });
@@ -92,7 +103,7 @@ public void setUpTestEnvironment() {
     }
 
     @AfterEach
-    public void tearDownTestEnvironment() {
+    public void tearDownTestEnvironment() throws URISyntaxException, IOException, InterruptedException {
         Keycloak adminClient = DC_1.kc().adminClient();
 
         if (adminClient.realms().realm(REALM_NAME).toRepresentation() != null) {
@@ -113,5 +124,9 @@ public void tearDownTestEnvironment() {
                 });
 
         MOCK_COOKIE_MANAGER.getCookieStore().removeAll();
+
+        DC_1.kc().markLBCheckUp();
+        DC_2.kc().markLBCheckUp();
+        DC_1.kc().waitToBeActive(LOAD_BALANCER_KEYCLOAK);
     }
 }
diff --git a/...ak-benchmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/FailoverTest.java b/...ak-benchmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/FailoverTest.java
@@ -0,0 +1,34 @@
+package org.keycloak.benchmark.crossdc;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.SESSIONS;
+
+public class FailoverTest extends AbstractCrossDCTest {
+
+    @Test
+    public void logoutUserWithFailoverTest() throws IOException, URISyntaxException, InterruptedException {
+        // Login and exchange code in DC1
+        String code = LOAD_BALANCER_KEYCLOAK.usernamePasswordLogin( REALM_NAME, USERNAME, MAIN_PASSWORD, CLIENTID);
+        Map<String, Object> tokensMap = LOAD_BALANCER_KEYCLOAK.exchangeCode(REALM_NAME, CLIENTID, CLIENT_SECRET, 200, code);
+
+        DC_1.kc().markLBCheckDown();
+        DC_2.kc().waitToBeActive(LOAD_BALANCER_KEYCLOAK);
+
+        // Verify if the user session UUID in code, we fetched from Keycloak exists in session cache key of external ISPN in DC2
+        Set<String> sessions = DC_2.ispn().cache(SESSIONS).keys();
+        assertTrue(sessions.contains(code.split("[.]")[1]));
+
+        tokensMap = LOAD_BALANCER_KEYCLOAK.refreshToken(REALM_NAME, (String) tokensMap.get("refresh_token"), CLIENTID, CLIENT_SECRET, 200);
+
+        LOAD_BALANCER_KEYCLOAK.logout(REALM_NAME, (String) tokensMap.get("id_token"), CLIENTID);
+
+        LOAD_BALANCER_KEYCLOAK.refreshToken(REALM_NAME, (String) tokensMap.get("refresh_token"), CLIENTID, CLIENT_SECRET, 400);
+    }
+}
diff --git a/...benchmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/LoginLogoutTest.java b/...benchmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/LoginLogoutTest.java
@@ -19,8 +19,8 @@ public void loginLogoutTest() throws URISyntaxException, IOException, Interrupte
         Map<String, Object> tokensMap = LOAD_BALANCER_KEYCLOAK.exchangeCode(REALM_NAME, CLIENTID, CLIENT_SECRET, 200, code);
 
         //Making sure the code cannot be reused in any of the DCs
-        DC_2.kc().exchangeCode(REALM_NAME, CLIENTID, CLIENT_SECRET, 400, code);
-        DC_1.kc().exchangeCode(REALM_NAME, CLIENTID, CLIENT_SECRET, 400, code);
+        DC_2.kc().exchangeCode(REALM_NAME, CLIENTID, CLIENT_SECRET, 400, code, LOAD_BALANCER_KEYCLOAK.getRedirectUri(REALM_NAME));
+        DC_1.kc().exchangeCode(REALM_NAME, CLIENTID, CLIENT_SECRET, 400, code, LOAD_BALANCER_KEYCLOAK.getRedirectUri(REALM_NAME));
 
         //Verify if the user session UUID in code, we fetched from Keycloak exists in session cache key of external ISPN in DC1
         String sessionId = code.split("[.]")[1];

diff --git a/...ark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/SessionExpirationTest.java b/...ark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/SessionExpirationTest.java
@@ -0,0 +1,59 @@
+package org.keycloak.benchmark.crossdc;
+
+import org.junit.jupiter.api.Test;
+import org.keycloak.representations.idm.RealmRepresentation;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.CLIENT_SESSIONS;
+import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.SESSIONS;
+
+public class SessionExpirationTest extends AbstractCrossDCTest {
+
+    @Test
+    public void sessionExpirationTest() throws IOException, URISyntaxException, InterruptedException {
+        // set user/client session lifespan to 5s
+        RealmRepresentation realm = LOAD_BALANCER_KEYCLOAK.adminClient().realm(REALM_NAME).toRepresentation();
+        realm.setSsoSessionMaxLifespan(5);
+        realm.setClientSessionMaxLifespan(5);
+        LOAD_BALANCER_KEYCLOAK.adminClient().realm(REALM_NAME).update(realm);
+
+        // create a user and client session
+        String code = LOAD_BALANCER_KEYCLOAK.usernamePasswordLogin(REALM_NAME, USERNAME, MAIN_PASSWORD, CLIENTID);
+        Map<String, Object> tokensMap = LOAD_BALANCER_KEYCLOAK.exchangeCode(REALM_NAME, CLIENTID, CLIENT_SECRET, 200, code);
+
+        // check the sessions are replicated in remote caches
+        assertEquals(1, DC_1.ispn().cache(SESSIONS).size());
+        assertEquals(1, DC_2.ispn().cache(SESSIONS).size());
+        assertEquals(1, DC_1.ispn().cache(CLIENT_SESSIONS).size());
+        assertEquals(1, DC_2.ispn().cache(CLIENT_SESSIONS).size());
+
+        // check the sessions are replicated in embedded caches
+        assertEquals(1, DC_1.kc().embeddedIspn().cache(SESSIONS).size());
+        assertEquals(1, DC_2.kc().embeddedIspn().cache(SESSIONS).size());
+        assertEquals(1, DC_1.kc().embeddedIspn().cache(CLIENT_SESSIONS).size());
+        assertEquals(1, DC_2.kc().embeddedIspn().cache(CLIENT_SESSIONS).size());
+
+        // let them expire
+        Thread.sleep(6000);
+
+        // check the remote caches are empty
+        assertEquals(0, DC_1.ispn().cache(SESSIONS).size());
+        assertEquals(0, DC_2.ispn().cache(SESSIONS).size());
+        assertEquals(0, DC_1.ispn().cache(CLIENT_SESSIONS).size());
+        assertEquals(0, DC_2.ispn().cache(CLIENT_SESSIONS).size());
+
+        // check the embedded caches are empty
+        assertEquals(0, DC_1.kc().embeddedIspn().cache(SESSIONS).size());
+        assertEquals(0, DC_2.kc().embeddedIspn().cache(SESSIONS).size());
+        assertEquals(0, DC_1.kc().embeddedIspn().cache(CLIENT_SESSIONS).size());
+        assertEquals(0, DC_2.kc().embeddedIspn().cache(CLIENT_SESSIONS).size());
+
+        // token refresh should fail
+        DC_2.kc().refreshToken(REALM_NAME, (String) tokensMap.get("refresh_token"), CLIENTID, CLIENT_SECRET, 400);
+        DC_1.kc().refreshToken(REALM_NAME, (String) tokensMap.get("refresh_token"), CLIENTID, CLIENT_SECRET, 400);
+    }
+}
diff --git a/...c-tests/src/test/java/org/keycloak/benchmark/crossdc/client/ExternalInfinispanClient.java b/...c-tests/src/test/java/org/keycloak/benchmark/crossdc/client/ExternalInfinispanClient.java
@@ -16,6 +16,7 @@
 import java.util.stream.Collectors;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.keycloak.benchmark.crossdc.AbstractCrossDCTest.ISPN_USERNAME;
 import static org.keycloak.benchmark.crossdc.AbstractCrossDCTest.MAIN_PASSWORD;
 import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.getBasicAuthenticationHeader;
@@ -29,6 +30,7 @@ public class ExternalInfinispanClient implements InfinispanClient {
     Pattern UUID_REGEX = Pattern.compile("[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}");
 
     public ExternalInfinispanClient(HttpClient httpClient, String infinispanUrl, String username, String password) {
+        assertNotNull(infinispanUrl, "Infinispan URL cannot be null");
         this.httpClient = httpClient;
         this.infinispanUrl = infinispanUrl;
         this.username = username;
@@ -61,7 +63,7 @@ public long size() {
                 HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
                 assertEquals(200, response.statusCode());
 
-                if (cacheName.equals(InfinispanUtils.SESSIONS)) {
+                if (cacheName.equals(InfinispanUtils.SESSIONS) || cacheName.equals(InfinispanUtils.CLIENT_SESSIONS)) {
                     return Long.parseLong(response.body()) - KeycloakClient.getCurrentlyInitializedAdminClients();
                 }
                 return Long.parseLong(response.body());