Fix wso2/product-apim#6466

bhathiya · bhathiya · commit b00cbad4d919 · 2019-10-16T01:07:49.000+05:30
diff --git a/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/main/java/org/wso2/carbon/apimgt/keymgt/issuers/APIMTokenIssuer.java b/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/main/java/org/wso2/carbon/apimgt/keymgt/issuers/APIMTokenIssuer.java
@@ -39,7 +39,6 @@
 import org.wso2.carbon.apimgt.keymgt.token.APIMJWTGenerator;
 import org.wso2.carbon.apimgt.keymgt.token.TokenGenerator;
 import org.wso2.carbon.apimgt.keymgt.util.APIKeyMgtDataHolder;
-import org.wso2.carbon.apimgt.keymgt.util.APIMTokenIssuerUtil;
 import org.wso2.carbon.identity.base.IdentityConstants;
 import org.wso2.carbon.identity.core.util.IdentityUtil;
 import org.wso2.carbon.identity.oauth.common.exception.InvalidOAuthClientException;
@@ -114,13 +113,17 @@ public String accessToken(OAuthTokenReqMessageContext tokReqMsgCtx) throws OAuth
                         if (jwtGenerator != null) {
                             TokenValidationContext validationContext = new TokenValidationContext();
                             APIKeyValidationInfoDTO apiKeyValidationInfoDTO = new APIKeyValidationInfoDTO();
-                            apiKeyValidationInfoDTO.setEndUserName(tokReqMsgCtx.getAuthorizedUser().toFullQualifiedUsername());
+                            apiKeyValidationInfoDTO.setEndUserName(tokReqMsgCtx.getAuthorizedUser()
+                                    .toFullQualifiedUsername());
                             apiKeyValidationInfoDTO.setSubscriber(application.getOwner());
                             apiKeyValidationInfoDTO.setApplicationName(application.getName());
                             apiKeyValidationInfoDTO.setApplicationId(String.valueOf(application.getId()));
                             apiKeyValidationInfoDTO.setType(application.getKeyType());
                             apiKeyValidationInfoDTO.setApplicationTier(application.getTier());
                             validationContext.setValidationInfoDTO(apiKeyValidationInfoDTO);
+                            validationContext.setUser(tokReqMsgCtx.getAuthorizedUser());
+                            validationContext.setAuthorizationCode(tokReqMsgCtx.getOauth2AccessTokenReqDTO()
+                                    .getAuthorizationCode());
                             jwtTokenInfoDTO.setBackendJwt(jwtGenerator.generateToken(validationContext));
                         }
                     }
diff --git a/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/main/java/org/wso2/carbon/apimgt/keymgt/service/TokenValidationContext.java b/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/main/java/org/wso2/carbon/apimgt/keymgt/service/TokenValidationContext.java
@@ -22,6 +22,7 @@
 
 import org.wso2.carbon.apimgt.api.model.AccessTokenInfo;
 import org.wso2.carbon.apimgt.impl.dto.APIKeyValidationInfoDTO;
+import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedUser;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -43,6 +44,8 @@ public class TokenValidationContext {
     private APIKeyValidationInfoDTO validationInfoDTO;
     private boolean isCacheHit;
     private AccessTokenInfo tokenInfo;
+    private AuthenticatedUser user;
+    private String authorizationCode;
 
     public AccessTokenInfo getTokenInfo() {
         return tokenInfo;
@@ -140,4 +143,19 @@ public Object getAttribute(String key){
         return this.attributeMap.get(key);
     }
 
+    public AuthenticatedUser getUser() {
+        return user;
+    }
+
+    public void setUser(AuthenticatedUser user) {
+        this.user = user;
+    }
+
+    public String getAuthorizationCode() {
+        return authorizationCode;
+    }
+
+    public void setAuthorizationCode(String authorizationCode) {
+        this.authorizationCode = authorizationCode;
+    }
 }
diff --git a/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/main/java/org/wso2/carbon/apimgt/keymgt/token/JWTGenerator.java b/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/main/java/org/wso2/carbon/apimgt/keymgt/token/JWTGenerator.java
@@ -21,12 +21,11 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-
 import org.wso2.carbon.apimgt.api.APIManagementException;
+import org.wso2.carbon.apimgt.api.model.Application;
 import org.wso2.carbon.apimgt.impl.internal.ServiceReferenceHolder;
 import org.wso2.carbon.apimgt.impl.token.ClaimsRetriever;
 import org.wso2.carbon.apimgt.impl.utils.APIUtil;
-import org.wso2.carbon.apimgt.api.model.Application;
 import org.wso2.carbon.apimgt.keymgt.MethodStats;
 import org.wso2.carbon.apimgt.keymgt.service.TokenValidationContext;
 import org.wso2.carbon.identity.application.common.model.ClaimMapping;
@@ -37,7 +36,9 @@
 import org.wso2.carbon.user.api.UserStoreManager;
 import org.wso2.carbon.utils.multitenancy.MultitenantUtils;
 
-import java.util.*;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
 
 import static org.apache.commons.collections.MapUtils.isNotEmpty;
 
@@ -113,11 +114,31 @@ public Map<String, String> populateCustomClaims(TokenValidationContext validatio
             throws APIManagementException {
         ClaimsRetriever claimsRetriever = getClaimsRetriever();
         if (claimsRetriever != null) {
+            Map<ClaimMapping, String> customClaimsWithMapping;
+            Map<String, String> customClaims;
             //fix for https://github.com/wso2/product-apim/issues/4112
             String accessToken = validationContext.getAccessToken();
-            AuthorizationGrantCacheKey cacheKey = new AuthorizationGrantCacheKey(accessToken);
+            String authCode = validationContext.getAuthorizationCode();
+            if (accessToken != null) {
+                AuthorizationGrantCacheEntry cacheEntry = AuthorizationGrantCache.getInstance()
+                        .getValueFromCacheByToken(new AuthorizationGrantCacheKey(accessToken));
+                if (cacheEntry == null) {
+                    return new HashMap<>();
+                }
+                customClaimsWithMapping = cacheEntry.getUserAttributes();
+            } else if (authCode != null) {
+                AuthorizationGrantCacheEntry cacheEntry = AuthorizationGrantCache.getInstance()
+                        .getValueFromCacheByCode(new AuthorizationGrantCacheKey(authCode));
+                if (cacheEntry == null) {
+                    return new HashMap<>();
+                }
+                customClaimsWithMapping = cacheEntry.getUserAttributes();
+            } else {
+                customClaimsWithMapping = validationContext.getUser().getUserAttributes();
+            }
+
+            customClaims = convertClaimMap(customClaimsWithMapping);
 
-            Map<String, String> customClaims = getClaimsFromCache(cacheKey);
             if (isNotEmpty(customClaims)) {
                 if (log.isDebugEnabled()) {
                     log.debug("The custom claims are retrieved from AuthorizationGrantCache for user : " +
@@ -163,16 +184,12 @@ public Map<String, String> populateCustomClaims(TokenValidationContext validatio
         }
         return null;
     }
-    protected Map<String, String> getClaimsFromCache(AuthorizationGrantCacheKey cacheKey) {
 
-        AuthorizationGrantCacheEntry cacheEntry = AuthorizationGrantCache.getInstance().getValueFromCacheByToken(cacheKey);
-        if (cacheEntry == null) {
-            return new HashMap<String, String>();
-        }
-        Map<ClaimMapping, String> userAttributes = cacheEntry.getUserAttributes();
-        Map<String, String> userClaims = new HashMap<String, String>();
+    protected Map<String, String> convertClaimMap(Map<ClaimMapping, String> userAttributes) {
+
+        Map<String, String> userClaims = new HashMap<>();
         for (Map.Entry<ClaimMapping, String> entry : userAttributes.entrySet()) {
-            userClaims.put(entry.getKey().getRemoteClaim().getClaimUri(), entry.getValue());
+            userClaims.put(entry.getKey().getLocalClaim().getClaimUri(), entry.getValue());
         }
         return userClaims;
     }
diff --git a/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/test/java/org/wso2/carbon/apimgt/keymgt/token/TokenGenTest.java b/components/apimgt/org.wso2.carbon.apimgt.keymgt/src/test/java/org/wso2/carbon/apimgt/keymgt/token/TokenGenTest.java
@@ -16,36 +16,36 @@
 
 package org.wso2.carbon.apimgt.keymgt.token;
 
-import java.io.FileInputStream;
-import java.io.InputStream;
-import java.security.KeyStore;
-import java.security.MessageDigest;
-import java.security.cert.Certificate;
-import java.security.cert.X509Certificate;
+import org.apache.axiom.util.base64.Base64Utils;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.junit.Assert;
+import org.junit.Before;
 import org.junit.Ignore;
+import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mockito;
 import org.powermock.api.mockito.PowerMockito;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
-import org.apache.axiom.util.base64.Base64Utils;
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
-import org.junit.Before;
-import org.junit.Test;
 import org.wso2.carbon.apimgt.impl.APIConstants;
 import org.wso2.carbon.apimgt.impl.APIManagerConfiguration;
 import org.wso2.carbon.apimgt.impl.APIManagerConfigurationServiceImpl;
 import org.wso2.carbon.apimgt.impl.dto.APIKeyValidationInfoDTO;
 import org.wso2.carbon.apimgt.impl.internal.ServiceReferenceHolder;
 import org.wso2.carbon.apimgt.impl.utils.APIUtil;
 import org.wso2.carbon.apimgt.keymgt.service.TokenValidationContext;
-import org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheKey;
+import org.wso2.carbon.core.util.KeyStoreManager;
+import org.wso2.carbon.identity.application.common.model.ClaimMapping;
 
+import java.io.FileInputStream;
+import java.io.InputStream;
+import java.security.KeyStore;
+import java.security.MessageDigest;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
 import java.util.HashMap;
 import java.util.Map;
-import org.wso2.carbon.core.util.KeyStoreManager;
 //import org.wso2.carbon.apimgt.impl.utils.TokenGenUtil;
 
 @RunWith(PowerMockRunner.class)
@@ -66,8 +66,8 @@ public void setUp() throws Exception {
     public void testAbstractJWTGenerator() throws Exception {
         JWTGenerator jwtGen = new JWTGenerator() {
             @Override
-            protected Map<String, String> getClaimsFromCache(AuthorizationGrantCacheKey cacheKey) {
-                return new HashMap<String, String>();
+            protected Map<String, String> convertClaimMap(Map<ClaimMapping, String> userAttributes) {
+                return new HashMap<>();
             }
         };
         APIKeyValidationInfoDTO dto=new APIKeyValidationInfoDTO();
@@ -127,8 +127,8 @@ protected Map<String, String> getClaimsFromCache(AuthorizationGrantCacheKey cach
     public void testJWTGeneration() throws Exception {
         JWTGenerator jwtGen = new JWTGenerator() {
             @Override
-            public Map<String, String> getClaimsFromCache(AuthorizationGrantCacheKey cacheKey) {
-                return new HashMap<String, String>();
+            public Map<String, String> convertClaimMap(Map<ClaimMapping, String> userAttributes) {
+                return new HashMap<>();
             }
         };
         APIKeyValidationInfoDTO dto=new APIKeyValidationInfoDTO();
