Merge pull request #56 from nhooyr/better

Fix remaining major bugs and UX issues

Author: nhooyr

.gitignore

@@ -1,2 +1,3 @@
 coverage.html
 wstest_reports
+websocket.test

accept.go

@@ -53,19 +53,19 @@ func AcceptInsecureOrigin() AcceptOption {
 
 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"))
+		err := xerrors.Errorf("websocket: protocol violation: Connection header %q does not contain Upgrade", r.Header.Get("Connection"))
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return err
 	}
 
 	if !headerValuesContainsToken(r.Header, "Upgrade", "WebSocket") {
-		err := xerrors.Errorf("websocket: protocol violation: Upgrade header does not contain websocket: %q", r.Header.Get("Upgrade"))
+		err := xerrors.Errorf("websocket: protocol violation: Upgrade header %q does not contain websocket", r.Header.Get("Upgrade"))
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return err
 	}
 
 	if r.Method != "GET" {
-		err := xerrors.Errorf("websocket: protocol violation: handshake request method is not GET: %q", r.Method)
+		err := xerrors.Errorf("websocket: protocol violation: handshake request method %q is not GET", r.Method)
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return err
 	}
@@ -88,7 +88,7 @@ func verifyClientRequest(w http.ResponseWriter, r *http.Request) error {
 // 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.
+// the AcceptInsecureOrigin option 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

bench_test.go

@@ -0,0 +1,87 @@
+package websocket_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"nhooyr.io/websocket"
+)
+
+func BenchmarkConn(b *testing.B) {
+	b.StopTimer()
+
+	s, closeFn := testServer(b, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c, err := websocket.Accept(w, r,
+			websocket.AcceptSubprotocols("echo"),
+		)
+		if err != nil {
+			b.Logf("server handshake failed: %+v", err)
+			return
+		}
+		echoLoop(r.Context(), c)
+	}))
+	defer closeFn()
+
+	wsURL := strings.Replace(s.URL, "http", "ws", 1)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*5)
+	defer cancel()
+
+	c, _, err := websocket.Dial(ctx, wsURL)
+	if err != nil {
+		b.Fatalf("failed to dial: %v", err)
+	}
+	defer c.Close(websocket.StatusInternalError, "")
+
+	runN := func(n int) {
+		b.Run(strconv.Itoa(n), func(b *testing.B) {
+			msg := []byte(strings.Repeat("2", n))
+			buf := make([]byte, len(msg))
+			b.SetBytes(int64(len(msg)))
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				w, err := c.Write(ctx, websocket.MessageText)
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				_, err = w.Write(msg)
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				err = w.Close()
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				_, r, err := c.Read(ctx)
+				if err != nil {
+					b.Fatal(err, b.N)
+				}
+
+				_, err = io.ReadFull(r, buf)
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+			b.StopTimer()
+		})
+	}
+
+	runN(32)
+	runN(128)
+	runN(512)
+	runN(1024)
+	runN(4096)
+	runN(16384)
+	runN(65536)
+	runN(131072)
+
+	c.Close(websocket.StatusNormalClosure, "")
+}

dial_test.go

@@ -7,6 +7,8 @@ import (
 )
 
 func Test_verifyServerHandshake(t *testing.T) {
+	t.Parallel()
+
 	testCases := []struct {
 		name     string
 		response func(w http.ResponseWriter)

example_test.go

@@ -34,10 +34,13 @@ func ExampleAccept_echo() {
 			if err != nil {
 				return err
 			}
-
 			r = io.LimitReader(r, 32768)
 
-			w := c.Write(ctx, typ)
+			w, err := c.Write(ctx, typ)
+			if err != nil {
+				return err
+			}
+
 			_, err = io.Copy(w, r)
 			if err != nil {
 				return err
@@ -76,7 +79,7 @@ func ExampleAccept() {
 			log.Printf("server handshake failed: %v", err)
 			return
 		}
-		defer c.Close(websocket.StatusInternalError, "") // TODO returning internal is incorect if its a timeout error.
+		defer c.Close(websocket.StatusInternalError, "")
 
 		jc := websocket.JSONConn{
 			Conn: c,

header.go

@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"math"
 
 	"golang.org/x/xerrors"
 )
@@ -55,7 +56,7 @@ func marshalHeader(h header) []byte {
 		panic(fmt.Sprintf("websocket: invalid header: negative length: %v", h.payloadLength))
 	case h.payloadLength <= 125:
 		b[1] = byte(h.payloadLength)
-	case h.payloadLength <= 1<<16:
+	case h.payloadLength <= math.MaxUint16:
 		b[1] = 126
 		b = b[:len(b)+2]
 		binary.BigEndian.PutUint16(b[len(b)-2:], uint16(h.payloadLength))
@@ -105,10 +106,8 @@ func readHeader(r io.Reader) (header, error) {
 	case payloadLength < 126:
 		h.payloadLength = int64(payloadLength)
 	case payloadLength == 126:
-		h.payloadLength = 126
 		extra += 2
 	case payloadLength == 127:
-		h.payloadLength = 127
 		extra += 8
 	}
 

header_test.go

@@ -3,6 +3,7 @@ package websocket
 import (
 	"bytes"
 	"math/rand"
+	"strconv"
 	"testing"
 	"time"
 
@@ -36,10 +37,38 @@ func TestHeader(t *testing.T) {
 			t.Fatalf("unexpected error value: %+v", err)
 		}
 	})
+
+	t.Run("lengths", func(t *testing.T) {
+		t.Parallel()
+
+		lengths := []int{
+			124,
+			125,
+			126,
+			4096,
+			16384,
+			65535,
+			65536,
+			65537,
+			131072,
+		}
+
+		for _, n := range lengths {
+			n := n
+			t.Run(strconv.Itoa(n), func(t *testing.T) {
+				t.Parallel()
+
+				testHeader(t, header{
+					payloadLength: int64(n),
+				})
+			})
+		}
+	})
+
 	t.Run("fuzz", func(t *testing.T) {
 		t.Parallel()
 
-		for i := 0; i < 1000; i++ {
+		for i := 0; i < 10000; i++ {
 			h := header{
 				fin:    randBool(),
 				rsv1:   randBool(),
@@ -55,20 +84,24 @@ func TestHeader(t *testing.T) {
 				rand.Read(h.maskKey[:])
 			}
 
-			b := marshalHeader(h)
-			r := bytes.NewReader(b)
-			h2, err := readHeader(r)
-			if err != nil {
-				t.Logf("header: %#v", h)
-				t.Logf("bytes: %b", b)
-				t.Fatalf("failed to read header: %v", err)
-			}
-
-			if !cmp.Equal(h, h2, cmp.AllowUnexported(header{})) {
-				t.Logf("header: %#v", h)
-				t.Logf("bytes: %b", b)
-				t.Fatalf("parsed and read header differ: %v", cmp.Diff(h, h2, cmp.AllowUnexported(header{})))
-			}
+			testHeader(t, h)
 		}
 	})
 }
+
+func testHeader(t *testing.T, h header) {
+	b := marshalHeader(h)
+	r := bytes.NewReader(b)
+	h2, err := readHeader(r)
+	if err != nil {
+		t.Logf("header: %#v", h)
+		t.Logf("bytes: %b", b)
+		t.Fatalf("failed to read header: %v", err)
+	}
+
+	if !cmp.Equal(h, h2, cmp.AllowUnexported(header{})) {
+		t.Logf("header: %#v", h)
+		t.Logf("bytes: %b", b)
+		t.Fatalf("parsed and read header differ: %v", cmp.Diff(h, h2, cmp.AllowUnexported(header{})))
+	}
+}

json.go

@@ -22,7 +22,7 @@ func (jc JSONConn) Read(ctx context.Context, v interface{}) error {
 	return nil
 }
 
-func (jc *JSONConn) read(ctx context.Context, v interface{}) error {
+func (jc JSONConn) read(ctx context.Context, v interface{}) error {
 	typ, r, err := jc.Conn.Read(ctx)
 	if err != nil {
 		return err
@@ -53,10 +53,13 @@ func (jc JSONConn) Write(ctx context.Context, v interface{}) error {
 }
 
 func (jc JSONConn) write(ctx context.Context, v interface{}) error {
-	w := jc.Conn.Write(ctx, MessageText)
+	w, err := jc.Conn.Write(ctx, MessageText)
+	if err != nil {
+		return xerrors.Errorf("failed to get message writer: %w", err)
+	}
 
 	e := json.NewEncoder(w)
-	err := e.Encode(v)
+	err = e.Encode(v)
 	if err != nil {
 		return xerrors.Errorf("failed to encode json: %w", err)
 	}

statuscode.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"math/bits"
-	"unicode/utf8"
 
 	"golang.org/x/xerrors"
 )
@@ -54,6 +53,12 @@ func (ce CloseError) Error() string {
 }
 
 func parseClosePayload(p []byte) (CloseError, error) {
+	if len(p) == 0 {
+		return CloseError{
+			Code: StatusNoStatusRcvd,
+		}, nil
+	}
+
 	if len(p) < 2 {
 		return CloseError{}, fmt.Errorf("close payload too small, cannot even contain the 2 byte status code")
 	}
@@ -63,9 +68,6 @@ func parseClosePayload(p []byte) (CloseError, error) {
 		Reason: string(p[2:]),
 	}
 
-	if !utf8.ValidString(ce.Reason) {
-		return CloseError{}, xerrors.Errorf("invalid utf-8: %q", ce.Reason)
-	}
 	if !validWireCloseCode(ce.Code) {
 		return CloseError{}, xerrors.Errorf("invalid code %v", ce.Code)
 	}