Improve API and test coverage

README.md

@@ -45,13 +45,17 @@ fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	}
 	defer c.Close(websocket.StatusInternalError, "")
 
+	jc := websocket.JSONConn{
+		Conn: c,
+	}
+
 	ctx, cancel := context.WithTimeout(r.Context(), time.Second*10)
 	defer cancel()
 
 	v := map[string]interface{}{
 		"my_field": "foo",
 	}
-	err = websocket.WriteJSON(ctx, c, v)
+	err = jc.Write(ctx, v)
 	if err != nil {
 		log.Printf("failed to write json: %v", err)
 		return
@@ -73,7 +77,7 @@ For a production quality example that shows off the low level API, see the [echo
 
 ```go
 ctx := context.Background()
-ctx, cancel := context.WithTimeout(ctx, time.Second*10)
+ctx, cancel := context.WithTimeout(ctx, time.Minute)
 defer cancel()
 
 c, _, err := websocket.Dial(ctx, "ws://localhost:8080",
@@ -84,8 +88,12 @@ if err != nil {
 }
 defer c.Close(websocket.StatusInternalError, "")
 
+jc := websocket.JSONConn{
+	Conn: c,
+}
+
 var v interface{}
-err = websocket.ReadJSON(ctx, c, v)
+err = jc.Read(ctx, v)
 if err != nil {
 	log.Fatalf("failed to read json: %v", err)
 }

accept.go

@@ -4,6 +4,7 @@ import (
 	"crypto/sha1"
 	"encoding/base64"
 	"net/http"
+	"net/textproto"
 	"net/url"
 	"strings"
 
@@ -45,56 +46,65 @@ func AcceptOrigins(origins ...string) AcceptOption {
 	return acceptOrigins(origins)
 }
 
-// Accept accepts a WebSocket handshake from a client and upgrades the
-// the connection to WebSocket.
-// Accept will reject the handshake if the Origin is not the same as the Host unless
-// InsecureAcceptOrigin is passed.
-// Accept uses w to write the handshake response so the timeouts on the http.Server apply.
-func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn, error) {
-	var subprotocols []string
-	origins := []string{r.Host}
-	for _, opt := range opts {
-		switch opt := opt.(type) {
-		case acceptOrigins:
-			origins = []string(opt)
-		case acceptSubprotocols:
-			subprotocols = []string(opt)
-		}
-	}
-
-	if !httpguts.HeaderValuesContainsToken(r.Header["Connection"], "Upgrade") {
+func verifyClientRequest(w http.ResponseWriter, r *http.Request) error {
+	if !headerValuesContainsToken(r.Header, "Connection", "Upgrade") {
 		err := xerrors.Errorf("websocket: protocol violation: Connection header does not contain Upgrade: %q", r.Header.Get("Connection"))
 		http.Error(w, err.Error(), http.StatusBadRequest)
-		return nil, err
+		return err
 	}
 
-	if !httpguts.HeaderValuesContainsToken(r.Header["Upgrade"], "websocket") {
+	if !headerValuesContainsToken(r.Header, "Upgrade", "WebSocket") {
 		err := xerrors.Errorf("websocket: protocol violation: Upgrade header does not contain websocket: %q", r.Header.Get("Upgrade"))
 		http.Error(w, err.Error(), http.StatusBadRequest)
-		return nil, err
+		return err
 	}
 
 	if r.Method != "GET" {
 		err := xerrors.Errorf("websocket: protocol violation: handshake request method is not GET: %q", r.Method)
 		http.Error(w, err.Error(), http.StatusBadRequest)
-		return nil, err
+		return err
 	}
 
 	if r.Header.Get("Sec-WebSocket-Version") != "13" {
 		err := xerrors.Errorf("websocket: unsupported protocol version: %q", r.Header.Get("Sec-WebSocket-Version"))
 		http.Error(w, err.Error(), http.StatusBadRequest)
-		return nil, err
+		return err
 	}
 
 	if r.Header.Get("Sec-WebSocket-Key") == "" {
 		err := xerrors.New("websocket: protocol violation: missing Sec-WebSocket-Key")
 		http.Error(w, err.Error(), http.StatusBadRequest)
+		return err
+	}
+
+	return nil
+}
+
+// Accept accepts a WebSocket handshake from a client and upgrades the
+// the connection to WebSocket.
+// Accept will reject the handshake if the Origin is not the same as the Host unless
+// InsecureAcceptOrigin is passed.
+// Accept uses w to write the handshake response so the timeouts on the http.Server apply.
+func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn, error) {
+	var subprotocols []string
+	origins := []string{r.Host}
+	for _, opt := range opts {
+		switch opt := opt.(type) {
+		case acceptOrigins:
+			origins = []string(opt)
+		case acceptSubprotocols:
+			subprotocols = []string(opt)
+		}
+	}
+
+	err := verifyClientRequest(w, r)
+	if err != nil {
 		return nil, err
 	}
 
 	origins = append(origins, r.Host)
 
-	err := authenticateOrigin(r, origins)
+	err = authenticateOrigin(r, origins)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusForbidden)
 		return nil, err
@@ -112,7 +122,10 @@ func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn
 
 	handleKey(w, r)
 
-	selectSubprotocol(w, r, subprotocols)
+	subproto := selectSubprotocol(r, subprotocols)
+	if subproto != "" {
+		w.Header().Set("Sec-WebSocket-Protocol", subproto)
+	}
 
 	w.WriteHeader(http.StatusSwitchingProtocols)
 
@@ -134,16 +147,18 @@ func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn
 	return c, nil
 }
 
-func selectSubprotocol(w http.ResponseWriter, r *http.Request, subprotocols []string) {
-	clientSubprotocols := strings.Split(r.Header.Get("Sec-WebSocket-Protocol"), ",")
+func headerValuesContainsToken(h http.Header, key, val string) bool {
+	key = textproto.CanonicalMIMEHeaderKey(key)
+	return httpguts.HeaderValuesContainsToken(h[key], val)
+}
+
+func selectSubprotocol(r *http.Request, subprotocols []string) string {
 	for _, sp := range subprotocols {
-		for _, cp := range clientSubprotocols {
-			if sp == strings.TrimSpace(cp) {
-				w.Header().Set("Sec-WebSocket-Protocol", sp)
-				return
-			}
+		if headerValuesContainsToken(r.Header, "Sec-WebSocket-Protocol", sp) {
+			return sp
 		}
 	}
+	return ""
 }
 
 var keyGUID = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11")

accept_test.go

@@ -0,0 +1,191 @@
+package websocket
+
+import (
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+func Test_verifyClientHandshake(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name    string
+		method  string
+		h       map[string]string
+		success bool
+	}{
+		{
+			name: "badConnection",
+			h: map[string]string{
+				"Connection": "notUpgrade",
+			},
+		},
+		{
+			name: "badUpgrade",
+			h: map[string]string{
+				"Connection": "Upgrade",
+				"Upgrade":    "notWebSocket",
+			},
+		},
+		{
+			name:   "badMethod",
+			method: "POST",
+			h: map[string]string{
+				"Connection": "Upgrade",
+				"Upgrade":    "websocket",
+			},
+		},
+		{
+			name: "badWebSocketVersion",
+			h: map[string]string{
+				"Connection":            "Upgrade",
+				"Upgrade":               "websocket",
+				"Sec-WebSocket-Version": "14",
+			},
+		},
+		{
+			name: "badWebSocketKey",
+			h: map[string]string{
+				"Connection":            "Upgrade",
+				"Upgrade":               "websocket",
+				"Sec-WebSocket-Version": "13",
+				"Sec-WebSocket-Key":     "",
+			},
+		},
+		{
+			name: "success",
+			h: map[string]string{
+				"Connection":            "Upgrade",
+				"Upgrade":               "websocket",
+				"Sec-WebSocket-Version": "13",
+				"Sec-WebSocket-Key":     "meow123",
+			},
+			success: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			w := httptest.NewRecorder()
+			r := httptest.NewRequest(tc.method, "/", nil)
+
+			for k, v := range tc.h {
+				r.Header.Set(k, v)
+			}
+
+			err := verifyClientRequest(w, r)
+			if (err == nil) != tc.success {
+				t.Fatalf("unexpected error value: %+v", err)
+			}
+		})
+	}
+}
+
+func Test_selectSubprotocol(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name            string
+		clientProtocols []string
+		serverProtocols []string
+		negotiated      string
+	}{
+		{
+			name:            "empty",
+			clientProtocols: nil,
+			serverProtocols: nil,
+			negotiated:      "",
+		},
+		{
+			name:            "basic",
+			clientProtocols: []string{"echo", "echo2"},
+			serverProtocols: []string{"echo2", "echo"},
+			negotiated:      "echo2",
+		},
+		{
+			name:            "none",
+			clientProtocols: []string{"echo", "echo3"},
+			serverProtocols: []string{"echo2", "echo4"},
+			negotiated:      "",
+		},
+		{
+			name:            "fallback",
+			clientProtocols: []string{"echo", "echo3"},
+			serverProtocols: []string{"echo2", "echo3"},
+			negotiated:      "echo3",
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			r := httptest.NewRequest("GET", "/", nil)
+			r.Header.Set("Sec-WebSocket-Protocol", strings.Join(tc.clientProtocols, ","))
+
+			negotiated := selectSubprotocol(r, tc.serverProtocols)
+			if tc.negotiated != negotiated {
+				t.Fatalf("expected %q but got %q", tc.negotiated, negotiated)
+			}
+		})
+	}
+}
+
+func Test_authenticateOrigin(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name              string
+		origin            string
+		authorizedOrigins []string
+		success           bool
+	}{
+		{
+			name:    "none",
+			success: true,
+		},
+		{
+			name:    "invalid",
+			origin:  "$#)(*)$#@*$(#@*$)#@*%)#(@*%)#(@%#@$#@$#$#@$#@}{}{}",
+			success: false,
+		},
+		{
+			name:              "unauthorized",
+			origin:            "https://example.com",
+			authorizedOrigins: []string{"example1.com"},
+			success:           false,
+		},
+		{
+			name:              "authorized",
+			origin:            "https://example.com",
+			authorizedOrigins: []string{"example.com"},
+			success:           true,
+		},
+		{
+			name:              "authorizedCaseInsensitive",
+			origin:            "https://examplE.com",
+			authorizedOrigins: []string{"example.com"},
+			success:           true,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			r := httptest.NewRequest("GET", "/", nil)
+			r.Header.Set("Origin", tc.origin)
+
+			err := authenticateOrigin(r, tc.authorizedOrigins)
+			if (err == nil) != tc.success {
+				t.Fatalf("unexpected error value: %+v", err)
+			}
+		})
+	}
+}

datatype.go

@@ -7,6 +7,6 @@ type DataType int
 
 // DataType constants.
 const (
-	Text   DataType = DataType(opText)
-	Binary DataType = DataType(opBinary)
+	DataText   DataType = DataType(opText)
+	DataBinary DataType = DataType(opBinary)
 )