Based KC 25.x - Create a new endpoint to create sessions

freenowtech · Oct 8, 2024 · 20e2e6d · 20e2e6d
1 parent 23e227b
commit 20e2e6d
Show file tree

Hide file tree

Showing 4 changed files with 191 additions and 9 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
@@ -24,11 +24,14 @@
 import jakarta.ws.rs.Produces;
 import jakarta.ws.rs.core.MediaType;
 import jakarta.ws.rs.core.Response;
+import java.util.Base64;
+import java.util.UUID;
 import org.jboss.logging.Logger;
 import org.jboss.resteasy.reactive.NoCache;
 import org.keycloak.benchmark.dataset.config.ConfigUtil;
 import org.keycloak.benchmark.dataset.config.DatasetConfig;
 import org.keycloak.benchmark.dataset.config.DatasetException;
+import org.keycloak.common.util.Time;
 import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.events.Event;
 import org.keycloak.events.EventStoreProvider;
@@ -48,6 +51,7 @@
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.cache.CacheRealmProvider;
+import org.keycloak.models.session.UserSessionPersisterProvider;
 import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.protocol.oidc.OIDCAdvancedConfigWrapper;
 import org.keycloak.protocol.oidc.OIDCLoginProtocol;
@@ -72,6 +76,7 @@
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_EVENTS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_OFFLINE_SESSIONS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_REALMS;
+import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_SESSIONS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_USERS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.LAST_CLIENT;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.LAST_REALM;
@@ -582,6 +587,144 @@ public Response createOfflineSessions() {
         }
     }
 
+
+    @GET
+    @Path("/create-sessions")
+    @NoCache
+    @Produces(MediaType.APPLICATION_JSON)
+    public Response createSessions() {
+        boolean started = false;
+        boolean taskAdded = false;
+        try {
+            DatasetConfig config = ConfigUtil.createConfigFromQueryParams(httpRequest, CREATE_SESSIONS);
+
+            int lastRealmIndex = ConfigUtil.findFreeEntityIndex(index -> {
+                String realmName = config.getRealmPrefix() + index;
+                return baseSession.getProvider(RealmProvider.class).getRealmByName(realmName) != null;
+            }) - 1;
+            if (lastRealmIndex < 0) {
+                throw new DatasetException("Not found any realm with prefix '" + config.getRealmName() + "'");
+            }
+
+            Task task = Task.start("Creation of " + config.getCount() + " sessions");
+            TaskManager taskManager = new TaskManager(baseSession);
+            Task existingTask = taskManager.addTaskIfNotInProgress(task, config.getTaskTimeout());
+            if (existingTask != null) {
+                return Response.status(400).entity(TaskResponse.errorSomeTaskInProgress(existingTask, getStatusUrl())).build();
+            } else {
+                taskAdded = true;
+            }
+
+            logger.infof("Trigger creating sessions with the configuration: %s", config);
+            logger.infof("Will create sessions in the realms '" + config.getRealmPrefix() + "0' - '" + config.getRealmPrefix() + lastRealmIndex + "'");
+
+            // Run this in separate thread to not block HTTP request
+            new Thread(() -> createSessionsImpl(task, config, lastRealmIndex)).start();
+            started = true;
+
+            return Response.ok(TaskResponse.taskStarted(task, getStatusUrl())).build();
+        } catch (DatasetException de) {
+            return handleDatasetException(de);
+        } finally {
+            if (taskAdded && !started) {
+                new TaskManager(baseSession).removeExistingTask(false);
+            }
+        }
+    }
+
+    // Implementation of creating many sessions. This is triggered outside of HTTP request to not block HTTP request
+    private void createSessionsImpl(Task task, DatasetConfig config, int lastRealmIndex) {
+        ExecutorHelper executor = new ExecutorHelper(config.getThreadsCount(), baseSession.getKeycloakSessionFactory(), config);
+        try {
+
+            // Create events now
+            int sessionsPerTransaction = 100;
+            for (int i = 0; i < config.getCount(); i += sessionsPerTransaction) {
+                final int sessionIndex = i + sessionsPerTransaction;
+
+                // Run this concurrently with multiple threads
+                executor.addTaskRunningInTransaction(session -> {
+
+                    for (int j = 0; j < sessionsPerTransaction; j++) {
+                        int realmIdx = new Random().nextInt(lastRealmIndex + 1);
+                        String realmName = config.getRealmPrefix() + realmIdx;
+                        RealmModel realm = session.realms().getRealmByName(realmName);
+                        if (realm == null) {
+                            throw new IllegalStateException("Not found realm with name '" + realmName + "'");
+                        }
+                        // Just use user like "user-0"
+                        String username = config.getUserPrefix() + new Random().ints(0, config.getUsersPerRealm()).findFirst().getAsInt();
+                        UserModel user = session.users().getUserByUsername(realm, username);
+                        if (user == null) {
+                            throw new IllegalStateException("Not found user with username '" + username + "' in the realm '" + realmName + "'");
+                        }
+                        // Just use client like "client-0"
+                        String clientId = config.getClientPrefix() + new Random().ints(0, config.getClientsPerRealm()).findFirst().getAsInt();
+                        ClientModel client = session.clients().getClientByClientId(realm, clientId);
+                        if (client == null) {
+                            throw new IllegalStateException("Not found client with clientId  '" + client + "' in the realm '" + clientId + "'");
+                        }
+
+                        UserSessionModel userSession = session.sessions().createUserSession(realm, user, username, "127.0.0.1", "openid-connect", false, null, null);
+                        String started = userSession.getStarted() + "";
+                        userSession.setNote("KC_DEVICE_NOTE", Base64.getEncoder().encodeToString("{  \"ipAddress\": \"127.0.0.1\",  \"os\": \"Mac OS X\",  \"osVersion\": \"10.8\",  \"browser\": \"Firefox/16.0\",  \"device\": \"Mac\",  \"lastAccess\": 0,  \"mobile\": false}".getBytes()));
+                        userSession.setNote("AUTH_TIME", started);
+                        userSession.setNote("authenticators-completed", "{\"" + UUID.randomUUID() + "\": \"" +  started + "\"}");
+                        userSession.setNote("state", "LOGGED_IN");
+                        AuthenticatedClientSessionModel clientSession = session.sessions().createClientSession(userSession.getRealm(), client, userSession);
+
+                        String redirectUri = uriInfo.getBaseUri() + "realms/" + realmName + "/account";
+                        clientSession.setRedirectUri(redirectUri);
+                        clientSession.setNote("clientId", client.getId());
+                        clientSession.setNote("scope", "openid profile");
+                        clientSession.setNote("userSessionStartedAt", started);
+                        clientSession.setNote("iss", uriInfo.getBaseUri() + "realms/" + realmName);
+                        clientSession.setNote("startedAt", started);
+                        clientSession.setNote("response_type", "code");
+                        clientSession.setNote("level-of-authentication", "-1");
+                        clientSession.setNote("redirect_uri", redirectUri);
+                        clientSession.setNote("state", UUID.randomUUID().toString());
+                        clientSession.setNote("client_request_param_login", "true");
+
+                        // Persist the user sessions
+                        createOrUpdateUserSession(session, clientSession, userSession);
+
+                    }
+
+                    if (sessionIndex % (config.getThreadsCount() * sessionsPerTransaction) == 0) {
+                        task.info(logger, "Created %d sessions", sessionIndex);
+                    }
+                });
+            }
+
+            executor.waitForAllToFinish();
+
+            task.info(logger, "Created all %d sessions", config.getCount());
+            success();
+
+        } catch (Throwable ex) {
+            logException(ex);
+        } finally {
+            cleanup(executor);
+        }
+    }
+
+    public void createOrUpdateUserSession(KeycloakSession kcSession, AuthenticatedClientSessionModel clientSession, UserSessionModel userSession) {
+        UserSessionPersisterProvider persister = kcSession.getProvider(UserSessionPersisterProvider.class);
+        UserSessionModel storedUserSession = kcSession.sessions().getUserSession(clientSession.getRealm(), userSession.getId());
+        if (storedUserSession == null) {
+            persister.createUserSession(userSession, false);
+        } else {
+            storedUserSession.setLastSessionRefresh(Time.currentTime());
+        }
+
+        AuthenticatedClientSessionModel storedClientSession = storedUserSession.getAuthenticatedClientSessionByClient(clientSession.getClient().getId());
+        if (storedClientSession == null) {
+            persister.createClientSession(clientSession, false);
+        }
+
+    }
+
     // Implementation of creating many offline sessions. This is triggered outside of HTTP request to not block HTTP request
     private void createOfflineSessionsImpl(Task task, DatasetConfig config, int lastRealmIndex) {
         ExecutorHelper executor = new ExecutorHelper(config.getThreadsCount(), baseSession.getKeycloakSessionFactory(), config);

diff --git a/dataset/src/main/java/org/keycloak/benchmark/dataset/config/DatasetConfig.java b/dataset/src/main/java/org/keycloak/benchmark/dataset/config/DatasetConfig.java
@@ -23,6 +23,7 @@
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_EVENTS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_OFFLINE_SESSIONS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_REALMS;
+import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_SESSIONS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_USERS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.LAST_CLIENT;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.LAST_REALM;
@@ -39,7 +40,7 @@ public class DatasetConfig {
 
     // Used when creating many realms as a prefix. For example when prefix us "foo", we will create realms like "foo0", "foo1" etc.
     // For many events, it will need the realm prefix as events are created randomly in all the already created realms
-    @QueryParamFill(paramName = "realm-prefix", defaultValue = "realm-", operations = { CREATE_REALMS, CREATE_EVENTS, CREATE_OFFLINE_SESSIONS,
+    @QueryParamFill(paramName = "realm-prefix", defaultValue = "realm-", operations = { CREATE_REALMS, CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, CREATE_SESSIONS,
             REMOVE_REALMS, LAST_REALM })
     private String realmPrefix;
 
@@ -63,7 +64,7 @@ public class DatasetConfig {
     private Integer start;
 
     // Count of entities to be created. Entity is realm, client or user based on the operation
-    @QueryParamIntFill(paramName = "count", required = true, operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_USERS, CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, CREATE_AUTHZ_CLIENT })
+    @QueryParamIntFill(paramName = "count", required = true, operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_USERS, CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, CREATE_SESSIONS, CREATE_AUTHZ_CLIENT })
     private Integer count;
 
     // Prefix for realm roles to create in every realm (in case of CREATE_REALMS) or to assign to users (in case of CREATE_USERS)
@@ -75,10 +76,10 @@ public class DatasetConfig {
     private Integer realmRolesPerRealm;
 
     // Prefix for newly created clients (in case of CREATE_REALMS and CREATE_CLIENTS). In case of CREATE_USERS it is used to find the clients with clientRoles, which will be assigned to users
-    @QueryParamFill(paramName = "client-prefix", defaultValue = "client-", operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_USERS, CREATE_OFFLINE_SESSIONS, LAST_CLIENT })
+    @QueryParamFill(paramName = "client-prefix", defaultValue = "client-", operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_USERS, CREATE_OFFLINE_SESSIONS, CREATE_SESSIONS, LAST_CLIENT })
     private String clientPrefix;
 
-    @QueryParamIntFill(paramName = "clients-per-realm", defaultValue = 30, operations = { CREATE_REALMS })
+    @QueryParamIntFill(paramName = "clients-per-realm", defaultValue = 30, operations = { CREATE_REALMS, CREATE_SESSIONS })
     private Integer clientsPerRealm;
 
     // Count of clients created in every DB transaction
@@ -131,11 +132,11 @@ public class DatasetConfig {
 
 
     // Prefix for newly created users
-    @QueryParamFill(paramName = "user-prefix", defaultValue = "user-", operations = { CREATE_REALMS, CREATE_USERS, CREATE_OFFLINE_SESSIONS, LAST_USER, CREATE_AUTHZ_CLIENT })
+    @QueryParamFill(paramName = "user-prefix", defaultValue = "user-", operations = { CREATE_REALMS, CREATE_USERS, CREATE_OFFLINE_SESSIONS, CREATE_SESSIONS, LAST_USER, CREATE_AUTHZ_CLIENT })
     private String userPrefix;
 
     // Count of users to be created in every realm (In case of CREATE_REALMS)
-    @QueryParamIntFill(paramName = "users-per-realm", defaultValue = 200, operations = { CREATE_REALMS, CREATE_AUTHZ_CLIENT })
+    @QueryParamIntFill(paramName = "users-per-realm", defaultValue = 200, operations = { CREATE_REALMS, CREATE_AUTHZ_CLIENT, CREATE_SESSIONS })
     private Integer usersPerRealm;
 
     // Count of groups assigned to every user
@@ -164,7 +165,7 @@ public class DatasetConfig {
 
     // Transaction timeout used for transactions for creating objects
     @QueryParamIntFill(paramName = "transaction-timeout", defaultValue = 300, operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_USERS,
-            CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, REMOVE_REALMS, CREATE_AUTHZ_CLIENT })
+            CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, CREATE_SESSIONS, REMOVE_REALMS, CREATE_AUTHZ_CLIENT })
     private Integer transactionTimeoutInSeconds;
 
     // Count of users created in every transaction
@@ -173,13 +174,13 @@ public class DatasetConfig {
 
     // Count of worker threads concurrently creating entities
     @QueryParamIntFill(paramName = "threads-count", operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_USERS,
-            CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, REMOVE_REALMS, CREATE_AUTHZ_CLIENT })
+            CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, CREATE_SESSIONS, REMOVE_REALMS, CREATE_AUTHZ_CLIENT })
     private Integer threadsCount;
 
     // Timeout for the whole task. If timeout expires, then the existing task may not be terminated immediatelly. However it will be permitted to start another task
     // (EG. Send another HTTP request for creating realms), which can cause conflicts
     @QueryParamIntFill(paramName = "task-timeout", defaultValue = 3600, operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_USERS,
-            CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, REMOVE_REALMS, CREATE_AUTHZ_CLIENT })
+            CREATE_EVENTS, CREATE_OFFLINE_SESSIONS, CREATE_SESSIONS, REMOVE_REALMS, CREATE_AUTHZ_CLIENT })
     private Integer taskTimeout;
 
     // The client id of a client to which data is going to be provisioned

diff --git a/dataset/src/main/java/org/keycloak/benchmark/dataset/config/DatasetOperation.java b/dataset/src/main/java/org/keycloak/benchmark/dataset/config/DatasetOperation.java
@@ -28,6 +28,7 @@ public enum DatasetOperation {
     CREATE_USERS,
     CREATE_EVENTS,
     CREATE_OFFLINE_SESSIONS,
+    CREATE_SESSIONS,
     REMOVE_REALMS,
     LAST_REALM,
     LAST_CLIENT,

diff --git a/provision/docker-compose/keycloak/keycloak-25-quarkus-postgres.yml b/provision/docker-compose/keycloak/keycloak-25-quarkus-postgres.yml
@@ -0,0 +1,37 @@
+version: '2.2'
+
+volumes:
+  postgres_data:
+      driver: local
+
+services:
+  postgres:
+      image: postgres
+      container_name: postgres
+      volumes:
+        - postgres_data:/var/lib/postgresql/data
+      environment:
+        POSTGRES_DB: keycloak
+        POSTGRES_USER: keycloak
+        POSTGRES_PASSWORD: password
+      ports:
+        - "5432:5432"
+  keycloak:
+      image: quay.io/keycloak/keycloak:25.0.2
+      container_name: keycloak
+      command:
+        - start-dev
+      volumes:
+        - ./providers:/opt/keycloak/providers:z
+      environment:
+        KC_DB_URL: jdbc:postgresql://postgres:5432/keycloak
+        KC_DB_USERNAME: keycloak
+        KC_DB_PASSWORD: password
+        KC_DB: postgres
+        KC_HOSTNAME: localhost
+        KEYCLOAK_ADMIN: admin
+        KEYCLOAK_ADMIN_PASSWORD: admin
+      ports:
+        - 8080:8080
+      depends_on:
+        - postgres