add ext claims to refresh token

cloudfoundry · Feb 8, 2018 · faf7403 · faf7403
1 parent 01da490
commit faf7403
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 4 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
@@ -897,9 +897,11 @@ private ExpiringOAuth2RefreshToken createRefreshToken(UaaUser user, String token
         ExpiringOAuth2RefreshToken token = new DefaultExpiringOAuth2RefreshToken(tokenId,
                                                                                  new Date(System.currentTimeMillis() + (validitySeconds * 1000L)));
 
-        Map<String,String> externalAttributes = null;
+        Map<String,Object> externalAttributes = null;
         if (uaaTokenEnhancer != null) {
-            externalAttributes = uaaTokenEnhancer.getExternalAttributes(authentication);
+            externalAttributes = new HashMap<>();
+            externalAttributes.putAll(uaaTokenEnhancer.enhance(emptyMap(), authentication));
+            //externalAttributes = uaaTokenEnhancer.getExternalAttributes(authentication);
         }
 
         String content;
@@ -943,7 +945,7 @@ protected String getUserId(OAuth2Authentication authentication) {
         Set<String> resourceIds,
         String revocableSignature,
         boolean revocable,
-        Map<String, String> externalAttributes) {
+        Map<String, Object> externalAttributes) {
 
         Map<String, Object> response = new LinkedHashMap<String, Object>();
 
@@ -954,7 +956,7 @@ protected String getUserId(OAuth2Authentication authentication) {
             response.put(ADDITIONAL_AZ_ATTR, additionalAuthorizationAttributes);
         }
         if (null != externalAttributes) {
-            response.put(EXTERNAL_ATTR, externalAttributes);
+            response.putAll(externalAttributes);
         }
 
         response.put(IAT, System.currentTimeMillis() / 1000);

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
@@ -697,6 +697,26 @@ protected void validateExternalAttributes(OAuth2AccessToken accessToken) {
         }
     }
 
+    @Test
+    public void testCreateAccessTokenExternalContext() throws InterruptedException {
+        OAuth2AccessToken accessToken = getOAuth2AccessToken();
+
+        TokenRequest refreshTokenRequest = getRefreshTokenRequest();
+        String xx = accessToken.getRefreshToken().getValue();
+        OAuth2AccessToken refreshedAccessToken = tokenServices.refreshAccessToken(xx, refreshTokenRequest);
+        Map<String, Object> extendedContext = (Map<String, Object>) refreshedAccessToken.getAdditionalInformation();
+
+        if (tokenEnhancer!=null) {
+            assertNotNull(extendedContext);
+            assertEquals("test", ((Map<String, String>)extendedContext.get("ext_attr")).get("purpose"));
+            assertNotNull(extendedContext.get("ex_groups"));
+            assertNotNull(extendedContext.get("ex_prop"));
+            assertEquals("nz", ((Map<String, String>) extendedContext.get("ex_prop")).get("country"));
+        } else {
+            assertNull("External attributes should not exist", extendedContext.get("ext_attr"));
+        }
+    }
+
     @Test
     public void testCreateAccessTokenRefreshGrant() throws InterruptedException {
         OAuth2AccessToken accessToken = getOAuth2AccessToken();