Refactor token enhancer call

- only once - add info only if interface used
cloudfoundry · Feb 8, 2018 · 995c559 · 995c559
1 parent 749b2c6
commit 995c559
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 20 deletions.
diff --git a/server/src/main/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenServices.java b/server/src/main/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenServices.java
@@ -445,10 +445,12 @@ private CompositeAccessToken createAccessToken(String tokenId,
             info.put(NONCE, nonce);
         }
 
-        additionalRootClaims
-            .entrySet()
-            .stream()
-            .forEach(entry -> info.putIfAbsent(entry.getKey(), entry.getValue()));
+        if(additionalRootClaims != null) {
+            additionalRootClaims
+                .entrySet()
+                .stream()
+                .forEach(entry -> info.putIfAbsent(entry.getKey(), entry.getValue()));
+        }
 
         accessToken.setAdditionalInformation(info);
 
@@ -676,9 +678,15 @@ public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication)
         boolean refreshTokenRevocable = accessTokenRevocable || TokenConstants.TokenFormat.OPAQUE.getStringValue().equals(IdentityZoneHolder.get().getConfig().getTokenPolicy().getRefreshTokenFormat());
 
         OAuth2RefreshToken refreshToken = null;
+        Map<String,Object> additionalRootClaims = null;
+
+        if (uaaTokenEnhancer != null) {
+            additionalRootClaims = new HashMap<>();
+            additionalRootClaims.putAll(uaaTokenEnhancer.enhance(emptyMap(), authentication));
+        }
 
         if(client.getAuthorizedGrantTypes().contains(GRANT_TYPE_REFRESH_TOKEN)){
-            refreshToken = createRefreshToken(user, refreshTokenId, authentication, revocableHashSignature, refreshTokenRevocable);
+            refreshToken = createRefreshToken(user, refreshTokenId, authentication, revocableHashSignature, refreshTokenRevocable, additionalRootClaims);
         }
 
         String clientId = authentication.getOAuth2Request().getClientId();
@@ -713,12 +721,6 @@ public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication)
         Integer validity = client.getAccessTokenValiditySeconds();
         Set<String> responseTypes = extractResponseTypes(authentication);
 
-        Map<String,Object> additionalRootClaims = new HashMap<>();
-
-        if (uaaTokenEnhancer != null) {
-            additionalRootClaims.putAll(uaaTokenEnhancer.enhance(emptyMap(), authentication));
-        }
-
         CompositeAccessToken accessToken =
             createAccessToken(
                 tokenId,
@@ -876,7 +878,8 @@ private Map<String, String> getAdditionalAuthorizationAttributes(String authorit
     private ExpiringOAuth2RefreshToken createRefreshToken(UaaUser user, String tokenId,
                                                           OAuth2Authentication authentication,
                                                           String revocableHashSignature,
-                                                          boolean revocable) {
+                                                          boolean revocable,
+                                                          Map<String,Object> externalAttributes) {
 
         String grantType = authentication.getOAuth2Request().getRequestParameters().get("grant_type");
         Set<String> scope = authentication.getOAuth2Request().getScope();
@@ -891,12 +894,6 @@ private ExpiringOAuth2RefreshToken createRefreshToken(UaaUser user, String token
         ExpiringOAuth2RefreshToken token = new DefaultExpiringOAuth2RefreshToken(tokenId,
                                                                                  new Date(System.currentTimeMillis() + (validitySeconds * 1000L)));
 
-        Map<String,Object> externalAttributes = null;
-        if (uaaTokenEnhancer != null) {
-            externalAttributes = new HashMap<>();
-            externalAttributes.putAll(uaaTokenEnhancer.enhance(emptyMap(), authentication));
-        }
-
         String content;
         try {
             content = JsonUtils.writeValueAsString(

diff --git a/server/src/test/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenServicesTests.java b/server/src/test/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenServicesTests.java
@@ -705,8 +705,7 @@ public void testCreateAccessTokenExternalContext() throws InterruptedException {
         OAuth2AccessToken accessToken = getOAuth2AccessToken();
 
         TokenRequest refreshTokenRequest = getRefreshTokenRequest();
-        String xx = accessToken.getRefreshToken().getValue();
-        OAuth2AccessToken refreshedAccessToken = tokenServices.refreshAccessToken(xx, refreshTokenRequest);
+        OAuth2AccessToken refreshedAccessToken = tokenServices.refreshAccessToken(accessToken.getRefreshToken().getValue(), refreshTokenRequest);
         Map<String, Object> extendedContext = (Map<String, Object>) refreshedAccessToken.getAdditionalInformation();
 
         if (tokenEnhancer!=null) {