add review comments

wso2 · Aug 7, 2024 · 7f491ad · 7f491ad
1 parent e74c1c1
commit 7f491ad
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 18 deletions.
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
@@ -101,8 +101,6 @@ 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.

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
@@ -51,6 +51,7 @@
 import java.text.ParseException;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.Map;
 import java.util.stream.Collectors;
 
 /**
@@ -79,13 +80,7 @@ public boolean canAuthenticate(RequestContext requestContext) {
                 ConfigHolder.getInstance().getConfig().getAuthHeader().getTestConsoleHeaderName().toLowerCase());
         if (apiType.equalsIgnoreCase("WS") && internalKey == null) {
             internalKey = extractInternalKeyInWSProtocolHeader(requestContext);
-
-            if (isAdditionalResponseHeadersRequired(requestContext)) {
-                HashMap<String, String> responseHeadersToAddMap = new HashMap<>();
-                responseHeadersToAddMap.put(
-                        HttpConstants.WEBSOCKET_PROTOCOL_HEADER, Constants.WS_API_KEY_IDENTIFIER);
-                requestContext.setResponseHeadersToAddMap(responseHeadersToAddMap);
-            }
+            addWSProtocolResponseHeaderIfRequired(requestContext);
         }
 
         return isAPIKey(internalKey);
@@ -305,13 +300,13 @@ private String extractInternalKey(RequestContext requestContext) {
         }
         if (requestContext.getMatchedAPI().getApiType().equalsIgnoreCase("WS")) {
             internalKey = extractInternalKeyInWSProtocolHeader(requestContext);
-            if (internalKey != null) {
+            if (internalKey != null && !internalKey.isEmpty()) {
                 String protocols = getProtocolsToSetInRequestHeaders(requestContext);
                 if (protocols != null) {
                     requestContext.addOrModifyHeaders(HttpConstants.WEBSOCKET_PROTOCOL_HEADER, protocols);
                 }
                 return internalKey.trim();
-            } 
+            }
         }
         return null;
     }
@@ -326,7 +321,7 @@ public String extractInternalKeyInWSProtocolHeader(RequestContext requestContext
                 return secProtocolHeaderValues[1].trim();
             }
         }
-        return null;
+        return "";
     }
 
     public String getProtocolsToSetInRequestHeaders(RequestContext requestContext) {
@@ -339,11 +334,22 @@ public String getProtocolsToSetInRequestHeaders(RequestContext requestContext) {
         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;
+    public void addWSProtocolResponseHeaderIfRequired(RequestContext requestContext) {
+        String secProtocolHeader =  requestContext.getHeaders().get(HttpConstants.WEBSOCKET_PROTOCOL_HEADER);
+        if (secProtocolHeader != null) {
+            String[] secProtocolHeaderValues = secProtocolHeader.split(",");
+            if (secProtocolHeaderValues[0].equals(Constants.WS_API_KEY_IDENTIFIER) &&
+                    secProtocolHeaderValues.length == 2) {
+                Map<String, String> responseHeadersToAddMap = requestContext.getResponseHeadersToAddMap();
+
+                if (responseHeadersToAddMap == null) {
+                    responseHeadersToAddMap = new HashMap<>();
+                }
+                responseHeadersToAddMap.put(
+                        HttpConstants.WEBSOCKET_PROTOCOL_HEADER, Constants.WS_API_KEY_IDENTIFIER);
+                requestContext.setResponseHeadersToAddMap(responseHeadersToAddMap);
+            }
+        }
     }
 
     @Override

diff --git a/...t/java/org/wso2/choreo/connect/enforcer/security/jwt/InternalAPIKeyAuthenticatorTest.java b/...t/java/org/wso2/choreo/connect/enforcer/security/jwt/InternalAPIKeyAuthenticatorTest.java
@@ -89,7 +89,7 @@ public void extractInternalKeyInWSProtocolHeaderTest() {
         headersMap2.put(securityParam, mockToken);
         builder2.headers(headersMap2);
         RequestContext requestContext2 = builder2.build();
-        Assert.assertEquals(internalAPIKeyAuthenticator.extractInternalKeyInWSProtocolHeader(requestContext2), null);
+        Assert.assertEquals(internalAPIKeyAuthenticator.extractInternalKeyInWSProtocolHeader(requestContext2), "");
 
     }