Add changes with tests

wso2 · Aug 6, 2024 · e74c1c1 · e74c1c1
1 parent f9b6ebc
commit e74c1c1
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 57 deletions.
diff --git a/.azure/product-microgateway-sca-scan.yaml b/.azure/product-microgateway-sca-scan.yaml
@@ -26,9 +26,6 @@ pr:
     - 'choreo'
     - 'main'
 
-pool:
-  vmImage: 'ubuntu-latest'
-
 resources:
   repositories:
   - repository: templates

diff --git a/.../main/java/org/wso2/choreo/connect/enforcer/security/jwt/InternalAPIKeyAuthenticator.java b/.../main/java/org/wso2/choreo/connect/enforcer/security/jwt/InternalAPIKeyAuthenticator.java
@@ -78,20 +78,12 @@ public boolean canAuthenticate(RequestContext requestContext) {
         String internalKey = requestContext.getHeaders().get(
                 ConfigHolder.getInstance().getConfig().getAuthHeader().getTestConsoleHeaderName().toLowerCase());
         if (apiType.equalsIgnoreCase("WS") && internalKey == null) {
-            String[] secProtocolHeaderValues = requestContext.getHeaders().get(
-                    HttpConstants.WEBSOCKET_PROTOCOL_HEADER).split(",");
-            if (internalKey == null || internalKey.isBlank()) {
-                if (secProtocolHeaderValues.length > 1 && secProtocolHeaderValues[0].equals(
-                    Constants.WS_API_KEY_IDENTIFIER)) {
-                    internalKey = secProtocolHeaderValues[1];
-                }
-            }
+            internalKey = extractInternalKeyInWSProtocolHeader(requestContext);
 
-            if (secProtocolHeaderValues[0].equals(Constants.WS_API_KEY_IDENTIFIER) &&
-                    secProtocolHeaderValues.length == 2) {
+            if (isAdditionalResponseHeadersRequired(requestContext)) {
                 HashMap<String, String> responseHeadersToAddMap = new HashMap<>();
                 responseHeadersToAddMap.put(
-                        HttpConstants.WEBSOCKET_PROTOCOL_HEADER, secProtocolHeaderValues[0]);
+                        HttpConstants.WEBSOCKET_PROTOCOL_HEADER, Constants.WS_API_KEY_IDENTIFIER);
                 requestContext.setResponseHeadersToAddMap(responseHeadersToAddMap);
             }
         }
@@ -312,18 +304,10 @@ private String extractInternalKey(RequestContext requestContext) {
             return internalKey.trim();
         }
         if (requestContext.getMatchedAPI().getApiType().equalsIgnoreCase("WS")) {
-            String[] secProtocolHeaderValues = requestContext.getHeaders().get(
-                HttpConstants.WEBSOCKET_PROTOCOL_HEADER).split(",");
-            if (secProtocolHeaderValues.length > 1 && secProtocolHeaderValues[0].equals(
-                Constants.WS_API_KEY_IDENTIFIER)) {
-                internalKey = secProtocolHeaderValues[1];
-            }
+            internalKey = extractInternalKeyInWSProtocolHeader(requestContext);
             if (internalKey != null) {
-                if (secProtocolHeaderValues.length > 2) {
-                    String protocols = Arrays.stream(
-                        secProtocolHeaderValues, 2,
-                        secProtocolHeaderValues.length)
-                            .collect(Collectors.joining(","));
+                String protocols = getProtocolsToSetInRequestHeaders(requestContext);
+                if (protocols != null) {
                     requestContext.addOrModifyHeaders(HttpConstants.WEBSOCKET_PROTOCOL_HEADER, protocols);
                 }
                 return internalKey.trim();
@@ -332,6 +316,36 @@ private String extractInternalKey(RequestContext requestContext) {
         return null;
     }
 
+    public String extractInternalKeyInWSProtocolHeader(RequestContext requestContext) {
+        String protocolHeader = requestContext.getHeaders().get(
+                HttpConstants.WEBSOCKET_PROTOCOL_HEADER);
+        if (protocolHeader != null) {
+            String[] secProtocolHeaderValues = protocolHeader.split(",");
+            if (secProtocolHeaderValues.length > 1 && secProtocolHeaderValues[0].equals(
+                    Constants.WS_API_KEY_IDENTIFIER)) {
+                return secProtocolHeaderValues[1].trim();
+            }
+        }
+        return null;
+    }
+
+    public String getProtocolsToSetInRequestHeaders(RequestContext requestContext) {
+        String[] secProtocolHeaderValues = requestContext.getHeaders().get(
+            HttpConstants.WEBSOCKET_PROTOCOL_HEADER).split(",");
+        if (secProtocolHeaderValues.length > 2) {
+            return Arrays.stream(secProtocolHeaderValues, 2, secProtocolHeaderValues.length)
+                    .collect(Collectors.joining(",")).trim();
+        }
+        return null;
+    }
+
+    public boolean isAdditionalResponseHeadersRequired(RequestContext requestContext) {
+        String[] secProtocolHeaderValues = requestContext.getHeaders().get(
+            HttpConstants.WEBSOCKET_PROTOCOL_HEADER).split(",");
+        return secProtocolHeaderValues[0].equals(Constants.WS_API_KEY_IDENTIFIER) &&
+                    secProtocolHeaderValues.length == 2;
+    }
+
     @Override
     public int getPriority() {
         return -10;

diff --git a/...jwt/InternalAPIKeyAuthenticatiorTest.java → .../jwt/InternalAPIKeyAuthenticatorTest.java b/...jwt/InternalAPIKeyAuthenticatiorTest.java → .../jwt/InternalAPIKeyAuthenticatorTest.java
@@ -18,10 +18,8 @@
 
 package org.wso2.choreo.connect.enforcer.security.jwt;
 
-import java.lang.reflect.InvocationTargetException;
 import java.util.HashMap;
 import java.util.Map;
-import java.lang.reflect.Method;
 
 import org.junit.Assert;
 import org.junit.Test;
@@ -40,10 +38,10 @@
 @RunWith(PowerMockRunner.class)
 @PrepareForTest({ConfigHolder.class})
 @PowerMockIgnore("javax.management.*")
-public class InternalAPIKeyAuthenticatiorTest {
+public class InternalAPIKeyAuthenticatorTest {
 
     @Test
-    public void extractInternalKeyTest() throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
+    public void extractInternalKeyInWSProtocolHeaderTest() {
         PowerMockito.mockStatic(ConfigHolder.class);
         ConfigHolder configHolder = Mockito.mock(ConfigHolder.class);
         EnforcerConfig enforcerConfig = Mockito.mock(EnforcerConfig.class);
@@ -65,7 +63,7 @@ public void extractInternalKeyTest() throws NoSuchMethodException, InvocationTar
         String secWebsocketProtocolHeader = "sec-websocket-protocol";
 
         // Test case to test for an Upgrade request sent from the choreo console
-        // The token will be set to the sec-websocket-protocol header with c keyword
+        // The token will be set to the sec-websocket-protocol header with choreo-internal-API-Key keyword
         // the value after choreo-internal-API-Key will be the token
         RequestContext.Builder builder = new RequestContext.Builder("/pets");
         builder.matchedAPI(new APIConfig.Builder("Petstore")
@@ -75,41 +73,61 @@ public void extractInternalKeyTest() throws NoSuchMethodException, InvocationTar
         Map<String, String> headersMap = new HashMap<>();
         headersMap.put(
                 secWebsocketProtocolHeader,
-            "choreo-internal-API-Key, " + mockToken);
+            "choreo-internal-API-Key," + mockToken);
         builder.headers(headersMap);
         RequestContext requestContext = builder.build();
         InternalAPIKeyAuthenticator internalAPIKeyAuthenticator = new InternalAPIKeyAuthenticator(securityParam);
-        Method extractInternalKeyMethod = InternalAPIKeyAuthenticator.class.getDeclaredMethod(
-                "extractInternalKey",RequestContext.class);
-        extractInternalKeyMethod.setAccessible(true);
-        String internalKey = (String) extractInternalKeyMethod.invoke(internalAPIKeyAuthenticator, requestContext);
-        Assert.assertEquals(internalKey, mockToken);
+        Assert.assertEquals(internalAPIKeyAuthenticator.extractInternalKeyInWSProtocolHeader(requestContext), mockToken);
 
-        // test case to test for Upgrade request sent with the API-Key header set
+        // Test case to test for an Upgrade request sent from a client with api-key
         RequestContext.Builder builder2 = new RequestContext.Builder("/pets");
         builder2.matchedAPI(new APIConfig.Builder("Petstore")
-                .basePath("/choreo")
-                .apiType("WS")
-                .build());
+               .basePath("/choreo")
+               .apiType("WS")
+               .build());
         Map<String, String> headersMap2 = new HashMap<>();
         headersMap2.put(securityParam, mockToken);
-        headersMap2.put(secWebsocketProtocolHeader, "chat");
         builder2.headers(headersMap2);
-        RequestContext requestContext2 = builder.build();
-        String internalKey2 = (String) extractInternalKeyMethod.invoke(internalAPIKeyAuthenticator, requestContext2);
-        Assert.assertEquals(internalKey2, mockToken);
-
-        // Test case to test for normal REST API HTTP Request
-        RequestContext.Builder builder3 = new RequestContext.Builder("/pets");
-        builder3.matchedAPI(new APIConfig.Builder("Petstore")
-                .basePath("/choreo")
-                .apiType("HTTP")
-                .build());
-        Map<String, String> headersMap3 = new HashMap<>();
-        headersMap3.put(securityParam, mockToken);
-        builder3.headers(headersMap3);
-        RequestContext requestContext3 = builder.build();
-        String internalKey3 = (String) extractInternalKeyMethod.invoke(internalAPIKeyAuthenticator, requestContext3);
-        Assert.assertEquals(internalKey3, mockToken);
+        RequestContext requestContext2 = builder2.build();
+        Assert.assertEquals(internalAPIKeyAuthenticator.extractInternalKeyInWSProtocolHeader(requestContext2), null);
+
+    }
+
+    @Test
+    public void getProtocolsToSetInRequestHeadersTest() {
+        PowerMockito.mockStatic(ConfigHolder.class);
+        ConfigHolder configHolder = Mockito.mock(ConfigHolder.class);
+        EnforcerConfig enforcerConfig = Mockito.mock(EnforcerConfig.class);
+        CacheDto cacheDto = Mockito.mock(CacheDto.class);
+        Mockito.when(cacheDto.isEnabled()).thenReturn(true);
+        Mockito.when(enforcerConfig.getCacheDto()).thenReturn(cacheDto);
+        JWTConfigurationDto jwtConfigurationDto = Mockito.mock(JWTConfigurationDto.class);
+        Mockito.when(jwtConfigurationDto.isEnabled()).thenReturn(false);
+        Mockito.when(enforcerConfig.getJwtConfigurationDto()).thenReturn(jwtConfigurationDto);
+        Mockito.when(configHolder.getConfig()).thenReturn(enforcerConfig);
+        Mockito.when(ConfigHolder.getInstance()).thenReturn(configHolder);
+
+        String securityParam = "API-Key";
+
+        String secWebsocketProtocolHeader = "sec-websocket-protocol";
+
+        String mockToken = "eyJraWQiOiJnYXRld2F5XUlMyNTYifQlzaGVyXC92MlwvYXBpc1wvaW50ZXJuYlzaGVyXC92XBpc1wvaW50ZXJuY." +
+        "eyJzdWIiOiJhMzllYGV2OjQ0M1wvYXBpXC9hbVwvcHVibGlzaGVyXC92MlwvYXBpc1wvaW50ZXJuYWwta2V5Iiwia2V5dHlwZcl." +
+        "cnZpY2VcL3YxLjAiLCJwdWJsaXNoZXIiOiJjaG9yZW9fZGV2X2FwaW1fYWRtaW4iLCJ2ZXJzaW9uIjoidj7MIXRnS-2UWHdrmd7";
+
+        RequestContext.Builder builder = new RequestContext.Builder("/pets");
+        builder.matchedAPI(new APIConfig.Builder("Petstore")
+               .basePath("/choreo")
+               .apiType("WS")
+               .build());
+        Map<String, String> headersMap = new HashMap<>();
+        headersMap.put(
+                secWebsocketProtocolHeader,
+            "choreo-internal-API-Key, " + mockToken + ", " + "chat, bar");
+        builder.headers(headersMap);
+        RequestContext requestContext = builder.build();
+        InternalAPIKeyAuthenticator internalAPIKeyAuthenticator = new InternalAPIKeyAuthenticator(securityParam);
+        Assert.assertEquals(internalAPIKeyAuthenticator.getProtocolsToSetInRequestHeaders(requestContext), "chat, bar");
+
     }
 }