Merge pull request #3610 from Thushani-Jayasekera/choreo

[Choreo] Add support to use test key keyword at WS Handshake Authentication
wso2 · Oct 28, 2024 · c4ee825 · c4ee825
2 parents 4af100b + 87431a0
commit c4ee825
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 3 deletions.
diff --git a/...r-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/constants/Constants.java b/...r-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/constants/Constants.java
@@ -73,6 +73,7 @@ public class Constants {
     public static final String DEFAULT_CON_FACTORY_JNDI_NAME = "TopicConnectionFactory";
 
     // keyword to identify API-Key sent in sec-websocket-protocol header
+    public static final String WS_TEST_KEY_IDENTIFIER = "choreo-test-key";
     public static final String WS_API_KEY_IDENTIFIER = "choreo-api-key";
     public static final String WS_OAUTH2_KEY_IDENTIFIED = "choreo-oauth2-token";
 

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
@@ -39,7 +39,9 @@
 import org.wso2.choreo.connect.enforcer.config.dto.ExtendedTokenIssuerDto;
 import org.wso2.choreo.connect.enforcer.constants.APIConstants;
 import org.wso2.choreo.connect.enforcer.constants.APISecurityConstants;
+import org.wso2.choreo.connect.enforcer.constants.Constants;
 import org.wso2.choreo.connect.enforcer.constants.GeneralErrorCodeConstants;
+import org.wso2.choreo.connect.enforcer.constants.HttpConstants;
 import org.wso2.choreo.connect.enforcer.dto.APIKeyValidationInfoDTO;
 import org.wso2.choreo.connect.enforcer.dto.JWTTokenPayloadInfo;
 import org.wso2.choreo.connect.enforcer.exception.APISecurityException;
@@ -53,7 +55,9 @@
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.text.ParseException;
+import java.util.Arrays;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 /**
  * Extends the APIKeyHandler to authenticate request using API Key.
@@ -117,6 +121,21 @@ private static String getAPIKeyFromRequest(RequestContext requestContext) {
                 }
             }
         }
+
+        // If an API Key is not found, check for the API Key in the WebSocket protocol
+        // header
+        if (requestContext.getMatchedAPI().getApiType().equalsIgnoreCase(APIConstants.ApiType.WEB_SOCKET) &&
+                requestContext.getHeaders().containsKey(HttpConstants.WEBSOCKET_PROTOCOL_HEADER)) {
+            String apiKey = extractAPIKeyInWSProtocolHeader(requestContext);
+            if (apiKey != null && !apiKey.isEmpty()) {
+                String protocols = getProtocolsToSetInRequestHeaders(requestContext);
+                if (protocols != null) {
+                    requestContext.addOrModifyHeaders(HttpConstants.WEBSOCKET_PROTOCOL_HEADER, protocols);
+                }
+                return apiKey.trim();
+            }
+        }
+
         return "";
     }
 
@@ -458,4 +477,29 @@ public String getName() {
     public int getPriority() {
         return 30;
     }
+
+    public static String extractAPIKeyInWSProtocolHeader(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)) {
+                AuthenticatorUtils.addWSProtocolResponseHeaderIfRequired(requestContext,
+                        Constants.WS_API_KEY_IDENTIFIER);
+                return secProtocolHeaderValues[1].trim();
+            }
+        }
+        return "";
+    }
+
+    public static 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;
+    }
 }
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
@@ -93,7 +93,11 @@ public boolean canAuthenticate(RequestContext requestContext) {
             if (internalKey == null) {
                 internalKey = extractInternalKeyInWSProtocolHeader(requestContext);
             }
-            AuthenticatorUtils.addWSProtocolResponseHeaderIfRequired(requestContext, Constants.WS_API_KEY_IDENTIFIER);
+            String protocolHeader = requestContext.getHeaders().get(HttpConstants.WEBSOCKET_PROTOCOL_HEADER);
+            if (protocolHeader != null) {
+                String[] secProtocolHeaderValues = protocolHeader.split(",");
+                AuthenticatorUtils.addWSProtocolResponseHeaderIfRequired(requestContext, secProtocolHeaderValues[0]);
+            }
         }
 
         return isAPIKey(internalKey);
@@ -335,8 +339,9 @@ public String extractInternalKeyInWSProtocolHeader(RequestContext requestContext
                 HttpConstants.WEBSOCKET_PROTOCOL_HEADER);
         if (protocolHeader != null) {
             String[] secProtocolHeaderValues = protocolHeader.split(",");
-            if (secProtocolHeaderValues.length > 1 && secProtocolHeaderValues[0].equals(
-                    Constants.WS_API_KEY_IDENTIFIER)) {
+            if (secProtocolHeaderValues.length > 1 && (secProtocolHeaderValues[0].equals(
+                    Constants.WS_TEST_KEY_IDENTIFIER)
+                    || secProtocolHeaderValues[0].equals(Constants.WS_API_KEY_IDENTIFIER))) {
                 return secProtocolHeaderValues[1].trim();
             }
         }