Add implementation to preserve websocket headers

wso2 · Dec 6, 2023 · e96d025 · e96d025
1 parent 5b7439a
commit e96d025
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/...ocket.transport/src/main/java/org/wso2/carbon/websocket/transport/WebsocketConstants.java b/...ocket.transport/src/main/java/org/wso2/carbon/websocket/transport/WebsocketConstants.java
@@ -62,6 +62,7 @@ public class WebsocketConstants {
 
     public static final String WEBSOCKET_CUSTOM_HEADER_PREFIX = "websocket.custom.header.";
     public static final String WEBSOCKET_CUSTOM_HEADER_CONFIG = "ws.custom.header";
+    public static final String WEBSOCKET_HEADERS_PRESERVE_CONFIG = "ws.headers.preserve";
     public static final String WEBSOCKET_HOSTNAME_VERIFICATION_CONFIG = "ws.client.enable.hostname.verification";
 
     public static final String WEBSOCKET_SUBPROTOCOL = "websocket.subprotocol";

diff --git a/...transport/src/main/java/org/wso2/carbon/websocket/transport/WebsocketTransportSender.java b/...transport/src/main/java/org/wso2/carbon/websocket/transport/WebsocketTransportSender.java
@@ -154,6 +154,23 @@ public void sendMessage(MessageContext msgCtx, String targetEPR, OutTransportInf
             tenantDomain = MultitenantConstants.SUPER_TENANT_DOMAIN_NAME;
         }
 
+        /*
+        * Preserve the specified headers from the original request by adding to the customHeaders map.
+        * Headers present in this map won't be set at Netty level.
+        */
+        Parameter preserveWebSocketHeadersParameter = msgCtx.getTransportOut()
+                .getParameter(WebsocketConstants.WEBSOCKET_HEADERS_PRESERVE_CONFIG);
+        if (preserveWebSocketHeadersParameter != null && preserveWebSocketHeadersParameter.getValue() != null &&
+                !preserveWebSocketHeadersParameter.getValue().toString().isEmpty()) {
+            String preservableHeaders = preserveWebSocketHeadersParameter.getValue().toString();
+            for (String header : preservableHeaders.split(",")) {
+                Object headerValue = msgCtx.getProperty(header);
+                if (headerValue != null) {
+                    customHeaders.put(header, headerValue);
+                }
+            }
+        }
+
         /*
         * Get all the message property names and check whether the properties with the websocket custom header
         * prefix are exist in the property map.