websocketproxy: fix and add test for subprotocols

Author: fatih

websocketproxy.go

@@ -80,11 +80,11 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	// the final destinations.
 	h := http.Header{}
 	h.Add("Origin", req.Header.Get("Origin"))
-	protocols := req.Header["Sec-WebSocket-Protocol"]
+	protocols := req.Header[http.CanonicalHeaderKey("Sec-WebSocket-Protocol")]
 	for _, prot := range protocols {
 		h.Add("Sec-WebSocket-Protocol", prot)
 	}
-	cookies := req.Header["Cookie"]
+	cookies := req.Header[http.CanonicalHeaderKey("Cookie")]
 	for _, cookie := range cookies {
 		h.Add("Cookie", cookie)
 	}
@@ -114,7 +114,8 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	// Connect to the backend URL, also pass the headers we prepared above.
 	// TODO: support multiplexing on the same backend connection instead of
 	// opening a new TCP connection time for each request. This should be
-	// optional.
+	// optional:
+	// http://tools.ietf.org/html/draft-ietf-hybi-websocket-multiplexing-01
 	connBackend, resp, err := dialer.Dial(backendURL.String(), h)
 	if err != nil {
 		log.Printf("websocketproxy: couldn't dial to remote backend url %s\n", err)
@@ -129,8 +130,10 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 
 	// Only pass those headers to the upgrader.
 	upgradeHeader := http.Header{}
-	upgradeHeader.Add("Sec-WebSocket-Protocol", resp.Header.Get("Sec-WebSocket-Protocol"))
-	upgradeHeader.Add("Set-Cookie", resp.Header.Get("Set-Cookie"))
+	upgradeHeader.Set("Sec-WebSocket-Protocol",
+		resp.Header.Get(http.CanonicalHeaderKey("Sec-WebSocket-Protocol")))
+	upgradeHeader.Set("Set-Cookie",
+		resp.Header.Get(http.CanonicalHeaderKey("Set-Cookie")))
 
 	// Now upgrade the existing incoming request to a WebSocket connection.
 	// Also pass the header that we gathered from the Dial handshake.

websocketproxy_test.go

@@ -17,12 +17,14 @@ var (
 
 func TestProxy(t *testing.T) {
 	// websocket proxy
+	supportedSubProtocols := []string{"test-protocol"}
 	upgrader := &websocket.Upgrader{
 		ReadBufferSize:  4096,
 		WriteBufferSize: 4096,
 		CheckOrigin: func(r *http.Request) bool {
 			return true
 		},
+		Subprotocols: supportedSubProtocols,
 	}
 
 	u, _ := url.Parse(backendURL)
@@ -67,13 +69,41 @@ func TestProxy(t *testing.T) {
 
 	time.Sleep(time.Millisecond * 100)
 
+	// let's us define two subprotocols, only one is supported by the server
+	clientSubProtocols := []string{"test-protocol", "test-notsupported"}
+	h := http.Header{}
+	for _, subprot := range clientSubProtocols {
+		h.Add("Sec-WebSocket-Protocol", subprot)
+	}
+
 	// frontend server, dial now our proxy, which will reverse proxy our
 	// message to the backend websocket server.
-	conn, _, err := websocket.DefaultDialer.Dial(serverURL+"/proxy", nil)
+	conn, resp, err := websocket.DefaultDialer.Dial(serverURL+"/proxy", h)
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	// check if the server really accepted only the first one
+	in := func(desired string) bool {
+		for _, prot := range resp.Header[http.CanonicalHeaderKey("Sec-WebSocket-Protocol")] {
+			if desired == prot {
+				return true
+			}
+		}
+
+		return false
+	}
+
+	if !in("test-protocol") {
+		t.Error("test-protocol should be available")
+	}
+
+	if in("test-notsupported") {
+		t.Error("test-notsupported should be not recevied from the server.")
+	}
+
+	// now write a message and send it to the backend server (which goes trough
+	// proxy..)
 	msg := "hello kite"
 	err = conn.WriteMessage(websocket.TextMessage, []byte(msg))
 	if err != nil {