Address review comments

wso2 · Oct 29, 2024 · bc49a1d · bc49a1d
1 parent 061b353
commit bc49a1d
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 17 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
@@ -51,7 +51,8 @@ public APIKeyAuthenticator() {
     public boolean canAuthenticate(RequestContext requestContext) {
 
         String apiKeyValue = getAPIKeyFromRequest(requestContext);
-        return apiKeyValue != null && apiKeyValue.startsWith(APIKeyConstants.API_KEY_PREFIX);
+        return apiKeyValue != null && apiKeyValue.startsWith(APIKeyConstants.API_KEY_PREFIX) &&
+                apiKeyValue.length() > 10;
     }
 
     @Override
@@ -72,9 +73,6 @@ private void dropAPIKeyDataFromAPIKeyHeader(RequestContext requestContext) throw
         // Update the header with the new API key data.
         String encodedKeyData = Base64.getEncoder().encodeToString(jsonObject.toJSONString().getBytes());
         String newAPIKeyHeaderValue = APIKeyConstants.API_KEY_PREFIX + encodedKeyData + checksum;
-        // Remove the existing header.
-        requestContext.getRemoveHeaders().add(ConfigHolder.getInstance().getConfig().getApiKeyConfig()
-                .getApiKeyInternalHeader().toLowerCase());
         // Add the new header.
         requestContext.addOrModifyHeaders(ConfigHolder.getInstance().getConfig().getApiKeyConfig()
                 .getApiKeyInternalHeader().toLowerCase(), newAPIKeyHeaderValue);
@@ -104,7 +102,7 @@ private JSONObject getDecodedAPIKeyData(String apiKeyHeaderValue) throws APISecu
     @Override
     protected String retrieveTokenFromRequestCtx(RequestContext requestContext) throws APISecurityException {
 
-        String apiKey = getAPIKeyFromRequest(requestContext);
+        String apiKey = getAPIKeyFromRequest(requestContext).trim();
         if (!APIKeyUtils.isValidAPIKey(apiKey)) {
             throw new APISecurityException(APIConstants.StatusCodes.UNAUTHENTICATED.getCode(),
                     APISecurityConstants.API_AUTH_INVALID_CREDENTIALS,

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
@@ -35,6 +35,9 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.wso2.choreo.connect.enforcer.config.ConfigHolder;
+import org.wso2.choreo.connect.enforcer.constants.APIConstants;
+import org.wso2.choreo.connect.enforcer.constants.APISecurityConstants;
+import org.wso2.choreo.connect.enforcer.exception.APISecurityException;
 import org.wso2.choreo.connect.enforcer.util.FilterUtils;
 
 import java.io.InputStream;
@@ -80,18 +83,24 @@ public static boolean isValidAPIKey(String apiKey) {
      * @param apiKey    API Key
      * @return key hash
      */
-    public static String generateAPIKeyHash(String apiKey) {
+    public static String generateAPIKeyHash(String apiKey) throws APISecurityException {
 
-        // Skipping the prefix(`chp_`) and checksum.
-        String keyData = apiKey.substring(4, apiKey.length() - 6);
-        // Base 64 decode key data.
-        String decodedKeyData = new String(Base64.getDecoder().decode(keyData));
-        // Convert data into JSON.
-        JSONObject jsonObject = (JSONObject) JSONValue.parse(decodedKeyData);
-        // Extracting the key.
-        String key = jsonObject.getAsString(APIKeyConstants.API_KEY_JSON_KEY);
-        // Return SHA256 hash of the key.
-        return DigestUtils.sha256Hex(key);
+        try {
+            // Skipping the prefix(`chp_`) and checksum.
+            String keyData = apiKey.substring(4, apiKey.length() - 6);
+            // Base 64 decode key data.
+            String decodedKeyData = new String(Base64.getDecoder().decode(keyData));
+            // Convert data into JSON.
+            JSONObject jsonObject = (JSONObject) JSONValue.parse(decodedKeyData);
+            // Extracting the key.
+            String key = jsonObject.getAsString(APIKeyConstants.API_KEY_JSON_KEY);
+            // Return SHA256 hash of the key.
+            return DigestUtils.sha256Hex(key);
+        } catch (Exception e) {
+            throw new APISecurityException(APIConstants.StatusCodes.UNAUTHENTICATED.getCode(),
+                    APISecurityConstants.API_AUTH_INVALID_CREDENTIALS,
+                    APISecurityConstants.API_AUTH_INVALID_CREDENTIALS_MESSAGE);
+        }
     }
 
     /**

diff --git a/...enforcer/src/test/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtilsTest.java b/...enforcer/src/test/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtilsTest.java
@@ -24,6 +24,7 @@
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.wso2.choreo.connect.enforcer.config.ConfigHolder;
+import org.wso2.choreo.connect.enforcer.exception.APISecurityException;
 
 @RunWith(PowerMockRunner.class)
 @PrepareForTest({ConfigHolder.class})
@@ -45,7 +46,7 @@ public void testIsValidAPIKey_invalid() {
     }
 
     @Test
-    public void testGenerateAPIKeyHash() {
+    public void testGenerateAPIKeyHash() throws APISecurityException {
         String apiKey = "chp_eyJrZXkiOiJlanp6am8yaGc5MnA2MTF6NTI2OXMzNzU1ZnJzbnFlNm9vb2hldWd0djBjbmQ3bXdobCJ9dknDJA";
         String expectedKeyHash = "62f73948188c9f773414d4ec77eae6e8caab21556e4ad18f94b7c6c5b018524c";
         String generatedAPIKeyHash = APIKeyUtils.generateAPIKeyHash(apiKey);