Merge pull request #48 from nhooyr/fixes

Improvements

Author: nhooyr

README.md

@@ -1,7 +1,6 @@
 # websocket
 
 [![GoDoc](https://godoc.org/nhooyr.io/websocket?status.svg)](https://godoc.org/nhooyr.io/websocket)
-[![Codecov](https://img.shields.io/codecov/c/github/nhooyr/websocket.svg)](https://codecov.io/gh/nhooyr/websocket)
 [![GitHub release](https://img.shields.io/github/release-pre/nhooyr/websocket.svg)](https://github.com/nhooyr/websocket/releases)
 
 websocket is a minimal and idiomatic WebSocket library for Go.
@@ -23,6 +22,7 @@ go get nhooyr.io/websocket
 - net/http is used for WebSocket dials and upgrades
 - Passes the [autobahn-testsuite](https://github.com/crossbario/autobahn-testsuite)
 - JSON helpers
+- Thoroughly testing
 
 ## Roadmap
 
@@ -112,7 +112,9 @@ c.Close(websocket.StatusNormalClosure, "")
 
 While I believe nhooyr/websocket has a better API than existing libraries, 
 both gorilla/websocket and gobwas/ws were extremely useful in implementing the
-WebSocket protocol correctly so big thanks to the authors of both.
+WebSocket protocol correctly so big thanks to the authors of both. In particular,
+I made sure to go through the issue tracker of gorilla/websocket to make sure
+I implemented details correctly.
 
 ### gorilla/websocket
 

accept.go

@@ -23,8 +23,8 @@ func (o acceptSubprotocols) acceptOption() {}
 
 // AcceptSubprotocols list the subprotocols that Accept will negotiate with a client.
 // The first protocol that a client supports will be negotiated.
-// Pass "" as a subprotocol if you would like to allow the default protocol along with
-// specific subprotocols.
+// The empty protocol will always be negotiated as per RFC 6455. If you would like to
+// reject it, close the connection is c.Subprotocol() == "".
 func AcceptSubprotocols(subprotocols ...string) AcceptOption {
 	return acceptSubprotocols(subprotocols)
 }
@@ -42,7 +42,7 @@ func (o acceptOrigins) acceptOption() {}
 // See https://stackoverflow.com/a/37837709/4283659
 // You can use a * for wildcards.
 func AcceptOrigins(origins ...string) AcceptOption {
-	return AcceptOrigins(origins...)
+	return acceptOrigins(origins)
 }
 
 // Accept accepts a WebSocket handshake from a client and upgrades the
@@ -118,7 +118,7 @@ func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn
 
 	netConn, brw, err := hj.Hijack()
 	if err != nil {
-		err = xerrors.Errorf("websocket: failed to hijack connection: %v", err)
+		err = xerrors.Errorf("websocket: failed to hijack connection: %w", err)
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return nil, err
 	}
@@ -135,7 +135,7 @@ func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn
 }
 
 func selectSubprotocol(w http.ResponseWriter, r *http.Request, subprotocols []string) {
-	clientSubprotocols := strings.Split(r.Header.Get("Sec-WebSocket-Protocol"), "\n")
+	clientSubprotocols := strings.Split(r.Header.Get("Sec-WebSocket-Protocol"), ",")
 	for _, sp := range subprotocols {
 		for _, cp := range clientSubprotocols {
 			if sp == strings.TrimSpace(cp) {
@@ -165,12 +165,12 @@ func authenticateOrigin(r *http.Request, origins []string) error {
 	}
 	u, err := url.Parse(origin)
 	if err != nil {
-		return xerrors.Errorf("failed to parse Origin header %q: %v", origin, err)
+		return xerrors.Errorf("failed to parse Origin header %q: %w", origin, err)
 	}
 	for _, o := range origins {
-		if u.Host == o {
+		if strings.EqualFold(u.Host, o) {
 			return nil
 		}
 	}
-	return xerrors.New("request origin is not authorized")
+	return xerrors.Errorf("request origin %q is not authorized", r.Header.Get("Origin"))
 }

datatype.go

@@ -1,9 +1,10 @@
 package websocket
 
 // DataType represents the Opcode of a WebSocket data frame.
-//go:generate go run golang.org/x/tools/cmd/stringer -type=DataType
 type DataType int
 
+//go:generate go run golang.org/x/tools/cmd/stringer -type=DataType
+
 // DataType constants.
 const (
 	Text   DataType = DataType(opText)

dial.go

@@ -74,19 +74,19 @@ func Dial(ctx context.Context, u string, opts ...DialOption) (_ *Conn, _ *http.R
 
 	parsedURL, err := url.Parse(u)
 	if err != nil {
-		return nil, nil, xerrors.Errorf("failed to parse websocket url: %v", err)
+		return nil, nil, xerrors.Errorf("failed to parse websocket url: %w", err)
 	}
 
 	switch parsedURL.Scheme {
-	case "ws", "http":
+	case "ws":
 		parsedURL.Scheme = "http"
-	case "wss", "https":
+	case "wss":
 		parsedURL.Scheme = "https"
 	default:
 		return nil, nil, xerrors.Errorf("unknown scheme in url: %q", parsedURL.Scheme)
 	}
 
-	req, _ := http.NewRequest("GET", u, nil)
+	req, _ := http.NewRequest("GET", parsedURL.String(), nil)
 	req = req.WithContext(ctx)
 	req.Header = header
 	req.Header.Set("Connection", "Upgrade")
@@ -113,7 +113,7 @@ func Dial(ctx context.Context, u string, opts ...DialOption) (_ *Conn, _ *http.R
 	}()
 
 	if resp.StatusCode != http.StatusSwitchingProtocols {
-		return nil, resp, xerrors.Errorf("websocket: expected status code %v but got %v", http.StatusSwitchingProtocols)
+		return nil, resp, xerrors.Errorf("websocket: expected status code %v but got %v", http.StatusSwitchingProtocols, resp.StatusCode)
 	}
 
 	if !httpguts.HeaderValuesContainsToken(resp.Header["Connection"], "Upgrade") {

json.go

@@ -12,17 +12,20 @@ import (
 func ReadJSON(ctx context.Context, c *Conn, v interface{}) error {
 	typ, r, err := c.ReadMessage(ctx)
 	if err != nil {
-		return xerrors.Errorf("failed to read json: %v", err)
+		return xerrors.Errorf("failed to read json: %w", err)
 	}
 
 	if typ != websocket.TextFrame {
 		return xerrors.Errorf("unexpected frame type for json (expected TextFrame): %v", typ)
 	}
 
+	r.Limit(131072)
+	r.SetContext(ctx)
+
 	d := json.NewDecoder(r)
 	err = d.Decode(v)
 	if err != nil {
-		return xerrors.Errorf("failed to read json: %v", err)
+		return xerrors.Errorf("failed to read json: %w", err)
 	}
 	return nil
 }
@@ -31,14 +34,16 @@ func ReadJSON(ctx context.Context, c *Conn, v interface{}) error {
 func WriteJSON(ctx context.Context, c *Conn, v interface{}) error {
 	w := c.MessageWriter(websocket.TextFrame)
 	w.SetContext(ctx)
+
 	e := json.NewEncoder(w)
 	err := e.Encode(v)
 	if err != nil {
-		return xerrors.Errorf("failed to write json: %v", err)
+		return xerrors.Errorf("failed to write json: %w", err)
 	}
+
 	err = w.Close()
 	if err != nil {
-		return xerrors.Errorf("failed to write json: %v", err)
+		return xerrors.Errorf("failed to write json: %w", err)
 	}
 	return nil
 }

opcode.go

@@ -1,16 +1,22 @@
 package websocket
 
 // opcode represents a WebSocket Opcode.
-//go:generate go run golang.org/x/tools/cmd/stringer -type=opcode
 type opcode int
 
+//go:generate go run golang.org/x/tools/cmd/stringer -type=opcode
+
 // opcode constants.
 const (
 	opContinuation opcode = iota
 	opText
 	opBinary
 	// 3 - 7 are reserved for further non-control frames.
-	opClose opcode = 8 + iota - 3
+	_
+	_
+	_
+	_
+	_
+	opClose
 	opPing
 	opPong
 	// 11-16 are reserved for further control frames.

statuscode.go

@@ -11,18 +11,19 @@ import (
 )
 
 // StatusCode represents a WebSocket status code.
-//go:generate go run golang.org/x/tools/cmd/stringer -type=StatusCode
 type StatusCode int
 
+//go:generate go run golang.org/x/tools/cmd/stringer -type=StatusCode
+
 // These codes were retrieved from:
 // https://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
 const (
 	StatusNormalClosure StatusCode = 1000 + iota
 	StatusGoingAway
 	StatusProtocolError
 	StatusUnsupportedData
-	// 1004 is reserved.
-	StatusNoStatusRcvd StatusCode = 1005 + iota - 4
+	_ // 1004 is reserved.
+	StatusNoStatusRcvd
 	StatusAbnormalClosure
 	StatusInvalidFramePayloadData
 	StatusPolicyViolation

websocket.go

@@ -58,7 +58,7 @@ func (c *Conn) getCloseErr() error {
 
 func (c *Conn) close(err error) {
 	if err != nil {
-		err = xerrors.Errorf("websocket: connection broken: %v", err)
+		err = xerrors.Errorf("websocket: connection broken: %w", err)
 	}
 
 	c.closeOnce.Do(func() {
@@ -102,20 +102,20 @@ func (c *Conn) writeFrame(h header, p []byte) {
 	b2 := marshalHeader(h)
 	_, err := c.bw.Write(b2)
 	if err != nil {
-		c.close(xerrors.Errorf("failed to write to connection: %v", err))
+		c.close(xerrors.Errorf("failed to write to connection: %w", err))
 		return
 	}
 
 	_, err = c.bw.Write(p)
 	if err != nil {
-		c.close(xerrors.Errorf("failed to write to connection: %v", err))
+		c.close(xerrors.Errorf("failed to write to connection: %w", err))
 		return
 	}
 
 	if h.opcode.controlOp() {
 		err := c.bw.Flush()
 		if err != nil {
-			c.close(xerrors.Errorf("failed to write to connection: %v", err))
+			c.close(xerrors.Errorf("failed to write to connection: %w", err))
 			return
 		}
 	}
@@ -139,7 +139,11 @@ messageLoop:
 				masked:        c.client,
 			}
 			c.writeFrame(h, control.payload)
-			c.writeDone <- struct{}{}
+			select {
+			case <-c.closed:
+				return
+			case c.writeDone <- struct{}{}:
+			}
 			continue
 		}
 
@@ -176,7 +180,7 @@ messageLoop:
 				if !ok {
 					err := c.bw.Flush()
 					if err != nil {
-						c.close(xerrors.Errorf("failed to write to connection: %v", err))
+						c.close(xerrors.Errorf("failed to write to connection: %w", err))
 						return
 					}
 				}
@@ -210,7 +214,7 @@ func (c *Conn) handleControl(h header) {
 	b := make([]byte, h.payloadLength)
 	_, err := io.ReadFull(c.br, b)
 	if err != nil {
-		c.close(xerrors.Errorf("failed to read control frame payload: %v", err))
+		c.close(xerrors.Errorf("failed to read control frame payload: %w", err))
 		return
 	}
 
@@ -226,7 +230,7 @@ func (c *Conn) handleControl(h header) {
 		if len(b) > 0 {
 			code, reason, err := parseClosePayload(b)
 			if err != nil {
-				c.close(xerrors.Errorf("read invalid close payload: %v", err))
+				c.close(xerrors.Errorf("read invalid close payload: %w", err))
 				return
 			}
 			c.Close(code, reason)
@@ -247,7 +251,7 @@ func (c *Conn) readLoop() {
 	for {
 		h, err := readHeader(c.br)
 		if err != nil {
-			c.close(xerrors.Errorf("failed to read header: %v", err))
+			c.close(xerrors.Errorf("failed to read header: %w", err))
 			return
 		}
 
@@ -280,6 +284,7 @@ func (c *Conn) readLoop() {
 				return
 			}
 		default:
+			// TODO send back protocol violation message or figure out what RFC wants.
 			c.close(xerrors.Errorf("unexpected opcode in header: %#v", h))
 			return
 		}
@@ -298,7 +303,7 @@ func (c *Conn) readLoop() {
 
 				_, err = io.ReadFull(c.br, b)
 				if err != nil {
-					c.close(xerrors.Errorf("failed to read from connection: %v", err))
+					c.close(xerrors.Errorf("failed to read from connection: %w", err))
 					return
 				}
 				left -= int64(len(b))
@@ -355,14 +360,14 @@ func (c *Conn) MessageWriter(dataType DataType) *MessageWriter {
 func (c *Conn) ReadMessage(ctx context.Context) (DataType, *MessageReader, error) {
 	select {
 	case <-c.closed:
-		return 0, nil, xerrors.Errorf("failed to read message: %v", c.getCloseErr())
+		return 0, nil, xerrors.Errorf("failed to read message: %w", c.getCloseErr())
 	case opcode := <-c.read:
 		return DataType(opcode), &MessageReader{
 			ctx: context.Background(),
 			c:   c,
 		}, nil
 	case <-ctx.Done():
-		return 0, nil, xerrors.Errorf("failed to read message: %v", ctx.Err())
+		return 0, nil, xerrors.Errorf("failed to read message: %w", ctx.Err())
 	}
 }
 
@@ -481,8 +486,14 @@ func (w *MessageWriter) Close() error {
 		}
 	}
 	close(w.c.writeBytes)
-	<-w.c.writeDone
-	return nil
+	select {
+	case <-w.c.closed:
+		return w.c.getCloseErr()
+	case <-w.ctx.Done():
+		return w.ctx.Err()
+	case <-w.c.writeDone:
+		return nil
+	}
 }
 
 // MessageReader enables reading a data frame from the WebSocket connection.
@@ -501,6 +512,14 @@ func (r *MessageReader) SetContext(ctx context.Context) {
 }
 
 // Limit limits the number of bytes read by the reader.
+//
+// Why not use io.LimitReader? io.LimitReader returns a io.EOF
+// after the limit bytes which means its not possible to tell
+// whether the message has been read or a limit has been hit.
+// This results in unclear error and log messages.
+// This function will cause the connection to be closed if the limit is hit
+// with a close reason explaining the error and also an error
+// indicating the limit was hit.
 func (r *MessageReader) Limit(bytes int) {
 	r.limit = bytes
 }