Merge pull request #1693 from DinithHerath/master

Add implementation to preserve websocket headers with null check for preserved header value
wso2 · Dec 6, 2023 · 348803f · 348803f
2 parents 5b7439a + e96d025
commit 348803f
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.