Add enable/disable deploying in-memory JWKS Api and Fix block subscription for production only bug

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/InMemoryAPIDeployer.java

@@ -46,6 +46,7 @@
 import org.wso2.carbon.apimgt.gateway.internal.ServiceReferenceHolder;
 import org.wso2.carbon.apimgt.gateway.service.APIGatewayAdmin;
 import org.wso2.carbon.apimgt.impl.APIConstants;
+import org.wso2.carbon.apimgt.impl.dto.ExtendedJWTConfigurationDto;
 import org.wso2.carbon.apimgt.impl.dto.GatewayArtifactSynchronizerProperties;
 import org.wso2.carbon.apimgt.impl.dto.GatewayCleanupSkipList;
 import org.wso2.carbon.apimgt.impl.gatewayartifactsynchronizer.ArtifactRetriever;
@@ -210,7 +211,11 @@ public boolean deployAllAPIs(Set<String> assignedGatewayLabels, String tenantDom
 
         if (!redeployChangedAPIs) {
             try {
-                deployJWKSSynapseAPI(tenantDomain); // Deploy JWKS API
+                boolean isJWKSApiEnabled =  ServiceReferenceHolder
+                        .getInstance().getAPIManagerConfiguration().getJwtConfigurationDto().isJWKSApiEnabled();
+                if(isJWKSApiEnabled) {
+                    deployJWKSSynapseAPI(tenantDomain); // Deploy JWKS API
+                }
                 if (APIConstants.SUPER_TENANT_DOMAIN.equalsIgnoreCase(tenantDomain)) {
                     deployHealthCheckSynapseAPI(tenantDomain); // Deploy HealthCheck API for the super tenant
                 }

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/DefaultAPIHandler.java

@@ -29,6 +29,7 @@
 import org.wso2.carbon.apimgt.common.gateway.constants.HealthCheckConstants;
 import org.wso2.carbon.apimgt.common.gateway.constants.JWTConstants;
 import org.wso2.carbon.apimgt.gateway.InMemoryAPIDeployer;
+import org.wso2.carbon.apimgt.gateway.internal.ServiceReferenceHolder;
 import org.wso2.carbon.apimgt.gateway.utils.GatewayUtils;
 import org.wso2.carbon.apimgt.impl.APIConstants;
 import org.wso2.carbon.apimgt.impl.gatewayartifactsynchronizer.exception.ArtifactSynchronizerException;
@@ -62,7 +63,10 @@ public boolean handleRequestInFlow(MessageContext messageContext) {
             }
         }
 
-        if (isJWKSEndpoint) {
+        boolean isJWKSApiEnabled =  ServiceReferenceHolder
+                .getInstance().getAPIManagerConfiguration().getJwtConfigurationDto().isJWKSApiEnabled();
+
+        if (isJWKSEndpoint && isJWKSApiEnabled) {
             try {
                 InMemoryAPIDeployer.deployJWKSSynapseAPI(tenantDomain);
             } catch(APIManagementException e){

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/security/APIKeyValidator.java

@@ -743,9 +743,9 @@
     }
 
     public APIKeyValidationInfoDTO validateSubscription(String context, String version, int appID,
-                                                        String tenantDomain)
+                                                        String tenantDomain, String keyType)
             throws APISecurityException {
-        return dataStore.validateSubscription(context, version, appID,tenantDomain);
+        return dataStore.validateSubscription(context, version, appID,tenantDomain, keyType);
     }
 
     /**

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/security/keys/APIKeyDataStore.java

@@ -107,7 +107,7 @@ APIKeyValidationInfoDTO validateSubscription(String context, String version, Str
      * @return an APIKeyValidationInfoDTO instance containing key validation data
      * @throws org.wso2.carbon.apimgt.gateway.handlers.security.APISecurityException on error
      */
-    APIKeyValidationInfoDTO validateSubscription(String context, String version, int appId, String tenantDomain)
+    APIKeyValidationInfoDTO validateSubscription(String context, String version, int appId, String tenantDomain, String keyType)
             throws APISecurityException;
     /**
      * Validate scopes bound to the resource of the API being invoked against the scopes of the token.

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/security/keys/APIKeyValidatorClient.java

@@ -76,12 +76,12 @@
     }
 
     public APIKeyValidationInfoDTO validateSubscription(String context, String version, int appId,
-                                                        String tenantDomain)
+                                                        String tenantDomain, String keyType)
             throws APISecurityException {
 
         try {
             return apiKeyValidationService
-                    .validateSubscription(context, version, appId, tenantDomain);
+                    .validateSubscription(context, version, appId, tenantDomain, keyType);
         } catch (APIKeyMgtException | APIManagementException e) {
             log.error("Error while  validate subscriptions", e);
             throw new APISecurityException(APISecurityConstants.API_AUTH_GENERAL_ERROR,

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/security/keys/WSAPIKeyDataStore.java

@@ -107,11 +107,11 @@
 
     @Override
     public APIKeyValidationInfoDTO validateSubscription(String context, String version, int appId,
-                                                        String tenantDomain)
+                                                        String tenantDomain, String keyType)
             throws APISecurityException {
         APIKeyValidatorClient client = new APIKeyValidatorClient();
         try {
-            return client.validateSubscription(context, version, appId, tenantDomain);
+            return client.validateSubscription(context, version, appId, tenantDomain, keyType);
         } catch (APISecurityException ex) {
             throw new APISecurityException(ex.getErrorCode(),
                     "Resource forbidden", ex);

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/utils/GatewayUtils.java

@@ -796,6 +796,7 @@
         APIKeyValidator apiKeyValidator = new APIKeyValidator();
         APIKeyValidationInfoDTO apiKeyValidationInfoDTO = null;
         JSONObject application;
+        String keyType = (String) payload.getClaim(APIConstants.JwtTokenConstants.KEY_TYPE);
         int appId = 0;
         if (payload.getClaim(APIConstants.JwtTokenConstants.APPLICATION) != null) {
             try {
@@ -813,7 +814,7 @@
         // if the appId is equal to 0 then it's a internal key
         if (appId != 0) {
             apiKeyValidationInfoDTO =
-                    apiKeyValidator.validateSubscription(apiContext, apiVersion, appId, getTenantDomain());
+                    apiKeyValidator.validateSubscription(apiContext, apiVersion, appId, getTenantDomain(), keyType);
         }
 
         if (payload.getClaim(APIConstants.JwtTokenConstants.SUBSCRIBED_APIS) != null) {
@@ -887,6 +888,7 @@
 
         APIKeyValidator apiKeyValidator = new APIKeyValidator();
         APIKeyValidationInfoDTO apiKeyValidationInfoDTO = null;
+        String keyType = (String) payload.getClaim(APIConstants.JwtTokenConstants.KEY_TYPE);
         int appId = 0;
         if (payload.getClaim(APIConstants.JwtTokenConstants.APPLICATION) != null) {
             try {
@@ -904,13 +906,12 @@
         // if the appId is equal to 0 then it's a internal key
         if (appId != 0) {
             apiKeyValidationInfoDTO =
-                    apiKeyValidator.validateSubscription(apiContext, apiVersion, appId, getTenantDomain());
+                    apiKeyValidator.validateSubscription(apiContext, apiVersion, appId, getTenantDomain(), keyType);
             if (apiKeyValidationInfoDTO.isAuthorized()) {
                 if (log.isDebugEnabled()) {
                     log.debug("User is subscribed to the API: " + apiContext + ", " +
                             "version: " + apiVersion + ". Token: " + getMaskedToken(token));
                 }
-                String keyType = (String) payload.getClaim(APIConstants.JwtTokenConstants.KEY_TYPE);
                 apiKeyValidationInfoDTO.setType(keyType);
             } else {
                 if (log.isDebugEnabled()) {

components/apimgt/org.wso2.carbon.apimgt.gateway/src/test/java/org/wso2/carbon/apimgt/gateway/handlers/DefaultAPIHandlerTest.java

@@ -29,23 +29,39 @@
 import org.powermock.api.mockito.PowerMockito;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
+import org.wso2.carbon.apimgt.gateway.internal.ServiceReferenceHolder;
 import org.wso2.carbon.apimgt.gateway.utils.GatewayUtils;
 import org.wso2.carbon.apimgt.impl.APIConstants;
+import org.wso2.carbon.apimgt.impl.APIManagerConfiguration;
+import org.wso2.carbon.apimgt.impl.dto.ExtendedJWTConfigurationDto;
 import org.wso2.carbon.apimgt.keymgt.model.entity.API;
 import org.wso2.carbon.inbound.endpoint.protocol.websocket.InboundWebsocketConstants;
 
 import java.util.Set;
 import java.util.TreeMap;
 
 @RunWith(PowerMockRunner.class)
-@PrepareForTest({ApiUtils.class, Utils.class, GatewayUtils.class})
+@PrepareForTest({ApiUtils.class, Utils.class, GatewayUtils.class, ServiceReferenceHolder.class, APIManagerConfiguration.class, ExtendedJWTConfigurationDto.class})
 public class DefaultAPIHandlerTest {
 
+    private ServiceReferenceHolder serviceReferenceHolder;
+    private APIManagerConfiguration apiManagerConfiguration;
+    private ExtendedJWTConfigurationDto extendedJWTConfigurationDto;
+
     @Before
     public void init() {
         PowerMockito.mockStatic(ApiUtils.class);
         PowerMockito.mockStatic(Utils.class);
         PowerMockito.mockStatic(GatewayUtils.class);
+        PowerMockito.mockStatic(ServiceReferenceHolder.class);
+        serviceReferenceHolder = Mockito.mock(ServiceReferenceHolder.class);
+        apiManagerConfiguration = Mockito.mock(APIManagerConfiguration.class);
+        extendedJWTConfigurationDto = Mockito.mock(ExtendedJWTConfigurationDto.class);
+        Mockito.when(ServiceReferenceHolder.getInstance()).thenReturn(serviceReferenceHolder);
+        Mockito.when(serviceReferenceHolder.getAPIManagerConfiguration()).thenReturn(apiManagerConfiguration);
+        Mockito.when(apiManagerConfiguration.getJwtConfigurationDto()).thenReturn(extendedJWTConfigurationDto);
+        Mockito.when(extendedJWTConfigurationDto.isJWKSApiEnabled()).thenReturn(true);
+
     }
 
     @Test

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

@@ -464,6 +464,7 @@ public final class APIConstants {
     public static final String BINDING_FEDERATED_USER_CLAIMS = "EnableBindingFederatedUserClaims";
     public static final String TOKEN_GENERATOR_IMPL = "JWTGeneratorImpl";
     public static final String ENABLE_JWT_GENERATION = "EnableJWTGeneration";
+    public static final String Enable_JWKS_API = "EnableJWKSApi";
     public static final String CLAIMS_RETRIEVER_CLASS = "ClaimsRetrieverImplClass";
     public static final String USE_KID = "UseKidProperty";
     public static final String CONSUMER_DIALECT_URI = "ConsumerDialectURI";

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

@@ -1787,6 +1787,10 @@ private void setJWTConfiguration(OMElement omElement) {
                 }
             }
         }
+
+        OMElement jwksApiEnableElement =
+                omElement.getFirstChildWithName(new QName(APIConstants.Enable_JWKS_API));
+        jwtConfigurationDto.setJWKSApiEnabled(Boolean.parseBoolean(jwksApiEnableElement.getText()));
     }
 
     public ThrottleProperties getThrottleProperties() {

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

@@ -8,6 +8,7 @@ public class ExtendedJWTConfigurationDto extends JWTConfigurationDto {
     private boolean tenantBasedSigningEnabled;
     private boolean enableUserClaimRetrievalFromUserStore;
     private boolean isBindFederatedUserClaims;
+    private boolean isJWKSApiEnabled;
 
     public String getClaimRetrieverImplClass() {
 
@@ -58,4 +59,12 @@ public void setBindFederatedUserClaims(boolean isBindFederatedUserClaims) {
 
         this.isBindFederatedUserClaims = isBindFederatedUserClaims;
     }
+
+    public boolean isJWKSApiEnabled() {
+        return isJWKSApiEnabled;
+    }
+
+    public void setJWKSApiEnabled(boolean JWKSApiEnabled) {
+        this.isJWKSApiEnabled = JWKSApiEnabled;
+    }
 }

components/apimgt/org.wso2.carbon.apimgt.impl/src/test/java/org/wso2/carbon/apimgt/impl/dao/test/APIMgtDAOTest.java

@@ -1818,12 +1818,13 @@ public void testRetrieveAllWorkflowFromInternalReference() throws Exception {
     }
 
     /**
-     * Test for getAPIRevisionDeploymentsByWorkflowStatusAndApiUUID method
-     * Checks whether the API revision deployment mapping details are retrieved correctly
+     * Test for testGetAndUpdateAPIRevisionDeploymentsByWorkflowStatusAndApiUUID method
+     * Checks whether the API revision deployment mapping details are retrieved correctly and
+     * Checks whether the API revision deployment status is updated correctly
      * @throws APIManagementException if an error occurs while retrieving revision deployment mapping details
      */
     @Test
-    public void testGetAPIRevisionDeploymentsByWorkflowStatusAndApiUUID() throws Exception {
+    public void testGetAndUpdateAPIRevisionDeploymentsByWorkflowStatusAndApiUUID() throws Exception {
         String workflowStatus = "CREATED";
         String apiUUID = "7af95c9d-6177-4191-ab3e-d3f6c1cdc4c2";
         String revisionUUID = "821b9664-eeca-4173-9f56-3dc6d46bd6eb";
@@ -1836,24 +1837,14 @@ public void testGetAPIRevisionDeploymentsByWorkflowStatusAndApiUUID() throws Exc
         Assert.assertNotNull(apiRevisionDeployment);
         Assert.assertEquals(apiRevisionDeployment.getDeployment(), deployment);
         Assert.assertEquals(apiRevisionDeployment.getRevisionUUID(), revisionUUID);
-    }
 
-    /**
-     * Test for updateAPIRevisionDeploymentStatus method
-     * Checks whether the API revision deployment status is updated correctly
-     * @throws APIManagementException if an error occurs while updating revision deployment status
-     */
-    @Test public void testUpdateAPIRevisionDeploymentStatus() throws Exception {
-        String workflowStatus = "APPROVED";
-        String revisionUUID = "821b9664-eeca-4173-9f56-3dc6d46bd6eb";
-        String apiId = "7af95c9d-6177-4191-ab3e-d3f6c1cdc4c2";
-        String deployment = "default";
-        apiMgtDAO.updateAPIRevisionDeploymentStatus(revisionUUID, workflowStatus, deployment);
-        List<APIRevisionDeployment> apiRevisionDeployments = apiMgtDAO.getAPIRevisionDeploymentByApiUUID(apiId);
-        Assert.assertNotNull(apiRevisionDeployments);
-        APIRevisionDeployment apiRevisionDeployment = apiRevisionDeployments.get(0);
-        Assert.assertNotNull(apiRevisionDeployment);
-        Assert.assertEquals(org.wso2.carbon.apimgt.api.WorkflowStatus.APPROVED,apiRevisionDeployment.getStatus());
+        String workflowStatus2 = "APPROVED";
+        apiMgtDAO.updateAPIRevisionDeploymentStatus(revisionUUID, workflowStatus2, deployment);
+        List<APIRevisionDeployment> apiRevisionDeployments2 = apiMgtDAO.getAPIRevisionDeploymentByApiUUID(apiUUID);
+        Assert.assertNotNull(apiRevisionDeployments2);
+        APIRevisionDeployment apiRevisionDeployment2 = apiRevisionDeployments2.get(0);
+        Assert.assertNotNull(apiRevisionDeployment2);
+        Assert.assertEquals(org.wso2.carbon.apimgt.api.WorkflowStatus.APPROVED,apiRevisionDeployment2.getStatus());
     }
 
     @Test

components/apimgt/org.wso2.carbon.apimgt.impl/src/test/resources/amConfig.xml

@@ -11,6 +11,7 @@
     </Database>
     <JWTConfiguration>
         <EnableJWTGeneration>false</EnableJWTGeneration>
+        <EnableJWKSApi>true</EnableJWKSApi>
 	    <SignatureAlgorithm>NONE</SignatureAlgorithm>
     </JWTConfiguration>
     <APIUsageTracking>

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

@@ -200,15 +200,15 @@
     }
 
     @Override
-    public APIKeyValidationInfoDTO validateSubscription(String apiContext, String apiVersion, int appId) {
+    public APIKeyValidationInfoDTO validateSubscription(String apiContext, String apiVersion, int appId, String keyType) {
         APIKeyValidationInfoDTO apiKeyValidationInfoDTO =  new APIKeyValidationInfoDTO();
 
         if (log.isDebugEnabled()) {
             log.debug("Before validating subscriptions");
             log.debug("Validation Info : { context : " + apiContext + " , " + "version : "
                     + apiVersion + " , appId : " + appId + " }");
         }
-        validateSubscriptionDetails(apiContext, apiVersion, appId, apiKeyValidationInfoDTO);
+        validateSubscriptionDetails(apiContext, apiVersion, appId, apiKeyValidationInfoDTO, keyType);
         if (log.isDebugEnabled()) {
             log.debug("After validating subscriptions");
         }
@@ -230,15 +230,15 @@
 
 
     private boolean validateSubscriptionDetails(String context, String version, int appId,
-                                                APIKeyValidationInfoDTO infoDTO) {
+                                                APIKeyValidationInfoDTO infoDTO, String keyType) {
 
         // Check if the api version has been prefixed with _default_
         if (version != null && version.startsWith(APIConstants.DEFAULT_VERSION_PREFIX)) {
             // Remove the prefix from the version.
             version = version.split(APIConstants.DEFAULT_VERSION_PREFIX)[1];
         }
 
-        validateSubscriptionDetails(infoDTO, context, version, appId);
+        validateSubscriptionDetails(infoDTO, context, version, appId, keyType);
         return infoDTO.isAuthorized();
     }
 
@@ -326,7 +326,7 @@
     }
 
     private APIKeyValidationInfoDTO validateSubscriptionDetails(APIKeyValidationInfoDTO infoDTO, String context,
-                                                                String version, int appId) {
+                                                                String version, int appId, String keyType) {
         String apiTenantDomain = MultitenantUtils.getTenantDomainFromRequestURL(context);
         if (apiTenantDomain == null) {
             apiTenantDomain = MultitenantConstants.SUPER_TENANT_DOMAIN_NAME;
@@ -367,7 +367,7 @@
         }
 
         if (api != null && sub != null) {
-            validate(infoDTO, apiTenantDomain, tenantId, datastore, api, app, sub);
+            validate(infoDTO, apiTenantDomain, tenantId, datastore, api, app, sub, keyType);
         } else if (!infoDTO.isAuthorized() && infoDTO.getValidationStatus() == 0) {
             //Scenario where validation failed and message is not set
             infoDTO.setValidationStatus(APIConstants.KeyValidationStatus.API_AUTH_RESOURCE_FORBIDDEN);
@@ -657,9 +657,8 @@
 
 
     private APIKeyValidationInfoDTO validate(APIKeyValidationInfoDTO infoDTO, String apiTenantDomain, int tenantId,
-                                             SubscriptionDataStore datastore, API api, Application app, Subscription sub) {
+                                             SubscriptionDataStore datastore, API api, Application app, Subscription sub, String keyType) {
         String subscriptionStatus = sub.getSubscriptionState();
-        String type = app.getTokenType();
         if (APIConstants.SubscriptionStatus.BLOCKED.equals(subscriptionStatus)) {
             infoDTO.setValidationStatus(APIConstants.KeyValidationStatus.API_BLOCKED);
             infoDTO.setAuthorized(false);
@@ -670,9 +669,9 @@
             infoDTO.setAuthorized(false);
             return infoDTO;
         } else if (APIConstants.SubscriptionStatus.PROD_ONLY_BLOCKED.equals(subscriptionStatus)
-                && !APIConstants.API_KEY_TYPE_SANDBOX.equals(type)) {
+                && !APIConstants.API_KEY_TYPE_SANDBOX.equals(keyType)) {
             infoDTO.setValidationStatus(APIConstants.KeyValidationStatus.API_BLOCKED);
-            infoDTO.setType(type);
+            infoDTO.setType(keyType);
             infoDTO.setAuthorized(false);
             return infoDTO;
         }
@@ -687,7 +686,7 @@
         infoDTO.setApplicationUUID(app.getUUID());
         infoDTO.setApplicationGroupIds(app.getGroupIds().stream().map(GroupId::getGroupId).collect(Collectors.toSet()));
         infoDTO.setAppAttributes(app.getAttributes());
-        infoDTO.setType(type);
+        infoDTO.setType(keyType);
 
         // Advanced Level Throttling Related Properties
         String apiTier = api.getApiTier();