Add Component ID to API key identifier

wso2 · Nov 6, 2024 · 6a53576 · 6a53576
1 parent dff646e
commit 6a53576
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 7 deletions.
diff --git a/...rcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyAuthenticator.java b/...rcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyAuthenticator.java
@@ -109,22 +109,27 @@ protected String retrieveTokenFromRequestCtx(RequestContext requestContext) thro
                     APISecurityConstants.API_AUTH_INVALID_CREDENTIALS_MESSAGE);
         }
         String keyHash = APIKeyUtils.generateAPIKeyHash(apiKey);
-        Object cachedJWT = CacheProvider.getGatewayAPIKeyJWTCache().getIfPresent(keyHash);
+        String componentId = requestContext.getMatchedAPI().getChoreoComponentInfo().getComponentID();
+        if (componentId == null) {
+            componentId = "";
+        }
+        String apiKeyId = keyHash + APIKeyConstants.API_KEY_ID_SEPARATOR + componentId;
+        Object cachedJWT = CacheProvider.getGatewayAPIKeyJWTCache().getIfPresent(apiKeyId);
         if (cachedJWT != null && !APIKeyUtils.isJWTExpired((String) cachedJWT)) {
             if (log.isDebugEnabled()) {
                 log.debug("Token retrieved from the cache. Token: " + FilterUtils.getMaskedToken(keyHash));
             }
             return (String) cachedJWT;
         }
         // Exchange the API Key to a JWT token.
-        Optional<String> jwt = APIKeyUtils.exchangeAPIKeyToJWT(keyHash);
+        Optional<String> jwt = APIKeyUtils.exchangeAPIKeyToJWT(apiKeyId);
         if (jwt.isEmpty()) {
             throw new APISecurityException(APIConstants.StatusCodes.UNAUTHENTICATED.getCode(),
                     APISecurityConstants.API_AUTH_INVALID_CREDENTIALS,
                     APISecurityConstants.API_AUTH_INVALID_CREDENTIALS_MESSAGE);
         }
         // Cache the JWT token.
-        CacheProvider.getGatewayAPIKeyJWTCache().put(keyHash, jwt.get());
+        CacheProvider.getGatewayAPIKeyJWTCache().put(apiKeyId, jwt.get());
         return jwt.get();
     }
 

diff --git a/...enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyConstants.java b/...enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyConstants.java
@@ -28,6 +28,8 @@ public class APIKeyConstants {
 
     public static final String API_KEY_JSON_KEY = "key";
 
+    public static final String API_KEY_ID_SEPARATOR = "#";
+
     public static final String PAT_EXCHANGE_ENDPOINT = "/internal/pat";
     public static final String API_KEY_EXCHANGE_ENDPOINT = "/internal/apiKey/token";
 }
diff --git a/...ent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtils.java b/...ent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtils.java
@@ -142,12 +142,12 @@ public static Optional<String> exchangePATToJWT(String patHash) {
     }
 
     /**
-     * Exchange a given API key hash to a JWT token.
+     * Exchange a given API key ID to a JWT token.
      *
-     * @param apiKeyHash    API Key Hash
+     * @param apiKeyId    API Key Hash + "#" + Target component ID.
      * @return JWT corresponding to given API Key.
      */
-    public static Optional<String> exchangeAPIKeyToJWT(String apiKeyHash) {
+    public static Optional<String> exchangeAPIKeyToJWT(String apiKeyId) {
 
         URL url = null;
         try {
@@ -162,7 +162,7 @@ public static Optional<String> exchangeAPIKeyToJWT(String apiKeyHash) {
             // Create a request to exchange API key to JWT.
             HttpPost exchangeRequest = new HttpPost(url.toURI());
             exchangeRequest.addHeader("Content-Type", ContentType.APPLICATION_JSON.toString());
-            exchangeRequest.setEntity(new StringEntity(createKeyHashExchangeRequest(apiKeyHash)));
+            exchangeRequest.setEntity(new StringEntity(createKeyHashExchangeRequest(apiKeyId)));
             try (CloseableHttpResponse response = httpClient.execute(exchangeRequest)) {
                 if (response.getStatusLine().getStatusCode() == 200) {
                     HttpEntity entity = response.getEntity();

diff --git a/.../src/test/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyAuthenticatorTest.java b/.../src/test/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyAuthenticatorTest.java
@@ -30,6 +30,7 @@
 import org.wso2.carbon.apimgt.common.gateway.dto.JWTConfigurationDto;
 import org.wso2.choreo.connect.enforcer.common.CacheProvider;
 import org.wso2.choreo.connect.enforcer.commons.model.APIConfig;
+import org.wso2.choreo.connect.enforcer.commons.model.ChoreoComponentInfo;
 import org.wso2.choreo.connect.enforcer.commons.model.RequestContext;
 import org.wso2.choreo.connect.enforcer.config.ConfigHolder;
 import org.wso2.choreo.connect.enforcer.config.EnforcerConfig;
@@ -68,9 +69,13 @@ public void setup() {
     public void retrieveTokenFromRequestCtxTest_invalidKey() {
 
         RequestContext.Builder requestContextBuilder = new RequestContext.Builder("/api-key");
+        ChoreoComponentInfo choreoComponentInfo = new ChoreoComponentInfo();
+        choreoComponentInfo.setComponentID("component_id");
         requestContextBuilder.matchedAPI(new APIConfig.Builder("Petstore")
                 .basePath("/test")
+                .uuid("6003a3b7-af0f-4fb3-853e-a6562b2345f2")
                 .apiType("REST")
+                .choreoComponentInfo(choreoComponentInfo)
                 .build());
         Map<String, String> headersMap = new HashMap<>();
         headersMap.put("choreo-api-key",
@@ -97,9 +102,13 @@ public void retrieveTokenFromRequestCtxTest_cached_validKey() throws APISecurity
         PowerMockito.when(CacheProvider.getGatewayAPIKeyJWTCache()).thenReturn(gatewayAPIKeyJWTCache);
         PowerMockito.when(gatewayAPIKeyJWTCache.getIfPresent(Mockito.anyString())).thenReturn(mockJWT);
 
+        ChoreoComponentInfo choreoComponentInfo = new ChoreoComponentInfo();
+        choreoComponentInfo.setComponentID("component_id");
         RequestContext.Builder requestContextBuilder = new RequestContext.Builder("/api-key");
         requestContextBuilder.matchedAPI(new APIConfig.Builder("Petstore")
                 .basePath("/test")
+                .uuid("6003a3b7-af0f-4fb3-853e-a6562b2345f2")
+                .choreoComponentInfo(choreoComponentInfo)
                 .apiType("REST")
                 .build());
         Map<String, String> headersMap = new HashMap<>();
@@ -128,10 +137,14 @@ public void retrieveTokenFromRequestCtxTest_validKey() throws APISecurityExcepti
         PowerMockito.when(CacheProvider.getGatewayAPIKeyJWTCache()).thenReturn(gatewayAPIKeyJWTCache);
         PowerMockito.when(gatewayAPIKeyJWTCache.getIfPresent(Mockito.anyString())).thenReturn(null);
 
+        ChoreoComponentInfo choreoComponentInfo = new ChoreoComponentInfo();
+        choreoComponentInfo.setComponentID("component_id");
         RequestContext.Builder requestContextBuilder = new RequestContext.Builder("/api-key");
         requestContextBuilder.matchedAPI(new APIConfig.Builder("Petstore")
                 .basePath("/test")
+                .uuid("6003a3b7-af0f-4fb3-853e-a6562b2345f2")
                 .apiType("REST")
+                .choreoComponentInfo(choreoComponentInfo)
                 .build());
         Map<String, String> headersMap = new HashMap<>();
         headersMap.put("choreo-api-key",