Skip to content

Commit

Permalink
Dataset for organizations
Browse files Browse the repository at this point in the history
Closes #856

Signed-off-by: Pedro Igor <[email protected]>
  • Loading branch information
pedroigor committed Jun 18, 2024
1 parent 427369a commit 7fce05d
Show file tree
Hide file tree
Showing 8 changed files with 702 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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 []";

Expand Down Expand Up @@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ public enum DatasetOperation {
CREATE_USERS,
CREATE_EVENTS,
CREATE_OFFLINE_SESSIONS,
CREATE_ORGS,
REMOVE_REALMS,
LAST_REALM,
LAST_CLIENT,
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}
Original file line number Diff line number Diff line change
@@ -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)));
}
}
Loading

0 comments on commit 7fce05d

Please sign in to comment.