Add subprotocol negotiation to Dialer.

client.go

@@ -67,6 +67,7 @@ func NewClient(netConn net.Conn, u *url.URL, requestHeader http.Header, readBufS
 		resp.Header.Get("Sec-Websocket-Accept") != acceptKey {
 		return nil, resp, ErrBadHandshake
 	}
+	c.subprotocol = resp.Header.Get("Sec-Websocket-Protocol")
 	return c, resp, nil
 }
 
@@ -85,6 +86,9 @@ type Dialer struct {
 	// Input and output buffer sizes. If the buffer size is zero, then a
 	// default value of 4096 is used.
 	ReadBufferSize, WriteBufferSize int
+
+	// Subprotocols specifies the client's requested subprotocols.
+	Subprotocols []string
 }
 
 var errMalformedURL = errors.New("malformed ws or wss URL")
@@ -206,6 +210,15 @@ func (d *Dialer) Dial(urlStr string, requestHeader http.Header) (*Conn, *http.Re
 		writeBufferSize = 4096
 	}
 
+	if len(d.Subprotocols) > 0 {
+		h := http.Header{}
+		for k, v := range requestHeader {
+			h[k] = v
+		}
+		h.Set("Sec-Websocket-Protocol", strings.Join(d.Subprotocols, ", "))
+		requestHeader = h
+	}
+
 	conn, resp, err := NewClient(
 		netConn,
 		&url.URL{Host: host + port, Opaque: opaque},

client_server_test.go

@@ -8,11 +8,11 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"io"
-	"io/ioutil"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"reflect"
 	"testing"
 	"time"
 )
@@ -24,15 +24,24 @@ type handshakeHandler struct {
 func (t handshakeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if r.Method != "GET" {
 		http.Error(w, "Method not allowed", 405)
-		t.Logf("bad method: %s", r.Method)
+		t.Logf("method = %s, want GET", r.Method)
 		return
 	}
-	if r.Header.Get("Origin") != "http://"+r.Host {
+	if origin := r.Header.Get("Origin"); origin != "http://"+r.Host {
 		http.Error(w, "Origin not allowed", 403)
-		t.Logf("bad origin: %s", r.Header.Get("Origin"))
+		t.Logf("Origin = %s, want %s", origin, r.Host)
 		return
 	}
-	ws, err := Upgrade(w, r, http.Header{"Set-Cookie": {"sessionID=1234"}}, 1024, 1024)
+	subprotos := Subprotocols(r)
+	if !reflect.DeepEqual(subprotos, handshakeDialer.Subprotocols) {
+		http.Error(w, "bad protocol", 400)
+		t.Logf("Subprotocols = %v, want %v", subprotos, handshakeDialer.Subprotocols)
+		return
+	}
+	ws, err := Upgrade(w, r, http.Header{
+		"Set-Cookie":             {"sessionID=1234"},
+		"Sec-Websocket-Protocol": {subprotos[0]},
+	}, 1024, 1024)
 	if _, ok := err.(HandshakeError); ok {
 		t.Logf("bad handshake: %v", err)
 		http.Error(w, "Not a websocket handshake", 400)
@@ -42,6 +51,12 @@ func (t handshakeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	defer ws.Close()
+
+	if ws.Subprotocol() != subprotos[0] {
+		t.Logf("ws.Subprotocol() = %s, want %s", ws.Subprotocol(), subprotos[0])
+		return
+	}
+
 	for {
 		op, r, err := ws.NextReader()
 		if err != nil {
@@ -66,18 +81,19 @@ func (t handshakeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+var handshakeDialer = &Dialer{
+	Subprotocols:    []string{"p1", "p2"},
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
+
 func TestHandshake(t *testing.T) {
 	s := httptest.NewServer(handshakeHandler{t})
 	defer s.Close()
-	u, _ := url.Parse(s.URL)
-	c, err := net.Dial("tcp", u.Host)
+	ws, resp, err := handshakeDialer.Dial(httpToWs(s.URL), http.Header{"Origin": {s.URL}})
 	if err != nil {
 		t.Fatalf("Dial: %v", err)
 	}
-	ws, resp, err := NewClient(c, u, http.Header{"Origin": {s.URL}}, 1024, 1024)
-	if err != nil {
-		t.Fatalf("NewClient: %v", err)
-	}
 	defer ws.Close()
 
 	var sessionID string
@@ -90,24 +106,11 @@ func TestHandshake(t *testing.T) {
 		t.Error("Set-Cookie not received from the server.")
 	}
 
-	w, _ := ws.NextWriter(TextMessage)
-	io.WriteString(w, "HELLO")
-	w.Close()
-	ws.SetReadDeadline(time.Now().Add(1 * time.Second))
-	op, r, err := ws.NextReader()
-	if err != nil {
-		t.Fatalf("NextReader: %v", err)
-	}
-	if op != TextMessage {
-		t.Fatalf("op=%d, want %d", op, TextMessage)
-	}
-	b, err := ioutil.ReadAll(r)
-	if err != nil {
-		t.Fatalf("ReadAll: %v", err)
-	}
-	if string(b) != "HELLO" {
-		t.Fatalf("message=%s, want %s", b, "HELLO")
+	if ws.Subprotocol() != handshakeDialer.Subprotocols[0] {
+		t.Errorf("ws.Subprotocol() = %s, want %s", ws.Subprotocol(), handshakeDialer.Subprotocols[0])
 	}
+
+	sendRecv(t, ws)
 }
 
 type dialHandler struct {
@@ -142,9 +145,15 @@ func (t dialHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 func sendRecv(t *testing.T, ws *Conn) {
 	const message = "Hello World!"
+	if err := ws.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatalf("SetWriteDeadline: %v", err)
+	}
 	if err := ws.WriteMessage(TextMessage, []byte(message)); err != nil {
 		t.Fatalf("WriteMessage: %v", err)
 	}
+	if err := ws.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
 	_, p, err := ws.ReadMessage()
 	if err != nil {
 		t.Fatalf("ReadMessage: %v", err)

conn.go

@@ -103,8 +103,9 @@ func newMaskKey() [4]byte {
 
 // Conn represents a WebSocket connection.
 type Conn struct {
-	conn     net.Conn
-	isServer bool
+	conn        net.Conn
+	isServer    bool
+	subprotocol string
 
 	// Write fields
 	mu        chan bool // used as mutex to protect write to conn and closeSent
@@ -151,6 +152,11 @@ func newConn(conn net.Conn, isServer bool, readBufSize, writeBufSize int) *Conn
 	return c
 }
 
+// Subprotocol returns the negotiated protocol for the connection.
+func (c *Conn) Subprotocol() string {
+	return c.subprotocol
+}
+
 // Close closes the underlying network connection without sending or waiting for a close frame.
 func (c *Conn) Close() error {
 	return c.conn.Close()

server.go

@@ -83,6 +83,9 @@ func Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header,
 	}
 
 	c := newConn(netConn, true, readBufSize, writeBufSize)
+	if responseHeader != nil {
+		c.subprotocol = responseHeader.Get("Sec-Websocket-Protocol")
+	}
 
 	p := c.writeBuf[:0]
 	p = append(p, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "...)