[Choreo] Support passing API-Key for WebSocket Requests using sec-websocket-protocol Header

enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/constants/Constants.java

@@ -71,4 +71,7 @@ public class Constants {
     public static final String PROP_CON_FACTORY = "connectionfactory.TopicConnectionFactory";
     public static final String DEFAULT_DESTINATION_TYPE = "Topic";
     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_API_KEY_IDENTIFIER = "choreo-internal-API-Key";
 }

enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/constants/HttpConstants.java

@@ -28,4 +28,5 @@ public class HttpConstants {
     public static final String X_REQUEST_ID_HEADER = "x-request-id";
     public static final String APPLICATION_JSON = "application/json";
     public static final String BASIC_LOWER = "basic";
+    public static final String WEBSOCKET_PROTOCOL_HEADER = "sec-websocket-protocol";
 }

enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/grpc/ExtAuthService.java

@@ -37,6 +37,7 @@
 import org.json.JSONObject;
 import org.wso2.choreo.connect.enforcer.api.ResponseObject;
 import org.wso2.choreo.connect.enforcer.constants.APIConstants;
+import org.wso2.choreo.connect.enforcer.constants.Constants;
 import org.wso2.choreo.connect.enforcer.constants.HttpConstants;
 import org.wso2.choreo.connect.enforcer.constants.MetadataConstants;
 import org.wso2.choreo.connect.enforcer.constants.RouterAccessLogConstants;
@@ -101,6 +102,8 @@ private CheckResponse buildResponse(CheckRequest request, ResponseObject respons
         DeniedHttpResponse.Builder responseBuilder = DeniedHttpResponse.newBuilder();
         HttpStatus status = HttpStatus.newBuilder().setCodeValue(responseObject.getStatusCode()).build();
         String traceKey = request.getAttributes().getRequest().getHttp().getId();
+        String[] secProtocolHeaderForWS = request.getAttributes().getRequest().getHttp().getHeadersOrDefault(
+            HttpConstants.WEBSOCKET_PROTOCOL_HEADER, "").split(",");
         Struct.Builder structBuilder = Struct.newBuilder();
         // Used to identify that the choreo-connect-enforcer handled the request. It is used to
         // provide local reply for authentication failures.
@@ -151,7 +154,15 @@ private CheckResponse buildResponse(CheckRequest request, ResponseObject respons
                     .build();
         } else {
             OkHttpResponse.Builder okResponseBuilder = OkHttpResponse.newBuilder();
-
+            if (secProtocolHeaderForWS[0].equals(Constants.WS_API_KEY_IDENTIFIER) &&
+                secProtocolHeaderForWS.length == 2) {
+                okResponseBuilder.addResponseHeadersToAdd(
+                    HeaderValueOption.newBuilder()
+                        .setHeader(HeaderValue.newBuilder()
+                        .setKey(HttpConstants.WEBSOCKET_PROTOCOL_HEADER)
+                        .setValue(Constants.WS_API_KEY_IDENTIFIER).build())
+                        .build());
+            }
             // If the user is sending the APIKey credentials within query parameters, those query parameters should
             // not be sent to the backend. Hence, the :path header needs to be constructed again removing the apiKey
             // query parameter. In this scenario, apiKey query parameter is sent within the property called

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

@@ -36,6 +36,8 @@
 import org.wso2.choreo.connect.enforcer.config.EnforcerConfig;
 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.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;
@@ -47,6 +49,8 @@
 import org.wso2.choreo.connect.enforcer.util.FilterUtils;
 
 import java.text.ParseException;
+import java.util.Arrays;
+import java.util.stream.Collectors;
 
 /**
  * Implements the authenticator interface to authenticate request using an Internal Key.
@@ -69,8 +73,20 @@ public InternalAPIKeyAuthenticator(String securityParam) {
 
     @Override
     public boolean canAuthenticate(RequestContext requestContext) {
+        String apiType = requestContext.getMatchedAPI().getApiType();
         String internalKey = requestContext.getHeaders().get(
-                ConfigHolder.getInstance().getConfig().getAuthHeader().getTestConsoleHeaderName().toLowerCase());
+            ConfigHolder.getInstance().getConfig().getAuthHeader().getTestConsoleHeaderName().toLowerCase());
+        if (apiType.equalsIgnoreCase("WS")) {
+            if (internalKey == null || internalKey.isBlank()) {
+                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];
+                }
+            }
+        }
+
         return isAPIKey(internalKey);
     }
 
@@ -281,10 +297,29 @@ public String getName() {
     }
 
     private String extractInternalKey(RequestContext requestContext) {
-        String internalKey = requestContext.getHeaders().get(securityParam);
+        String internalKey;
+        internalKey = requestContext.getHeaders().get(securityParam);
         if (internalKey != null) {
             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];
+            }
+            if (internalKey != null) {
+                if (secProtocolHeaderValues.length > 2) {
+                    String protocols = Arrays.stream(
+                        secProtocolHeaderValues, 2,
+                        secProtocolHeaderValues.length)
+                            .collect(Collectors.joining(","));
+                    requestContext.addOrModifyHeaders(HttpConstants.WEBSOCKET_PROTOCOL_HEADER, protocols);
+                }
+                return internalKey.trim();
+            } 
+        }
         return null;
     }