Merge pull request #3520 from renuka-fernando/lahiru-feature-flag

Retrieve subscription policies from SubscriptionDataStore for organizations the feature is enabled

enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/features/FeatureFlags.java

@@ -0,0 +1,51 @@
+/*
+ *  Copyright (c) 2024, WSO2 LLC. (http://www.wso2.org) All Rights Reserved.
+ *
+ *  WSO2 LLC. licenses this file to you under the Apache License,
+ *  Version 2.0 (the "License"); you may not use this file except
+ *  in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.choreo.connect.enforcer.features;
+
+import org.wso2.choreo.connect.enforcer.constants.APIConstants;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Feature flags for Choreo Connect.
+ */
+public class FeatureFlags {
+    private static final Set<String> CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ORG;
+    private static final boolean ENABLE_CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ALL_ORG;
+
+    static {
+        final String orgEnvVar = System.getenv().getOrDefault("CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ORG", "");
+        CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ORG = new HashSet<>(Arrays.asList(orgEnvVar.split(",")));
+        ENABLE_CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ALL_ORG = orgEnvVar.equals("*");
+    }
+
+    public static boolean isCustomSubscriptionPolicyHandlingEnabled(String orgId) {
+        return ENABLE_CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ALL_ORG
+                || CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ORG.contains(orgId);
+    }
+
+    public static String getCustomSubscriptionPolicyHandlingOrg(String orgId) {
+        if (isCustomSubscriptionPolicyHandlingEnabled(orgId)) {
+            return orgId;
+        }
+        return APIConstants.SUPER_TENANT_DOMAIN_NAME;
+    }
+}

enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/JWTAuthenticator.java

@@ -45,6 +45,7 @@
 import org.wso2.choreo.connect.enforcer.dto.APIKeyValidationInfoDTO;
 import org.wso2.choreo.connect.enforcer.exception.APISecurityException;
 import org.wso2.choreo.connect.enforcer.exception.EnforcerException;
+import org.wso2.choreo.connect.enforcer.features.FeatureFlags;
 import org.wso2.choreo.connect.enforcer.keymgt.KeyManagerHolder;
 import org.wso2.choreo.connect.enforcer.models.SubscriptionPolicy;
 import org.wso2.choreo.connect.enforcer.security.Authenticator;
@@ -72,8 +73,6 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 /**
  * Implements the authenticator interface to authenticate request using a JWT token.
@@ -85,18 +84,12 @@ public class JWTAuthenticator implements Authenticator {
     private final boolean isGatewayTokenCacheEnabled;
     private AbstractAPIMgtGatewayJWTGenerator jwtGenerator;
     private static final Set<String> prodTokenNonProdAllowedOrgs = new HashSet<>();
-    private static final String orgList = System.getenv("CUSTOM_SUBSCRIPTION_POLICY_HANDLING_ORG");
-    private static Set<String> orgSet = new HashSet<>();
 
     static {
         if (System.getenv("PROD_TOKEN_NONPROD_ALLOWED_ORGS") != null) {
             Collections.addAll(prodTokenNonProdAllowedOrgs,
                     System.getenv("PROD_TOKEN_NONPROD_ALLOWED_ORGS").split("\\s+"));
         }
-        if (orgList != null) {
-            orgSet = Stream.of(orgList.trim().split("\\s*,\\s*"))
-                    .collect(Collectors.toSet());
-        }
     }
 
     public JWTAuthenticator() {
@@ -340,9 +333,8 @@ public AuthenticationContext authenticate(RequestContext requestContext) throws
                                 != null) {
                             SubscriptionPolicy subPolicy = datastore.getSubscriptionPolicyByOrgIdAndName
                                     (matchedApiOrganizationId, subPolicyName);
-                            String metaDataOrgId = StringUtils.isNotEmpty(orgList) &&
-                                    (orgSet.contains(subPolicy.getOrganization()) || orgList.equals("*")) ?
-                                    subPolicy.getOrganization() : APIConstants.SUPER_TENANT_DOMAIN_NAME;
+                            String metaDataOrgId =
+                                    FeatureFlags.getCustomSubscriptionPolicyHandlingOrg(subPolicy.getOrganization());
                             log.debug("Subscription rate-limiting will be evaluated for the organization: " +
                                     metaDataOrgId);
                             requestContext.addMetadataToMap("ratelimit:organization", metaDataOrgId);

enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/subscription/SubscriptionDataStoreImpl.java

@@ -22,12 +22,12 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.wso2.choreo.connect.discovery.subscription.APIs;
-import org.wso2.choreo.connect.enforcer.constants.APIConstants;
 import org.wso2.choreo.connect.enforcer.discovery.ApplicationDiscoveryClient;
 import org.wso2.choreo.connect.enforcer.discovery.ApplicationKeyMappingDiscoveryClient;
 import org.wso2.choreo.connect.enforcer.discovery.ApplicationPolicyDiscoveryClient;
 import org.wso2.choreo.connect.enforcer.discovery.SubscriptionDiscoveryClient;
 import org.wso2.choreo.connect.enforcer.discovery.SubscriptionPolicyDiscoveryClient;
+import org.wso2.choreo.connect.enforcer.features.FeatureFlags;
 import org.wso2.choreo.connect.enforcer.models.API;
 import org.wso2.choreo.connect.enforcer.models.ApiPolicy;
 import org.wso2.choreo.connect.enforcer.models.Application;
@@ -111,7 +111,7 @@ public API getApiByContextAndVersion(String uuid) {
 
     @Override
     public SubscriptionPolicy getSubscriptionPolicyByOrgIdAndName(String orgId, String policyName) {
-        String organizationId = StringUtils.isEmpty(orgId) ? APIConstants.SUPER_TENANT_DOMAIN_NAME : orgId;
+        String organizationId = FeatureFlags.getCustomSubscriptionPolicyHandlingOrg(orgId);
         String key = PolicyType.SUBSCRIPTION +
                 SubscriptionDataStoreUtil.DELEM_PERIOD + organizationId +
                 SubscriptionDataStoreUtil.DELEM_PERIOD + policyName;