Skip to content

Commit

Permalink
Assign instance_admin during User getOrCreateByAuthId if appropriate …
Browse files Browse the repository at this point in the history 
…(#8997)
  • Loading branch information
@pmossman
pmossman committed Sep 24, 2023
1 parent c371e48 commit 418af74
Show file tree
Hide file tree
Showing 2 changed files with 190 additions and 92 deletions.
42 changes: 38 additions & 4 deletions airbyte-commons-server/src/main/java/io/airbyte/commons/server/handlers/UserHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import io.airbyte.api.model.generated.WorkspaceIdRequestBody;
import io.airbyte.api.model.generated.WorkspaceUserRead;
import io.airbyte.api.model.generated.WorkspaceUserReadList;
import io.airbyte.commons.auth.config.InitialUserConfiguration;
import io.airbyte.commons.enums.Enums;
import io.airbyte.commons.server.support.JwtUserResolver;
import io.airbyte.config.ConfigSchema;
Expand Down Expand Up @@ -59,26 +60,26 @@ public class UserHandler {
private final PermissionPersistence permissionPersistence;
private final PermissionHandler permissionHandler;
private final OrganizationPersistence organizationPersistence;
private final OrganizationsHandler organizationsHandler;

private final Optional<JwtUserResolver> jwtUserResolver;
private final Optional<InitialUserConfiguration> initialUserConfiguration;

@VisibleForTesting
public UserHandler(
final UserPersistence userPersistence,
final PermissionPersistence permissionPersistence,
final OrganizationPersistence organizationPersistence,
final PermissionHandler permissionHandler,
final OrganizationsHandler organizationsHandler,
final Supplier<UUID> uuidGenerator,
final Optional<JwtUserResolver> jwtUserResolver) {
final Optional<JwtUserResolver> jwtUserResolver,
final Optional<InitialUserConfiguration> initialUserConfiguration) {
this.uuidGenerator = uuidGenerator;
this.userPersistence = userPersistence;
this.organizationPersistence = organizationPersistence;
this.organizationsHandler = organizationsHandler;
this.permissionPersistence = permissionPersistence;
this.permissionHandler = permissionHandler;
this.jwtUserResolver = jwtUserResolver;
this.initialUserConfiguration = initialUserConfiguration;
}

/**
Expand Down Expand Up @@ -309,6 +310,10 @@ public UserRead getOrCreateUserByAuthId(final UserAuthIdRequestBody userAuthIdRe
.authProvider(userAuthIdRequestBody.getAuthProvider())
.email(incomingUser.getEmail()));

// If new user's email matches the initial user config email, create instance_admin permission for
// them.
createInstanceAdminPermissionIfInitialUser(createdUser);

// If incoming SSO Config matches with existing org, find that org and add user to it;
final String ssoRealm = jwtUserResolver.get().resolveSsoRealm();
if (ssoRealm != null) {
Expand All @@ -326,6 +331,35 @@ public UserRead getOrCreateUserByAuthId(final UserAuthIdRequestBody userAuthIdRe
return createdUser;
}

private void createInstanceAdminPermissionIfInitialUser(final UserRead createdUser) throws IOException {
if (initialUserConfiguration.isEmpty()) {
// do nothing if initial_user bean is not present.
return;
}

final String initialEmailFromConfig = initialUserConfiguration.get().getEmail();

if (initialEmailFromConfig == null || initialEmailFromConfig.isEmpty()) {
// do nothing if there is no initial_user email configured.
return;
}

// compare emails with case insensitivity because different email cases should be treated as the
// same user.
if (!initialEmailFromConfig.equalsIgnoreCase(createdUser.getEmail())) {
return;
}

LOGGER.info("creating instance_admin permission for user ID {} because their email matches this instance's configured initial_user",
createdUser.getUserId());

permissionHandler.createPermission(new io.airbyte.api.model.generated.PermissionCreate()
.workspaceId(null)
.organizationId(null)
.userId(createdUser.getUserId())
.permissionType(PermissionType.INSTANCE_ADMIN));
}

private WorkspaceUserReadList buildWorkspaceUserReadList(final List<UserPermission> userPermissions, final UUID workspaceId) {

return new WorkspaceUserReadList().users(
Expand Down
240 changes: 152 additions & 88 deletions airbyte-commons-server/src/test/java/io/airbyte/commons/server/handlers/UserHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,16 @@

package io.airbyte.commons.server.handlers;

import static io.airbyte.config.User.AuthProvider.GOOGLE_IDENTITY_PLATFORM;
import static io.airbyte.config.User.AuthProvider.KEYCLOAK;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import io.airbyte.api.model.generated.OrganizationIdRequestBody;
import io.airbyte.api.model.generated.OrganizationRead;
import io.airbyte.api.model.generated.OrganizationUserRead;
import io.airbyte.api.model.generated.OrganizationUserReadList;
import io.airbyte.api.model.generated.PermissionCreate;
Expand All @@ -26,6 +25,8 @@
import io.airbyte.api.model.generated.WorkspaceIdRequestBody;
import io.airbyte.api.model.generated.WorkspaceUserRead;
import io.airbyte.api.model.generated.WorkspaceUserReadList;
import io.airbyte.commons.auth.config.InitialUserConfiguration;
import io.airbyte.commons.enums.Enums;
import io.airbyte.commons.server.support.JwtUserResolver;
import io.airbyte.config.Organization;
import io.airbyte.config.Permission;
Expand All @@ -40,13 +41,21 @@
import io.airbyte.config.persistence.UserPersistence;
import io.airbyte.validation.json.JsonValidationException;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import java.util.function.Supplier;
import java.util.stream.Stream;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtensionContext;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.ArgumentsProvider;
import org.junit.jupiter.params.provider.ArgumentsSource;
import org.junit.jupiter.params.provider.EnumSource;

class UserHandlerTest {

Expand All @@ -59,6 +68,7 @@ class UserHandlerTest {
OrganizationPersistence organizationPersistence;
OrganizationsHandler organizationsHandler;
JwtUserResolver jwtUserResolver;
InitialUserConfiguration initialUserConfiguration;

private static final UUID USER_ID = UUID.randomUUID();
private static final String USER_NAME = "user 1";
Expand All @@ -84,9 +94,10 @@ void setUp() {
organizationsHandler = mock(OrganizationsHandler.class);
uuidSupplier = mock(Supplier.class);
jwtUserResolver = mock(JwtUserResolver.class);
initialUserConfiguration = mock(InitialUserConfiguration.class);

userHandler = new UserHandler(userPersistence, permissionPersistence, organizationPersistence, permissionHandler, organizationsHandler,
uuidSupplier, Optional.of(jwtUserResolver));
userHandler = new UserHandler(userPersistence, permissionPersistence, organizationPersistence, permissionHandler,
uuidSupplier, Optional.of(jwtUserResolver), Optional.of(initialUserConfiguration));
}

@Test
Expand Down Expand Up @@ -173,98 +184,151 @@ void testListInstanceAdminUser() throws Exception {
@Nested
class GetOrCreateUserByAuthIdTest {

private static final String SSO_REALM = "airbyte-realm";
@ParameterizedTest
@EnumSource(value = AuthProvider.class)
void authIdExists(final AuthProvider authProvider) throws Exception {
// set the auth provider for the existing user to match the test case
user.setAuthProvider(authProvider);

private static final String KEY_CLOAK_AUTH_ID = "key_cloack_auth_id";
private static final String GOOGLE_AUTH_ID = "google_auth_id";
// authUserId is for the existing user
final String authUserId = user.getAuthUserId();
final io.airbyte.api.model.generated.AuthProvider apiAuthProvider =
Enums.convertTo(authProvider, io.airbyte.api.model.generated.AuthProvider.class);

private static final User RESOLVED_USER_TEMPLATE = new User().withEmail(USER_EMAIL).withName(USER_NAME);
when(userPersistence.getUserByAuthId(authUserId, authProvider)).thenReturn(Optional.of(user));

@Test
void testGetOrCreateUserByAuthId_authIdExists() throws Exception {
when(userPersistence.getUserByAuthId(KEY_CLOAK_AUTH_ID, KEYCLOAK)).thenReturn(Optional.of(user));
UserRead userRead = userHandler.getOrCreateUserByAuthId(
new UserAuthIdRequestBody().authProvider(io.airbyte.api.model.generated.AuthProvider.KEYCLOAK).authUserId(KEY_CLOAK_AUTH_ID));

assertEquals(userRead.getUserId(), user.getUserId());
}

@Test
void testGetOrCreateUserByAuthId_firebaseUser_authIdNotExists() throws Exception {
when(userPersistence.getUserByAuthId(GOOGLE_AUTH_ID, GOOGLE_IDENTITY_PLATFORM)).thenReturn(Optional.empty());

final User resolvedUser =
RESOLVED_USER_TEMPLATE.withAuthUserId(GOOGLE_AUTH_ID).withAuthProvider(AuthProvider.GOOGLE_IDENTITY_PLATFORM);

when(jwtUserResolver.resolveUser()).thenReturn(resolvedUser);
when(jwtUserResolver.resolveSsoRealm()).thenReturn(null);
when(uuidSupplier.get()).thenReturn(USER_ID);
when(userPersistence.getUser(USER_ID)).thenReturn(Optional.of(user));

when(organizationsHandler.createOrganization(any())).thenReturn(new OrganizationRead().organizationId(ORGANIZATION.getOrganizationId()));

final UserRead userRead = userHandler.getOrCreateUserByAuthId(
new UserAuthIdRequestBody().authProvider(io.airbyte.api.model.generated.AuthProvider.GOOGLE_IDENTITY_PLATFORM).authUserId(GOOGLE_AUTH_ID));

assertEquals(userRead.getUserId(), USER_ID);
assertEquals(userRead.getEmail(), USER_EMAIL);

verify(userPersistence).writeUser(any());

verify(organizationsHandler, never()).createOrganization(any());
verify(permissionHandler, never()).createPermission(any());
}

@Test
void testGetOrCreateUserByAuthId_SsoUser_authIdNotExists() throws Exception {
when(userPersistence.getUserByAuthId(KEY_CLOAK_AUTH_ID, KEYCLOAK)).thenReturn(Optional.empty());

final User resolvedUser =
RESOLVED_USER_TEMPLATE.withAuthUserId(KEY_CLOAK_AUTH_ID).withAuthProvider(AuthProvider.KEYCLOAK);

when(jwtUserResolver.resolveUser()).thenReturn(resolvedUser);
when(jwtUserResolver.resolveSsoRealm()).thenReturn(SSO_REALM);
when(organizationPersistence.getOrganizationBySsoConfigRealm(SSO_REALM)).thenReturn(Optional.of(ORGANIZATION));
when(uuidSupplier.get()).thenReturn(USER_ID);

when(userPersistence.getUser(USER_ID)).thenReturn(Optional.of(user));

UserRead userRead = userHandler.getOrCreateUserByAuthId(
new UserAuthIdRequestBody().authProvider(io.airbyte.api.model.generated.AuthProvider.KEYCLOAK).authUserId(KEY_CLOAK_AUTH_ID));

verify(userPersistence).writeUser(any());

verify(permissionHandler).createPermission(new PermissionCreate()
.permissionType(io.airbyte.api.model.generated.PermissionType.ORGANIZATION_ADMIN).organizationId(ORGANIZATION.getOrganizationId())
.userId(USER_ID));
final UserRead userRead = userHandler.getOrCreateUserByAuthId(new UserAuthIdRequestBody()
.authProvider(apiAuthProvider)
.authUserId(authUserId));

assertEquals(userRead.getUserId(), USER_ID);
assertEquals(userRead.getEmail(), USER_EMAIL);
assertEquals(userRead.getAuthUserId(), authUserId);
assertEquals(userRead.getAuthProvider(), apiAuthProvider);
}

@Test
void testGetOrCreateUserByAuthId_SsoUser_noAuthIdNorOrg() throws Exception {
when(userPersistence.getUserByAuthId(KEY_CLOAK_AUTH_ID, KEYCLOAK)).thenReturn(Optional.empty());

final User resolvedUser =
RESOLVED_USER_TEMPLATE.withAuthUserId(KEY_CLOAK_AUTH_ID).withAuthProvider(AuthProvider.KEYCLOAK);

when(jwtUserResolver.resolveUser()).thenReturn(resolvedUser);
when(jwtUserResolver.resolveSsoRealm()).thenReturn(SSO_REALM);
when(organizationPersistence.getOrganizationBySsoConfigRealm(SSO_REALM)).thenReturn(Optional.empty());
when(uuidSupplier.get()).thenReturn(USER_ID);
when(userPersistence.getUser(USER_ID)).thenReturn(Optional.of(user));
when(organizationsHandler.createOrganization(any())).thenReturn(new OrganizationRead().organizationId(ORGANIZATION.getOrganizationId()));

final UserRead userRead = userHandler.getOrCreateUserByAuthId(
new UserAuthIdRequestBody().authProvider(io.airbyte.api.model.generated.AuthProvider.KEYCLOAK).authUserId(KEY_CLOAK_AUTH_ID));

verify(userPersistence).writeUser(any());
assertEquals(userRead.getUserId(), USER_ID);
assertEquals(userRead.getEmail(), USER_EMAIL);
@Nested
class NewUser {

private static final String NEW_AUTH_USER_ID = "new_auth_user_id";
private static final UUID NEW_USER_ID = UUID.randomUUID();
private static final String NEW_EMAIL = "[email protected]";

private User newUser;

// this class provides the arguments for the parameterized test below, by returning all
// permutations of auth provider, sso realm, initial user email, and deployment mode
static class NewUserArgumentsProvider implements ArgumentsProvider {

@Override
public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
List<AuthProvider> authProviders = Arrays.asList(AuthProvider.values());
List<String> ssoRealms = Arrays.asList("airbyte-realm", null);
List<String> initialUserEmails = Arrays.asList(null, "", "[email protected]", NEW_EMAIL);
List<Boolean> initialUserConfigPresent = Arrays.asList(true, false);

// return all permutations of auth provider, sso realm, and initial user email that we want to test
return authProviders.stream()
.flatMap(authProvider -> ssoRealms.stream().flatMap(ssoRealm -> initialUserEmails.stream().flatMap(email -> initialUserConfigPresent
.stream().flatMap(initialUserPresent -> Stream.of(Arguments.of(authProvider, ssoRealm, email, initialUserPresent))))));
}

}

@BeforeEach
void setUp() throws IOException {
newUser = new User().withUserId(NEW_USER_ID).withEmail(NEW_EMAIL).withAuthUserId(NEW_AUTH_USER_ID);

when(userPersistence.getUserByAuthId(anyString(), any())).thenReturn(Optional.empty());
when(jwtUserResolver.resolveUser()).thenReturn(newUser);
when(uuidSupplier.get()).thenReturn(NEW_USER_ID);
when(userPersistence.getUser(NEW_USER_ID)).thenReturn(Optional.of(newUser));
}

@ParameterizedTest
@ArgumentsSource(NewUserArgumentsProvider.class)
void testNewUserCreation(final AuthProvider authProvider,
final String ssoRealm,
final String initialUserEmail,
final boolean initialUserPresent)
throws Exception {

newUser.setAuthProvider(authProvider);

when(jwtUserResolver.resolveSsoRealm()).thenReturn(ssoRealm);
if (ssoRealm != null) {
when(organizationPersistence.getOrganizationBySsoConfigRealm(ssoRealm)).thenReturn(Optional.of(ORGANIZATION));
}

if (initialUserPresent) {
if (initialUserEmail != null) {
when(initialUserConfiguration.getEmail()).thenReturn(initialUserEmail);
}
} else {
// replace default user handler with one that doesn't use initial user config (ie to test what
// happens in Cloud)
userHandler = new UserHandler(userPersistence, permissionPersistence, organizationPersistence, permissionHandler,
uuidSupplier, Optional.of(jwtUserResolver), Optional.empty());
}

final io.airbyte.api.model.generated.AuthProvider apiAuthProvider =
Enums.convertTo(authProvider, io.airbyte.api.model.generated.AuthProvider.class);

final UserRead userRead = userHandler.getOrCreateUserByAuthId(
new UserAuthIdRequestBody().authProvider(apiAuthProvider).authUserId(NEW_AUTH_USER_ID));

verifyCreatedUser(authProvider);
verifyUserRead(userRead, apiAuthProvider);
verifyInstanceAdminPermissionCreation(initialUserEmail, initialUserPresent);
verifyOrganizationPermissionCreation(ssoRealm);
}

private void verifyCreatedUser(final AuthProvider expectedAuthProvider) throws IOException {
verify(userPersistence).writeUser(argThat(user -> user.getUserId().equals(NEW_USER_ID) &&
user.getEmail().equals(NEW_EMAIL) &&
user.getAuthUserId().equals(NEW_AUTH_USER_ID) &&
user.getAuthProvider().equals(expectedAuthProvider)));
}

private void verifyUserRead(final UserRead userRead, final io.airbyte.api.model.generated.AuthProvider expectedAuthProvider) {
assertEquals(userRead.getUserId(), NEW_USER_ID);
assertEquals(userRead.getEmail(), NEW_EMAIL);
assertEquals(userRead.getAuthUserId(), NEW_AUTH_USER_ID);
assertEquals(userRead.getAuthProvider(), expectedAuthProvider);
}

private void verifyInstanceAdminPermissionCreation(final String initialUserEmail, final boolean initialUserPresent) throws IOException {
// instance_admin permissions should only ever be created when the initial user config is present
// (which should never be true in Cloud).
// also, if the initial user email is null or doesn't match the new user's email, no instance_admin
// permission should be created
if (!initialUserPresent || initialUserEmail == null || !initialUserEmail.equalsIgnoreCase(NEW_EMAIL)) {
verify(permissionHandler, never()).createPermission(
argThat(permission -> permission.getPermissionType().equals(io.airbyte.api.model.generated.PermissionType.INSTANCE_ADMIN)));
} else {
// otherwise, instance_admin permission should be created
verify(permissionHandler).createPermission(new PermissionCreate()
.permissionType(io.airbyte.api.model.generated.PermissionType.INSTANCE_ADMIN)
.workspaceId(null)
.organizationId(null)
.userId(NEW_USER_ID));
}
}

private void verifyOrganizationPermissionCreation(final String ssoRealm) throws IOException {
// if the SSO Realm is null, no organization permission should be created
if (ssoRealm == null) {
verify(permissionHandler, never()).createPermission(
argThat(permission -> permission.getPermissionType().equals(io.airbyte.api.model.generated.PermissionType.ORGANIZATION_ADMIN)));
} else {
// otherwise, organization permission should be created for the associated user and org.
verify(permissionHandler).createPermission(new PermissionCreate()
.permissionType(io.airbyte.api.model.generated.PermissionType.ORGANIZATION_ADMIN)
.organizationId(ORGANIZATION.getOrganizationId())
.userId(NEW_USER_ID));
}
}

verify(organizationsHandler, never()).createOrganization(any());
verify(permissionHandler, never()).createPermission(any());
}

}
Expand Down

0 comments on commit 418af74

Please sign in to comment.