support for handling headers of websocket handshake

src/main/java/com/github/nkzawa/engineio/client/transports/WebSocket.java

@@ -14,6 +14,7 @@ import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Date;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 
 public class WebSocket extends Transport {
@@ -33,14 +34,25 @@ public class WebSocket extends Transport {
             return;
         }
 
+        Map<String, String> headers = new HashMap<String, String>();
+        this.emit(EVENT_REQUEST_HEADERS, headers);
+
         final WebSocket self = this;
         try {
-            this.socket = new WebSocketClient(new URI(this.uri()), new Draft_17()) {
+            this.socket = new WebSocketClient(new URI(this.uri()), new Draft_17(), headers, 0) {
                 @Override
-                public void onOpen(ServerHandshake serverHandshake) {
+                public void onOpen(final ServerHandshake serverHandshake) {
                     EventThread.exec(new Runnable() {
                         @Override
                         public void run() {
+                            Map<String, String> headers = new HashMap<String, String>();
+                            Iterator<String> it = serverHandshake.iterateHttpFields();
+                            while (it.hasNext()) {
+                                String field = it.next();
+                                headers.put(field, serverHandshake.getFieldValue(field));
+                            }
+                            self.emit(EVENT_RESPONSE_HEADERS, headers);
+
                             self.onOpen();
                         }
                     });

src/test/java/com/github/nkzawa/engineio/client/ServerConnectionTest.java

@@ -2,7 +2,7 @@ package com.github.nkzawa.engineio.client;
 
 import com.github.nkzawa.emitter.Emitter;
 import com.github.nkzawa.engineio.client.transports.Polling;
-import com.github.nkzawa.engineio.client.transports.PollingXHR;
+import com.github.nkzawa.engineio.client.transports.WebSocket;
 import com.github.nkzawa.engineio.parser.HandshakeData;
 import org.junit.After;
 import org.junit.Before;
@@ -244,8 +244,6 @@ public class ServerConnectionTest {
             @Override
             public void call(Object... args) {
                 Transport transport = (Transport)args[0];
-                if (!(transport instanceof PollingXHR)) return;
-
                 transport.on(Transport.EVENT_REQUEST_HEADERS, new Emitter.Listener() {
                     @Override
                     public void call(Object... args) {
@@ -269,4 +267,48 @@ public class ServerConnectionTest {
         assertThat(messages.take(), is("foo"));
         socket.close();
     }
+
+    @Test(timeout = TIMEOUT)
+    public void websocketHandshakeHeaders() throws URISyntaxException, InterruptedException {
+        final BlockingQueue<String> messages = new LinkedBlockingQueue<String>();
+
+        Socket.Options opts = new Socket.Options();
+        opts.transports = new String[] {WebSocket.NAME};
+
+        socket = new Socket("ws://localhost:" + PORT, opts) {
+            @Override
+            public void onopen() {}
+            @Override
+            public void onmessage(String data) {}
+            @Override
+            public void onclose() {}
+            @Override
+            public void onerror(Exception err) {}
+        };
+        socket.on(Socket.EVENT_TRANSPORT, new Emitter.Listener() {
+            @Override
+            public void call(Object... args) {
+                Transport transport = (Transport)args[0];
+                transport.on(Transport.EVENT_REQUEST_HEADERS, new Emitter.Listener() {
+                    @Override
+                    public void call(Object... args) {
+                        @SuppressWarnings("unchecked")
+                        Map<String, String> headers = (Map<String, String>)args[0];
+                        headers.put("X-EngineIO", "foo");
+                    }
+                }).on(Transport.EVENT_RESPONSE_HEADERS, new Emitter.Listener() {
+                    @Override
+                    public void call(Object... args) {
+                        @SuppressWarnings("unchecked")
+                        Map<String, String> headers = (Map<String, String>)args[0];
+                        messages.offer(headers.get("X-EngineIO"));
+                    }
+                });
+            }
+        });
+        socket.open();
+
+        assertThat(messages.take(), is("foo"));
+        socket.close();
+    }
 }

src/test/resources/index.js

@@ -19,10 +19,26 @@ server.on('connection', function(socket) {
 
 var handleRequest = server.handleRequest;
 server.handleRequest = function(req, res) {
-  var header = req.headers['x-engineio'];
+  // echo a header value
-  if (header) {
+  var value = req.headers['x-engineio'];
-    res.setHeader('X-EngineIO', header);
+  if (value) {
+    res.setHeader('X-EngineIO', value);
   }
 
   handleRequest.call(this, req, res);
 };
+
+var headerValue;
+var handleUpgrade = server.handleUpgrade;
+server.handleUpgrade = function(req, socket, head) {
+  // echo a header value for websocket handshake
+  headerValue = req.headers['x-engineio'];
+  handleUpgrade.call(this, req, socket, head);
+};
+
+// FIXME: support parallel requests
+server.ws.on('headers', function(headers) {
+  if (headerValue) {
+    headers.push('X-EngineIO: ' + headerValue);
+  }
+});