diff --git a/client.go b/client.go index c3f7825..a4090e1 100644 --- a/client.go +++ b/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}, diff --git a/client_server_test.go b/client_server_test.go index 5ddeaad..50ee84d 100644 --- a/client_server_test.go +++ b/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) diff --git a/conn.go b/conn.go index e37970d..3c841b5 100644 --- a/conn.go +++ b/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() diff --git a/server.go b/server.go index ebc43dd..10470dd 100644 --- a/server.go +++ b/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: "...)