diff --git a/api/src/main/java/io/kafbat/ui/service/rbac/extractor/CognitoAuthorityExtractor.java b/api/src/main/java/io/kafbat/ui/service/rbac/extractor/CognitoAuthorityExtractor.java index cc0e419bf..02825884b 100644 --- a/api/src/main/java/io/kafbat/ui/service/rbac/extractor/CognitoAuthorityExtractor.java +++ b/api/src/main/java/io/kafbat/ui/service/rbac/extractor/CognitoAuthorityExtractor.java @@ -3,21 +3,25 @@ import static io.kafbat.ui.model.rbac.provider.Provider.Name.COGNITO; import com.google.common.collect.Sets; +import io.kafbat.ui.config.auth.OAuthProperties; import io.kafbat.ui.model.rbac.Role; import io.kafbat.ui.model.rbac.provider.Provider; import io.kafbat.ui.service.rbac.AccessControlService; import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; import lombok.extern.slf4j.Slf4j; import org.springframework.security.oauth2.core.user.DefaultOAuth2User; +import org.springframework.util.Assert; import reactor.core.publisher.Mono; @Slf4j public class CognitoAuthorityExtractor implements ProviderAuthorityExtractor { + public static final String ROLES_FIELD_PARAM_NAME = "roles-field"; private static final String COGNITO_GROUPS_ATTRIBUTE_NAME = "cognito:groups"; @Override @@ -38,7 +42,7 @@ public Mono> extract(AccessControlService acs, Object value, Map extractUsernameRoles(AccessControlService acs, DefaultOAuth2 return rolesByUsername; } - private Set extractGroupRoles(AccessControlService acs, DefaultOAuth2User principal) { - List groups = principal.getAttribute(COGNITO_GROUPS_ATTRIBUTE_NAME); + private Set extractGroupRoles(AccessControlService acs, DefaultOAuth2User principal, + Map additionalParams) { + var provider = (OAuthProperties.OAuth2Provider) additionalParams.get("provider"); + Assert.notNull(provider, "provider is null"); + + var rolesFieldName = Optional.ofNullable(provider.getCustomParams().get(ROLES_FIELD_PARAM_NAME)) + .orElse(COGNITO_GROUPS_ATTRIBUTE_NAME); + + List groups = principal.getAttribute(rolesFieldName); if (groups == null) { log.debug("Cognito groups param is not present"); return Collections.emptySet();