Dataset for organizations

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.benchmark.dataset.organization.OrganizationProvisioner;
 import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.events.Event;
 import org.keycloak.events.EventStoreProvider;
@@ -871,6 +872,11 @@ public AuthorizationProvisioner authz() {
         return new AuthorizationProvisioner(baseSession);
     }
 
+    @Path("/orgs")
+    public OrganizationProvisioner orgs() {
+        return new OrganizationProvisioner(baseSession);
+    }
+
     @GET
     @Path("/take-dc-down")
     @NoCache
@@ -1103,7 +1109,8 @@ private void createUsers(RealmContext context, KeycloakSession session, int star
             int groupIndexStartForCurrentUser = (i * config.getGroupsPerUser());
             for (int j = groupIndexStartForCurrentUser ; j < groupIndexStartForCurrentUser + config.getGroupsPerUser() ; j++) {
                 int groupIndex = j % context.getGroups().size();
-                user.joinGroup(context.getGroups().get(groupIndex));
+                GroupModel group = context.getGroups().get(groupIndex);
+                user.joinGroup(session.groups().getGroupById(realm, group.getId()));
 
                 logger.tracef("Assigned group %s to the user %s", context.getGroups().get(groupIndex).getName(), user.getUsername());
             }

dataset/src/main/java/org/keycloak/benchmark/dataset/config/DatasetConfig.java

@@ -22,6 +22,7 @@
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_CLIENTS;
 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_ORGS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_REALMS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_USERS;
 import static org.keycloak.benchmark.dataset.config.DatasetOperation.LAST_CLIENT;
@@ -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_AUTHZ_CLIENT, CREATE_ORGS })
     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)
@@ -86,7 +87,7 @@ public class DatasetConfig {
     private Integer clientsPerTransaction;
 
     // default number of entries created per DB transaction
-    @QueryParamIntFill(paramName = "entries-per-transaction", defaultValue = 10, operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_AUTHZ_CLIENT })
+    @QueryParamIntFill(paramName = "entries-per-transaction", defaultValue = 10, operations = { CREATE_REALMS, CREATE_CLIENTS, CREATE_AUTHZ_CLIENT, CREATE_ORGS })
     private Integer entriesPerTransaction;
 
     // Prefix of clientRoles to be created (in case of CREATE_REALMS and CREATE_CLIENTS). In case of CREATE_USERS it is used to find the clientRoles, which will be assigned to users
@@ -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, REMOVE_REALMS, CREATE_AUTHZ_CLIENT, CREATE_ORGS })
     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, REMOVE_REALMS, CREATE_AUTHZ_CLIENT, CREATE_ORGS })
     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, REMOVE_REALMS, CREATE_AUTHZ_CLIENT, CREATE_ORGS })
     private Integer taskTimeout;
 
     // The client id of a client to which data is going to be provisioned
@@ -202,6 +203,24 @@ public class DatasetConfig {
     @QueryParamIntFill(paramName = "users-per-user-policy", defaultValue = 1, operations = CREATE_AUTHZ_CLIENT)
     private int usersPerUserPolicy;
 
+    @QueryParamFill(paramName = "name", operations = { CREATE_ORGS })
+    private String name;
+
+    @QueryParamFill(paramName = "domains", operations = { CREATE_ORGS })
+    private String domains;
+
+    @QueryParamIntFill(paramName = "domains-count", operations = { CREATE_ORGS })
+    private int domainsCount;
+
+    @QueryParamFill(paramName = "org-prefix", defaultValue = "org-", operations = { CREATE_ORGS })
+    private String orgPrefix;
+
+    @QueryParamIntFill(paramName = "unmanaged-members-count", operations = CREATE_ORGS)
+    private int unManagedMembersCount;
+
+    @QueryParamIntFill(paramName = "identity-providers-count", operations = CREATE_ORGS)
+    private int identityProvidersCount;
+
     // String representation of this configuration (cached here to not be computed in runtime)
     private String toString = "DatasetConfig []";
 
@@ -380,4 +399,28 @@ public int getUsersPerUserPolicy() {
     public String toString() {
         return toString;
     }
+
+    public String getName() {
+        return name;
+    }
+
+    public String getDomains() {
+        return domains;
+    }
+
+    public int getDomainsCount() {
+        return domainsCount;
+    }
+
+    public String getOrgPrefix() {
+        return orgPrefix;
+    }
+
+    public int getUnManagedMembersCount() {
+        return unManagedMembersCount;
+    }
+
+    public int getIdentityProvidersCount() {
+        return identityProvidersCount;
+    }
 }

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_ORGS,
     REMOVE_REALMS,
     LAST_REALM,
     LAST_CLIENT,

dataset/src/main/java/org/keycloak/benchmark/dataset/organization/AbstractOrganizationProvisioner.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2024 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.benchmark.dataset.organization;
+
+import static org.keycloak.benchmark.dataset.config.DatasetOperation.CREATE_ORGS;
+
+import jakarta.ws.rs.NotFoundException;
+import jakarta.ws.rs.core.Response;
+import org.keycloak.benchmark.dataset.DatasetResourceProvider;
+import org.keycloak.benchmark.dataset.ExecutorHelper;
+import org.keycloak.benchmark.dataset.Task;
+import org.keycloak.benchmark.dataset.TaskManager;
+import org.keycloak.benchmark.dataset.TaskResponse;
+import org.keycloak.benchmark.dataset.config.ConfigUtil;
+import org.keycloak.benchmark.dataset.config.DatasetConfig;
+import org.keycloak.benchmark.dataset.config.DatasetException;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionTask;
+import org.keycloak.models.RealmModel;
+import org.keycloak.models.utils.KeycloakModelUtils;
+import org.keycloak.organization.OrganizationProvider;
+
+public class AbstractOrganizationProvisioner extends DatasetResourceProvider {
+
+    private final String realmName;
+
+    public AbstractOrganizationProvisioner(KeycloakSession session) {
+        super(session);
+        realmName = session.getContext().getRealm().getName();
+    }
+
+    protected Response start(String name, Runnable runnable) {
+        DatasetConfig config = getDatasetConfig();
+        Task task = Task.start(name);
+        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();
+        }
+
+        try {
+            new Thread(runnable).start();
+            return Response.ok(TaskResponse.taskStarted(task, getStatusUrl())).build();
+        } catch (DatasetException de) {
+            return handleException(handleDatasetException(de));
+        } catch (Exception e) {
+            logger.error("Unexpected error", e);
+            return handleException(Response.status(500).entity(TaskResponse.error("Unexpected error")).build());
+        }
+    }
+
+    protected OrganizationProvider getOrganizationProvider(KeycloakSession session) {
+        session.getContext().setRealm(getRealm(session));
+        OrganizationProvider orgProvider = session.getProvider(OrganizationProvider.class);
+
+        if (orgProvider == null) {
+            throw new NotFoundException();
+        }
+
+        return orgProvider;
+    }
+
+    protected DatasetConfig getDatasetConfig() {
+        return ConfigUtil.createConfigFromQueryParams(httpRequest, CREATE_ORGS);
+    }
+
+    private Response handleException(Response response) {
+        new TaskManager(baseSession).removeExistingTask(false);
+        return response;
+    }
+
+    protected String getRealmName() {
+        return realmName;
+    }
+
+    protected RealmModel getRealm(KeycloakSession session) {
+        RealmModel realm = session.realms().getRealmByName(realmName);
+        session.getContext().setRealm(realm);
+        return realm;
+    }
+
+    protected void runJobInTransactionWithTimeout(KeycloakSessionTask task) {
+        KeycloakModelUtils.runJobInTransactionWithTimeout(baseSession.getKeycloakSessionFactory(), task::run, getDatasetConfig().getTransactionTimeoutInSeconds());
+    }
+
+    protected void addTaskRunningInTransaction(ExecutorHelper executor, KeycloakSessionTask task) {
+        executor.addTaskRunningInTransaction(task::run);
+    }
+}

dataset/src/main/java/org/keycloak/benchmark/dataset/organization/OrganizationIdentityProviderProvisioner.java

@@ -0,0 +1,144 @@
+/*
+ * Copyright 2024 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.benchmark.dataset.organization;
+
+import java.util.HashMap;
+import java.util.concurrent.CountDownLatch;
+
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.Produces;
+import jakarta.ws.rs.core.MediaType;
+import jakarta.ws.rs.core.Response;
+import org.jboss.resteasy.reactive.NoCache;
+import org.keycloak.benchmark.dataset.ExecutorHelper;
+import org.keycloak.benchmark.dataset.TaskManager;
+import org.keycloak.benchmark.dataset.config.ConfigUtil;
+import org.keycloak.benchmark.dataset.config.DatasetConfig;
+import org.keycloak.broker.oidc.KeycloakOIDCIdentityProviderFactory;
+import org.keycloak.models.ClientModel;
+import org.keycloak.models.IdentityProviderModel;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.OrganizationModel;
+import org.keycloak.models.RealmModel;
+import org.keycloak.models.utils.KeycloakModelUtils;
+import org.keycloak.organization.OrganizationProvider;
+import org.keycloak.protocol.oidc.OIDCLoginProtocol;
+
+public class OrganizationIdentityProviderProvisioner extends AbstractOrganizationProvisioner {
+
+    private final OrganizationModel organization;
+    private final DatasetConfig config;
+
+    public OrganizationIdentityProviderProvisioner(KeycloakSession session, OrganizationModel organization, DatasetConfig config) {
+        super(session);
+        this.organization = organization;
+        this.config = config;
+    }
+
+    @Path("create")
+    @GET
+    @NoCache
+    @Produces(MediaType.APPLICATION_JSON)
+    public Response create() {
+        return start("Creation of " + config.getCount() + " identity providers in organization " + organization.getName(), addIdentityProviders());
+    }
+
+    private Runnable addIdentityProviders() {
+        return () -> {
+            String orgId = organization.getId();
+            CountDownLatch latch = new CountDownLatch(config.getCount());
+            KeycloakSessionFactory sessionFactory = baseSession.getKeycloakSessionFactory();
+            ExecutorHelper executor = new ExecutorHelper(config.getThreadsCount(), sessionFactory, config);
+
+            executor.addTask(() -> runJobInTransactionWithTimeout(session -> addIdentityProviders(session, orgId, latch)));
+
+            try {
+                executor.waitForAllToFinish();
+                success();
+                logger.infof("Added %d identity providers to organization %s", config.getCount() - latch.getCount(), organization.getName());
+            } catch (Exception e) {
+                KeycloakModelUtils.runJobInTransaction(sessionFactory, session
+                        -> new TaskManager(session).removeExistingTask(false));
+            }
+        };
+    }
+
+    protected void addIdentityProviders(KeycloakSession session, String orgId, CountDownLatch latch) {
+        OrganizationProvider provider = getOrganizationProvider(session);
+        OrganizationModel organization = provider.getById(orgId);
+        int startIndex = getLastIndex(session, organization.getName());
+        RealmModel realm = session.getContext().getRealm();
+        ClientModel client = realm.getClientByClientId("org-broker-client");
+
+        if (client == null) {
+            client = realm.addClient("org-broker-client");
+            client.setSecret("secret");
+            client.setProtocol(OIDCLoginProtocol.LOGIN_PROTOCOL);
+            client.setPublicClient(false);
+        }
+
+        long count = latch.getCount();
+
+        for (int i = startIndex; i < count; i++) {
+            String idpAlias = "idp-" + organization.getName() + "-" + i;
+
+            IdentityProviderModel identityProvider = realm.getIdentityProviderByAlias(idpAlias);
+
+            if (identityProvider != null && identityProvider.getInternalId() != null && identityProvider.getOrganizationId() != null) {
+                continue;
+            }
+
+            if (identityProvider == null || identityProvider.getInternalId() == null) {
+                identityProvider = new IdentityProviderModel();
+
+                identityProvider.setAlias(idpAlias);
+                identityProvider.setProviderId(KeycloakOIDCIdentityProviderFactory.PROVIDER_ID);
+                identityProvider.setLoginHint(true);
+                identityProvider.setEnabled(true);
+                HashMap<String, String> idpConfig = new HashMap<>();
+                identityProvider.setConfig(idpConfig);
+                idpConfig.put("issuer", "http://localhost:8180/realms/" + config.getRealmName());
+                idpConfig.put("jwksUrl", "http://localhost:8180/realms/realm-0/protocol/openid-connect/certs");
+                idpConfig.put("logoutUrl", "http://localhost:8180/realms/realm-0/protocol/openid-connect/auth");
+                idpConfig.put("metadataDescriptorUrl", "http://localhost:8180/realms/realm-0/.well-known/openid-configuration");
+                idpConfig.put("tokenUrl", "http://localhost:8180/realms/realm-0/protocol/openid-connect/token");
+                idpConfig.put("authorizationUrl", "http://localhost:8180/realms/realm-0/protocol/openid-connect/auth");
+                idpConfig.put("useJwksUrl", "true");
+                idpConfig.put("userInfoUrl", "http://localhost:8180/realms/realm-0/protocol/openid-connect/userinfo");
+                idpConfig.put("validateSignature", "true");
+                idpConfig.put("clientId", "org-broker-client");
+                idpConfig.put("clientSecret", "secret");
+                idpConfig.put("clientAuthMethod", "client_secret_post");
+                realm.addIdentityProvider(identityProvider);
+                client.addRedirectUri("http://localhost:8180/realms/" + config.getRealmName() + "/broker/" + identityProvider.getAlias() + "/endpoint");
+            }
+
+            if (provider.addIdentityProvider(organization, identityProvider)) {
+                latch.countDown();
+            }
+        }
+    }
+
+    private int getLastIndex(KeycloakSession session, String orgName) {
+        RealmModel realm = session.getContext().getRealm();
+        return ConfigUtil.findFreeEntityIndex(index -> realm.getIdentityProvidersStream().anyMatch(idp -> idp.getAlias().equals("idp-" + orgName + "-" + index)));
+    }
+}