Skip to content

refactor: allow multiple UaaTokenEnhancers #3904

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
peterhaochen47 wants to merge 2 commits into
base: develop
Choose a base branch
from
+172 −5
Open

refactor: allow multiple UaaTokenEnhancers #3904

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
28feee4
refactor: allow multiple UaaTokenEnhancers
peterhaochen47 May 5, 2026
af0b878
feat: pass accumulated claims to subsequent token enhancers
peterhaochen47 May 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 20 additions & 5 deletions server/src/main/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenServices.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,10 @@
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import static java.util.Collections.emptyList;
import static java.util.Collections.emptyMap;
import static java.util.Collections.emptySet;
import static java.util.Collections.singletonList;
import static java.util.Optional.ofNullable;
import static org.cloudfoundry.identity.uaa.oauth.client.ClientConstants.REQUIRED_USER_GROUPS;
import static org.cloudfoundry.identity.uaa.oauth.openid.IdToken.ACR_VALUES_KEY;
Expand Down Expand Up @@ -143,7 +145,7 @@ public class UaaTokenServices implements AuthorizationServerTokenServices, Resou
private final TokenPolicy tokenPolicy;
private final RevocableTokenProvisioning tokenProvisioning;
private Set<String> excludedClaims;
private UaaTokenEnhancer uaaTokenEnhancer;
private List<UaaTokenEnhancer> uaaTokenEnhancers = new ArrayList<>();
private final IdTokenCreator idTokenCreator;
private final RefreshTokenCreator refreshTokenCreator;
private TokenEndpointBuilder tokenEndpointBuilder;
Expand Down Expand Up @@ -192,8 +194,13 @@ public void setExcludedClaims(Set<String> excludedClaims) {
}

@Autowired(required = false)
public void setUaaTokenEnhancers(List<UaaTokenEnhancer> uaaTokenEnhancers) {
this.uaaTokenEnhancers = new ArrayList<>(uaaTokenEnhancers == null ? emptyList() : uaaTokenEnhancers);
}

@Deprecated
public void setUaaTokenEnhancer(UaaTokenEnhancer uaaTokenEnhancer) {
this.uaaTokenEnhancer = uaaTokenEnhancer;
this.setUaaTokenEnhancers(uaaTokenEnhancer == null ? emptyList() : singletonList(uaaTokenEnhancer));
}

@Override
Expand Down Expand Up @@ -349,7 +356,7 @@ Claims getClaims(Map<String, Object> refreshTokenClaims) {

private Map<String, Object> getAdditionalRootClaims(Map<String, Object> refreshTokenClaims) {
Map<String, Object> additionalRootClaims = new HashMap<>();
if (uaaTokenEnhancer != null) {
if (!uaaTokenEnhancers.isEmpty()) {
refreshTokenClaims.entrySet()
.stream()
.filter(entry -> !NON_ADDITIONAL_ROOT_CLAIMS.contains(entry.getKey()))
Expand Down Expand Up @@ -625,8 +632,16 @@ public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication)
boolean isRefreshTokenRevocable = isAccessTokenRevocable || OPAQUE.getStringValue().equals(getActiveTokenPolicy().getRefreshTokenFormat());

Map<String, Object> additionalRootClaims = null;
if (uaaTokenEnhancer != null) {
additionalRootClaims = new HashMap<>(uaaTokenEnhancer.enhance(emptyMap(), authentication));
if (!uaaTokenEnhancers.isEmpty()) {
additionalRootClaims = new HashMap<>();
for (UaaTokenEnhancer enhancer : uaaTokenEnhancers) {
if (enhancer != null) {
Map<String, Object> claims = enhancer.enhance(additionalRootClaims, authentication);
if (claims != null) {
additionalRootClaims.putAll(claims);
Copy link
Copy Markdown
Member Author

@peterhaochen47 peterhaochen47 May 8, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure merging the claim values from different token enhancers is a good idea; that would result in unpredictable/unintended final outcome. It seems more rigorous that an enhancer has deterministic power of the final outcome of the claim it touches, and if there is a conflict, the later enhancer wins.

If we do wanna let two enhancers touch on the same claim (not a use case now), we could pass in the previously added claims resulted from prior enhancers to the current enhancer, like an assembly line, so the current enhancer knows what has already been added to what claims (and merge if needed):

enhancer.enhance(additionalRootClaims, authentication);

But we don't need this for now; we could always add this when we need to. But once we add this, we can't remove it without a breaking change. So I'm holding off on that for now.

Copy link
Copy Markdown
Member

@duanemay duanemay May 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this idea, so each successive enhancer (ordered by @order) can see previous changes and update as needed.

Copy link
Copy Markdown
Member Author

@peterhaochen47 peterhaochen47 May 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, we can add that if you think this is the time to do so. Added a new commit.

}
}
}
}

String clientAuthentication = getAuthenticationMethod(oAuth2Request);
Expand Down
48 changes: 48 additions & 0 deletions ...er/src/test/java/org/cloudfoundry/identity/uaa/oauth/DeprecatedUaaTokenServicesTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -582,6 +582,54 @@ void createOpaqueAccessTokenForAClient(TestTokenEnhancer enhancer) {
assertThat(accessToken.getRefreshToken()).isNull();
}

@MethodSource("data")
@ParameterizedTest(name = "{index}: {0}")
void multipleTokenEnhancersAreSupported(TestTokenEnhancer enhancer) {
initDeprecatedUaaTokenServicesTests(enhancer);
UaaTokenEnhancer enhancer1 = new UaaTokenEnhancer() {
@Override
public Map<String, String> getExternalAttributes(OAuth2Authentication authentication) {
return Map.of();
}

@Override
public Map<String, Object> enhance(Map<String, Object> claims, OAuth2Authentication authentication) {
return Map.of("claim1", "value1");
}
};

UaaTokenEnhancer enhancer2 = new UaaTokenEnhancer() {
@Override
public Map<String, String> getExternalAttributes(OAuth2Authentication authentication) {
return Map.of();
}

@Override
public Map<String, Object> enhance(Map<String, Object> claims, OAuth2Authentication authentication) {
if (claims.containsKey("claim1")) {
return Map.of("claim2", claims.get("claim1") + "_modified");
}
return Map.of("claim2", "value2");
}
};

tokenServices.setUaaTokenEnhancers(java.util.Arrays.asList(enhancer1, enhancer2));

AuthorizationRequest authorizationRequest = new AuthorizationRequest(CLIENT_ID, tokenSupport.clientScopes);
authorizationRequest.setResourceIds(new java.util.HashSet<>(tokenSupport.resourceIds));
authorizationRequest.setRequestParameters(new java.util.HashMap<>());
OAuth2Authentication authentication = new OAuth2Authentication(authorizationRequest.createOAuth2Request(), null);

OAuth2AccessToken accessToken = tokenServices.createAccessToken(authentication);

String jwt = accessToken.getValue();
org.cloudfoundry.identity.uaa.oauth.jwt.Jwt parsedToken = org.cloudfoundry.identity.uaa.oauth.jwt.JwtHelper.decode(jwt);
Map<String, Object> claims = org.cloudfoundry.identity.uaa.util.JsonUtils.readValue(parsedToken.getClaims(), new com.fasterxml.jackson.core.type.TypeReference<Map<String, Object>>() {});

assertThat(claims).containsEntry("claim1", "value1");
assertThat(claims).containsEntry("claim2", "value1_modified");
}

@MethodSource("data")
@ParameterizedTest(name = "{index}: {0}")
void createAccessTokenForAClientInAnotherIdentityZone(TestTokenEnhancer enhancer) {
Expand Down
104 changes: 104 additions & 0 deletions uaa/src/test/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenServicesTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,110 @@ void ensureAnIdTokenIsNotReturned(String grantType) {
}
}

@Nested
@DisplayName("when multiple token enhancers are provided")
@DefaultTestContext
@TestPropertySource(properties = {"uaa.url=https://uaa.some.test.domain.com:555/uaa"})
class WhenMultipleTokenEnhancersAreProvided {

@DisplayName("claims from all enhancers are merged into the token")
@ParameterizedTest
@ValueSource(strings = {GRANT_TYPE_PASSWORD, GRANT_TYPE_AUTHORIZATION_CODE})
void claimsAreMerged(String grantType) {
UaaTokenEnhancer enhancer1 = new UaaTokenEnhancer() {
@Override
public Map<String, String> getExternalAttributes(OAuth2Authentication authentication) {
return Map.of();
}

@Override
public Map<String, Object> enhance(Map<String, Object> claims, OAuth2Authentication authentication) {
return Map.of("claim1", "value1");
}
};

UaaTokenEnhancer enhancer2 = new UaaTokenEnhancer() {
@Override
public Map<String, String> getExternalAttributes(OAuth2Authentication authentication) {
return Map.of();
}

@Override
public Map<String, Object> enhance(Map<String, Object> claims, OAuth2Authentication authentication) {
return Map.of("claim2", "value2");
}
};

tokenServices.setUaaTokenEnhancers(Arrays.asList(enhancer1, enhancer2));

try {
AuthorizationRequest authorizationRequest = constructAuthorizationRequest(clientId, grantType, "openid", "user_attributes");
OAuth2Authentication auth2Authentication = constructUserAuthenticationFromAuthzRequest(authorizationRequest, "admin", "uaa");

CompositeToken accessToken = (CompositeToken) tokenServices.createAccessToken(auth2Authentication);

String jwt = accessToken.getValue();
Jwt parsedToken = JwtHelper.decode(jwt);
Map<String, Object> claims = JsonUtils.readValue(parsedToken.getClaims(), new TypeReference<Map<String, Object>>() {});

assertThat(claims).containsEntry("claim1", "value1");
assertThat(claims).containsEntry("claim2", "value2");
} finally {
tokenServices.setUaaTokenEnhancers(new ArrayList<>());
}
}

@DisplayName("claims are passed to subsequent enhancers")
@ParameterizedTest
@ValueSource(strings = {GRANT_TYPE_PASSWORD, GRANT_TYPE_AUTHORIZATION_CODE})
void claimsArePassedToSubsequentEnhancers(String grantType) {
UaaTokenEnhancer enhancer1 = new UaaTokenEnhancer() {
@Override
public Map<String, String> getExternalAttributes(OAuth2Authentication authentication) {
return Map.of();
}

@Override
public Map<String, Object> enhance(Map<String, Object> claims, OAuth2Authentication authentication) {
return Map.of("claim1", "value1");
}
};

UaaTokenEnhancer enhancer2 = new UaaTokenEnhancer() {
@Override
public Map<String, String> getExternalAttributes(OAuth2Authentication authentication) {
return Map.of();
}

@Override
public Map<String, Object> enhance(Map<String, Object> claims, OAuth2Authentication authentication) {
if (claims.containsKey("claim1")) {
return Map.of("claim2", claims.get("claim1") + "_modified");
}
return Map.of("claim2", "value2");
}
};

tokenServices.setUaaTokenEnhancers(Arrays.asList(enhancer1, enhancer2));

try {
AuthorizationRequest authorizationRequest = constructAuthorizationRequest(clientId, grantType, "openid", "user_attributes");
OAuth2Authentication auth2Authentication = constructUserAuthenticationFromAuthzRequest(authorizationRequest, "admin", "uaa");

CompositeToken accessToken = (CompositeToken) tokenServices.createAccessToken(auth2Authentication);

String jwt = accessToken.getValue();
Jwt parsedToken = JwtHelper.decode(jwt);
Map<String, Object> claims = JsonUtils.readValue(parsedToken.getClaims(), new TypeReference<Map<String, Object>>() {});

assertThat(claims).containsEntry("claim1", "value1");
assertThat(claims).containsEntry("claim2", "value1_modified");
} finally {
tokenServices.setUaaTokenEnhancers(new ArrayList<>());
}
}
}

@Nested
@DisplayName("when the hasn't approved the 'openid' scope")
@DefaultTestContext
Expand Down
Loading