Disable generating jwt on claim retrieval failure by a config

Author: RusJaI

Diff for: components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/security/jwt/JWTValidator.java

@@ -427,7 +427,7 @@ private String generateAndRetrieveJWTToken(String tokenSignature, JWTInfoDto jwt
         return endUserToken;
     }
 
-    private void includeUserStoreClaimsIntoClaims(JWTInfoDto jwtInfoDto) {
+    private void includeUserStoreClaimsIntoClaims(JWTInfoDto jwtInfoDto) throws JWTGeneratorException {
 
         JWTInfoDto localJWTInfoDto = new JWTInfoDto(jwtInfoDto);
         Map<String, String> userClaimsFromKeyManager = getUserClaimsFromKeyManager(localJWTInfoDto);
@@ -859,7 +859,7 @@ protected Cache getGatewayJWTTokenCache() {
 
         return CacheProvider.getGatewayJWTTokenCache();
     }
-    private Map<String, String> getUserClaimsFromKeyManager(JWTInfoDto jwtInfoDto) {
+    private Map<String, String> getUserClaimsFromKeyManager(JWTInfoDto jwtInfoDto) throws JWTGeneratorException {
 
         if (jwtConfigurationDto.isEnableUserClaimRetrievalFromUserStore()) {
             String tenantDomain = CarbonContext.getThreadLocalCarbonContext().getTenantDomain();
@@ -881,7 +881,11 @@ private Map<String, String> getUserClaimsFromKeyManager(JWTInfoDto jwtInfoDto) {
                     try {
                         return keyManagerInstance.getUserClaims(jwtInfoDto.getEndUser(), properties);
                     } catch (APIManagementException e) {
-                        log.error("Error while retrieving User claims from Key Manager ", e);
+                        if (jwtConfigurationDto.isContinueOnClaimRetrievalFailure()) {
+                            log.error("Error while retrieving User Claims from Key Manager ", e);
+                        } else {
+                            throw new JWTGeneratorException("Error while retrieving User Claims from Key Manager", e);
+                        }
                     }
                 }
             }

Diff for: components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java

@@ -676,6 +676,7 @@ private AI() {
     public static final String VALUE = "value";
     public static final String GATEWAY_INTROSPECT_CACHE_NAME = "GatewayIntrospectCache";
     public static final String ENABLE_USER_CLAIMS_RETRIEVAL_FROM_KEY_MANAGER = "EnableUserClaimRetrievalFromKeyManager";
+    public static final String CONTINUE_ON_CLAIM_RETRIEVAL_FAILURE = "ContinueOnClaimRetrievalFailure";
 
     public static final String DELEM_COLON = ":";
     public static final String DELEM_COMMA = ",";

Diff for: components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIManagerConfiguration.java

@@ -1845,6 +1845,12 @@ private void setJWTConfiguration(OMElement omElement) {
                         }
                     }
                 }
+                OMElement continueOnCustomClaimRetrievalFailureElement =
+                        omElement.getFirstChildWithName(new QName(APIConstants.CONTINUE_ON_CLAIM_RETRIEVAL_FAILURE));
+                if (continueOnCustomClaimRetrievalFailureElement != null) {
+                    jwtConfigurationDto.setContinueOnClaimRetrievalFailure(
+                            Boolean.parseBoolean(continueOnCustomClaimRetrievalFailureElement.getText()));
+                }
                 OMElement enableBase64PaddingElement = gatewayJWTConfigurationElement.getFirstChildWithName(
                         new QName(APIConstants.ENABLE_BASE64_PADDING));
                 if (enableBase64PaddingElement != null) {

Diff for: components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/dto/ExtendedJWTConfigurationDto.java

@@ -6,6 +6,7 @@ public class ExtendedJWTConfigurationDto extends JWTConfigurationDto {
     private String jwtGeneratorImplClass = "org.wso2.carbon.apimgt.keymgt.token.JWTGenerator";
     private String claimRetrieverImplClass;
     private boolean tenantBasedSigningEnabled;
+    private boolean continueOnClaimRetrievalFailure;
     private boolean enableUserClaimRetrievalFromUserStore;
     private boolean isBindFederatedUserClaims;
     private boolean isJWKSApiEnabled;
@@ -50,6 +51,14 @@ public boolean isEnableUserClaimRetrievalFromUserStore() {
         return enableUserClaimRetrievalFromUserStore;
     }
 
+    public boolean isContinueOnClaimRetrievalFailure() {
+        return continueOnClaimRetrievalFailure;
+    }
+
+    public void setContinueOnClaimRetrievalFailure(boolean continueOnClaimRetrievalFailure) {
+        this.continueOnClaimRetrievalFailure = continueOnClaimRetrievalFailure;
+    }
+
     public boolean isBindFederatedUserClaims() {
 
         return isBindFederatedUserClaims;

Diff for: components/apimgt/org.wso2.carbon.apimgt.keymgt/src/main/java/org/wso2/carbon/apimgt/keymgt/handlers/AbstractKeyValidationHandler.java

@@ -133,17 +133,20 @@ protected boolean hasTokenRequiredAuthLevel(String authScheme,
     @Override
     public boolean generateConsumerToken(TokenValidationContext validationContext) throws APIKeyMgtException {
 
-
+        String jwt;
         try {
-            String jwt = getCachedJWTToken(validationContext);
-            validationContext.getValidationInfoDTO().setEndUserToken(jwt);
-            return true;
-
+            jwt = getCachedJWTToken(validationContext);
         } catch (APIManagementException e) {
-            log.error("Error occurred while generating JWT. ", e);
+            if (!(ServiceReferenceHolder.getInstance().getAPIManagerConfigurationService().getAPIManagerConfiguration()
+                    .getJwtConfigurationDto().isContinueOnClaimRetrievalFailure())) {
+                throw new APIKeyMgtException("Error occurred while generating JWT", e);
+            } else {
+                log.error("Error occurred while generating JWT. ", e);
+                return false;
+            }
         }
-
-        return false;
+        validationContext.getValidationInfoDTO().setEndUserToken(jwt);
+        return true;
     }
 
     private String getCachedJWTToken(TokenValidationContext validationContext) throws APIManagementException {

Diff for: features/apimgt/org.wso2.carbon.apimgt.core.feature/src/main/resources/conf_templates/org.wso2.carbon.apimgt.core.default.json

@@ -32,6 +32,7 @@
   "apim.jwt.encode_x5t_without_padding": false,
   "apim.jwt.enable_tenant_based_signing": false,
   "apim.jwt.gateway_generator.enable_claim_retrieval": false,
+  "apim.jwt.continue_on_claim_retrieval_failure": true,
   "apim.jwt.binding_federated_user_claims": false,
   "apim.hashing.hashing_algorithm": "SHA-256",
   "apim.cache.gateway_token.enable": true,

Diff for: features/apimgt/org.wso2.carbon.apimgt.core.feature/src/main/resources/conf_templates/templates/repository/conf/api-manager.xml.j2

@@ -86,6 +86,8 @@
         <!-- This parameter specifies which implementation should be used for generating the Token. For URL safe JWT
              Token generation the implementation is provided in URLSafeJWTGenerator -->
         <!--<JWTGeneratorImpl>org.wso2.carbon.apimgt.keymgt.token.URLSafeJWTGenerator</JWTGeneratorImpl>-->
+        <ContinueOnClaimRetrievalFailure>{{apim.jwt.continue_on_claim_retrieval_failure}}</ContinueOnClaimRetrievalFailure>
+
        {% if apim.jwt.enable_tenant_based_signing is defined %}
         <EnableTenantBasedSigning>{{apim.jwt.enable_tenant_based_signing}}</EnableTenantBasedSigning>
         {% endif %}