Merge pull request #123 from rshelley/master

Added support for custom headers on websocket client connections

src/main/java/io/netty/example/http/websocketx/client/WebSocketClient.java

@@ -51,4 +51,14 @@ public interface WebSocketClient {
      * @return Write future. Will fire when the data is sent.
      */
     ChannelFuture send(WebSocketFrame frame);
+    
+    /**
+     * Adds a custom header to this client request
+     * 
+     * @param header
+     * 			Name of header field to add to request
+     * @param value
+     * 			Value of header field added to request
+     */
+    void addCustomHeader(String header, String value);
 }

src/main/java/io/netty/example/http/websocketx/client/WebSocketClientHandler.java

@@ -24,6 +24,8 @@ package io.netty.example.http.websocketx.client;
 
 import java.net.InetSocketAddress;
 import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
 
 import io.netty.bootstrap.ClientBootstrap;
 import io.netty.channel.Channel;
@@ -54,7 +56,8 @@ public class WebSocketClientHandler extends SimpleChannelUpstreamHandler impleme
     private Channel channel;
     private WebSocketClientHandshaker handshaker = null;
     private final WebSocketSpecificationVersion version;
-
+    private Map<String,String> customHeaders = null;
+    
     public WebSocketClientHandler(ClientBootstrap bootstrap, URI url, WebSocketSpecificationVersion version, WebSocketCallback callback) {
         this.bootstrap = bootstrap;
         this.url = url;
@@ -65,7 +68,7 @@ public class WebSocketClientHandler extends SimpleChannelUpstreamHandler impleme
     @Override
     public void channelConnected(ChannelHandlerContext ctx, ChannelStateEvent e) throws Exception {
         channel = e.getChannel();
-        this.handshaker = new WebSocketClientHandshakerFactory().newHandshaker(url, version, null, false);
+        this.handshaker = new WebSocketClientHandshakerFactory().newHandshaker(url, version, null, false, customHeaders);
         handshaker.performOpeningHandshake(channel);
     }
 
@@ -121,4 +124,11 @@ public class WebSocketClientHandler extends SimpleChannelUpstreamHandler impleme
     public void setUrl(URI url) {
         this.url = url;
     }
+
+    public void addCustomHeader(String header, String value){
+    	if(customHeaders == null){
+    		customHeaders = new HashMap<String,String>();
+    	}
+    	customHeaders.put(header, value);
+    }
 }

src/main/java/io/netty/handler/codec/http/websocketx/WebSocketClientHandshaker.java

@@ -18,6 +18,7 @@ package io.netty.handler.codec.http.websocketx;
 import java.net.URI;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.Map;
 
 import io.netty.buffer.ChannelBuffer;
 import io.netty.buffer.ChannelBuffers;
@@ -41,16 +42,19 @@ public abstract class WebSocketClientHandshaker {
 
     private String subProtocolResponse = null;
 
+    protected Map<String,String> customHeaders = null;
+    
     /**
      * 
      * @param webSocketURL
      * @param version
      * @param subProtocol
      */
-    public WebSocketClientHandshaker(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol) {
+    public WebSocketClientHandshaker(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, Map<String,String> customHeaders) {
         this.webSocketURL = webSocketURL;
         this.version = version;
         this.subProtocolRequest = subProtocol;
+        this.customHeaders = customHeaders;
     }
 
     /**

src/main/java/io/netty/handler/codec/http/websocketx/WebSocketClientHandshaker00.java

@@ -18,6 +18,7 @@ package io.netty.handler.codec.http.websocketx;
 import java.net.URI;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.Map;
 
 import io.netty.buffer.ChannelBuffers;
 import io.netty.channel.Channel;
@@ -60,9 +61,12 @@ public class WebSocketClientHandshaker00 extends WebSocketClientHandshaker {
      *            server
      * @param subProtocol
      *            Sub protocol request sent to the server.
+     * @param customHeaders
+     *            Map of custom headers to add to the client request
      */
-    public WebSocketClientHandshaker00(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol) {
+    public WebSocketClientHandshaker00(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, Map<String,String> customHeaders) {
-        super(webSocketURL, version, subProtocol);
+        super(webSocketURL, version, subProtocol, customHeaders);
+        
     }
 
     /**
@@ -142,6 +146,13 @@ public class WebSocketClientHandshaker00 extends WebSocketClientHandshaker {
         if (this.getSubProtocolRequest() != null && !this.getSubProtocolRequest().equals("")) {
             request.addHeader(Names.SEC_WEBSOCKET_PROTOCOL, this.getSubProtocolRequest());
         }
+
+        if(customHeaders != null){
+        	for(String header: customHeaders.keySet()){
+        		request.addHeader(header, customHeaders.get(header));
+        	}
+        }
+        
         request.setContent(ChannelBuffers.copiedBuffer(key3));
 
         channel.write(request);

src/main/java/io/netty/handler/codec/http/websocketx/WebSocketClientHandshaker10.java

@@ -16,6 +16,7 @@
 package io.netty.handler.codec.http.websocketx;
 
 import java.net.URI;
+import java.util.Map;
 
 import io.netty.channel.Channel;
 import io.netty.handler.codec.http.DefaultHttpRequest;
@@ -68,9 +69,11 @@ public class WebSocketClientHandshaker10 extends WebSocketClientHandshaker {
      * @param allowExtensions
      *            Allow extensions to be used in the reserved bits of the web
      *            socket frame
+     * @param customHeaders
+     *            Map of custom headers to add to the client request
      */
-    public WebSocketClientHandshaker10(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, boolean allowExtensions) {
+    public WebSocketClientHandshaker10(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, boolean allowExtensions, Map<String,String> customHeaders) {
-        super(webSocketURL, version, subProtocol);
+        super(webSocketURL, version, subProtocol, customHeaders);
         this.allowExtensions = allowExtensions;
     }
 
@@ -127,6 +130,12 @@ public class WebSocketClientHandshaker10 extends WebSocketClientHandshaker {
         }
         request.addHeader(Names.SEC_WEBSOCKET_VERSION, "8");
 
+        if(customHeaders != null){
+        	for(String header: customHeaders.keySet()){
+        		request.addHeader(header, customHeaders.get(header));
+        	}
+        }
+        
         channel.write(request);
 
         channel.getPipeline().replace(HttpRequestEncoder.class, "ws-encoder", new WebSocket08FrameEncoder(true));

src/main/java/io/netty/handler/codec/http/websocketx/WebSocketClientHandshaker17.java

@@ -16,6 +16,7 @@
 package io.netty.handler.codec.http.websocketx;
 
 import java.net.URI;
+import java.util.Map;
 
 import io.netty.channel.Channel;
 import io.netty.handler.codec.http.DefaultHttpRequest;
@@ -51,7 +52,7 @@ public class WebSocketClientHandshaker17 extends WebSocketClientHandshaker {
     private static final String protocol = null;
 
     private boolean allowExtensions = false;
-
+    
     /**
      * Constructor specifying the destination web socket location and version to
      * initiate
@@ -68,9 +69,11 @@ public class WebSocketClientHandshaker17 extends WebSocketClientHandshaker {
      * @param allowExtensions
      *            Allow extensions to be used in the reserved bits of the web
      *            socket frame
+     * @param customHeaders
+     *            Map of custom headers to add to the client request
      */
-    public WebSocketClientHandshaker17(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, boolean allowExtensions) {
+    public WebSocketClientHandshaker17(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, boolean allowExtensions, Map<String,String> customHeaders) {
-        super(webSocketURL, version, subProtocol);
+        super(webSocketURL, version, subProtocol, customHeaders);
         this.allowExtensions = allowExtensions;
     }
 
@@ -127,6 +130,11 @@ public class WebSocketClientHandshaker17 extends WebSocketClientHandshaker {
         }
         request.addHeader(Names.SEC_WEBSOCKET_VERSION, "13");
 
+        if(customHeaders != null){
+        	for(String header: customHeaders.keySet()){
+        		request.addHeader(header, customHeaders.get(header));
+        	}
+        }
         channel.write(request);
 
         channel.getPipeline().replace(HttpRequestEncoder.class, "ws-encoder", new WebSocket13FrameEncoder(true));

src/main/java/io/netty/handler/codec/http/websocketx/WebSocketClientHandshakerFactory.java

@@ -16,6 +16,7 @@
 package io.netty.handler.codec.http.websocketx;
 
 import java.net.URI;
+import java.util.Map;
 
 /**
  * Instances the appropriate handshake class to use for clients
@@ -40,15 +41,15 @@ public class WebSocketClientHandshakerFactory {
      *            socket frame
      * @throws WebSocketHandshakeException
      */
-    public WebSocketClientHandshaker newHandshaker(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, boolean allowExtensions) throws WebSocketHandshakeException {
+    public WebSocketClientHandshaker newHandshaker(URI webSocketURL, WebSocketSpecificationVersion version, String subProtocol, boolean allowExtensions, Map<String,String> customHeaders) throws WebSocketHandshakeException {
         if (version == WebSocketSpecificationVersion.V17) {
-            return new WebSocketClientHandshaker17(webSocketURL, version, subProtocol, allowExtensions);
+            return new WebSocketClientHandshaker17(webSocketURL, version, subProtocol, allowExtensions, customHeaders);
         }
         if (version == WebSocketSpecificationVersion.V10) {
-            return new WebSocketClientHandshaker10(webSocketURL, version, subProtocol, allowExtensions);
+            return new WebSocketClientHandshaker10(webSocketURL, version, subProtocol, allowExtensions, customHeaders);
         }
         if (version == WebSocketSpecificationVersion.V00) {
-            return new WebSocketClientHandshaker00(webSocketURL, version, subProtocol);
+            return new WebSocketClientHandshaker00(webSocketURL, version, subProtocol, customHeaders);
         }
 
         throw new WebSocketHandshakeException("Protocol version " + version.toString() + " not supported.");