Improve speed and add a benchmark

Author: nhooyr

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,77 @@
+package websocket_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"nhooyr.io/websocket"
+	"strings"
+	"testing"
+	"time"
+)
+
+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, "")
+
+	msg := strings.Repeat("2", 4096*16)
+	buf := make([]byte, len(msg))
+	b.SetBytes(int64(len(msg)))
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		w, err := c.Write(ctx, websocket.MessageText)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		_, err = io.WriteString(w, 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)
+		}
+
+		// TODO jank
+		_, err = r.Read(nil)
+		if err != io.EOF {
+			b.Fatalf("wtf %q", err)
+		}
+	}
+	b.StopTimer()
+	c.Close(websocket.StatusNormalClosure, "")
+}

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

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)
 	}