Add changes related to extracting API key from enforcer

wso2 · Aug 1, 2024 · a97bb09 · a97bb09
1 parent 9ae6d76
commit a97bb09
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 4 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
@@ -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";
 }
diff --git a/...rent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/constants/HttpConstants.java b/...rent/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";
 }
diff --git a/...r-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/grpc/ExtAuthService.java b/...r-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,14 @@ 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

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
@@ -36,6 +36,7 @@
 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.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 +48,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.
@@ -55,6 +58,7 @@ public class InternalAPIKeyAuthenticator extends APIKeyHandler {
 
     private static final Log log = LogFactory.getLog(InternalAPIKeyAuthenticator.class);
     private String securityParam;
+    private String secProtocolHeader = "sec-webSocket-protocol";
     private AbstractAPIMgtGatewayJWTGenerator jwtGenerator;
     private final boolean isGatewayTokenCacheEnabled;
 
@@ -68,9 +72,20 @@ public InternalAPIKeyAuthenticator(String securityParam) {
     }
 
     @Override
+    // check if sec-websocket-protocol header is there
     public boolean canAuthenticate(RequestContext requestContext) {
-        String internalKey = requestContext.getHeaders().get(
-                ConfigHolder.getInstance().getConfig().getAuthHeader().getTestConsoleHeaderName().toLowerCase());
+        String apiType = requestContext.getMatchedAPI().getApiType();
+        String internalKey;
+        if (apiType.equalsIgnoreCase("WS")) {
+            internalKey = requestContext.getHeaders().get(
+                    ConfigHolder.getInstance().getConfig().getAuthHeader().getTestConsoleHeaderName().toLowerCase());
+            if (internalKey == null || internalKey.isBlank()) {
+                internalKey = requestContext.getHeaders().get(HttpConstants.WEBSOCKET_PROTOCOL_HEADER).split(",") [1];
+            }
+        } else  {
+            internalKey = requestContext.getHeaders().get(
+                    ConfigHolder.getInstance().getConfig().getAuthHeader().getTestConsoleHeaderName().toLowerCase());
+        }
         return isAPIKey(internalKey);
     }
 
@@ -281,10 +296,25 @@ 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 secProtocolHeaderValue = requestContext.getHeaders().get(HttpConstants.WEBSOCKET_PROTOCOL_HEADER);
+            internalKey = secProtocolHeaderValue.split(",")[1];
+            if (internalKey != null) {
+                if (secProtocolHeaderValue.split(",").length > 2) {
+                    String protocols = Arrays.stream(
+                        secProtocolHeaderValue.split(","), 2,
+                        secProtocolHeaderValue.split(",").length)
+                            .collect(Collectors.joining(","));
+                    requestContext.addOrModifyHeaders(HttpConstants.WEBSOCKET_PROTOCOL_HEADER, protocols);
+                }
+                return internalKey.trim();
+            } 
+        }
         return null;
     }