Saving WIP rewrite

Author: nhooyr

close.go

@@ -130,10 +130,10 @@ func (ce CloseError) bytes() []byte {
 }
 
 func (ce CloseError) bytesErr() ([]byte, error) {
-	// TODO move check into frame write
-	// if len(ce.Reason) > wsframe.MaxControlFramePayload-2 {
-	// 	return nil, fmt.Errorf("reason string max is %v but got %q with length %v", wsframe.MaxControlFramePayload-2, ce.Reason, len(ce.Reason))
-	// }
+	const maxReason = maxControlPayload-2
+	if len(ce.Reason) > maxReason {
+		return nil, fmt.Errorf("reason string max is %v but got %q with length %v", maxReason, ce.Reason, len(ce.Reason))
+	}
 	if !validWireCloseCode(ce.Code) {
 		return nil, fmt.Errorf("status code %v cannot be set", ce.Code)
 	}

conn.go

@@ -116,6 +116,7 @@ func (c *Conn) close(err error) {
 	// closeErr.
 	c.closer.Close()
 
+	// TODO acquire looks here in a new goroutine to ensure we don't get blocked.
 	c.r.close()
 	c.w.close()
 }
@@ -450,9 +451,6 @@ func (c *Conn) exportedClose(code StatusCode, reason string, handshake bool) err
 		Reason: reason,
 	}
 
-	// This function also will not wait for a close frame from the peer like the RFC
-	// wants because that makes no sense and I don't think anyone actually follows that.
-	// Definitely worth seeing what popular browsers do later.
 	p := ce.bytes()
 	return c.writeClose(p, fmt.Errorf("sent close: %w", ce), handshake)
 }

internal/wsframe/mask.go renamed to frame.go

@@ -1,11 +1,189 @@
-package wsframe
+package websocket
 
 import (
+	"bufio"
 	"encoding/binary"
+	"fmt"
+	"math"
 	"math/bits"
 )
 
-// Mask applies the WebSocket masking algorithm to p
+// opcode represents a WebSocket opcode.
+type opcode int
+
+// List at https://tools.ietf.org/html/rfc6455#section-11.8.
+const (
+	opContinuation opcode = iota
+	opText
+	opBinary
+	// 3 - 7 are reserved for further non-control frames.
+	_
+	_
+	_
+	_
+	_
+	opClose
+	opPing
+	opPong
+	// 11-16 are reserved for further control frames.
+)
+
+// header represents a WebSocket frame header.
+// See https://tools.ietf.org/html/rfc6455#section-5.2.
+type header struct {
+	fin    bool
+	rsv1   bool
+	rsv2   bool
+	rsv3   bool
+	opcode opcode
+
+	payloadLength int64
+
+	masked  bool
+	maskKey uint32
+}
+
+// ReadFrameHeader reads a header from the reader.
+// See https://tools.ietf.org/html/rfc6455#section-5.2.
+func ReadFrameHeader(r *bufio.Reader) (header, error) {
+	h, err := readFrameHeader(r)
+	if err != nil {
+		return header{}, fmt.Errorf("failed to read frame header: %w", err)
+	}
+	return h, nil
+}
+
+func readFrameHeader(r *bufio.Reader) (header, error) {
+	b, err := r.ReadByte()
+	if err != nil {
+		return header{}, err
+	}
+
+	var h header
+	h.fin = b&(1<<7) != 0
+	h.rsv1 = b&(1<<6) != 0
+	h.rsv2 = b&(1<<5) != 0
+	h.rsv3 = b&(1<<4) != 0
+
+	h.opcode = opcode(b & 0xf)
+
+	b, err = r.ReadByte()
+	if err != nil {
+		return header{}, err
+	}
+
+	h.masked = b&(1<<7) != 0
+
+	payloadLength := b &^ (1 << 7)
+	switch {
+	case payloadLength < 126:
+		h.payloadLength = int64(payloadLength)
+	case payloadLength == 126:
+		var pl uint16
+		err = binary.Read(r, binary.BigEndian, &pl)
+		h.payloadLength = int64(pl)
+	case payloadLength == 127:
+		err = binary.Read(r, binary.BigEndian, &h.payloadLength)
+	}
+	if err != nil {
+		return header{}, err
+	}
+
+	if h.masked {
+		err = binary.Read(r, binary.LittleEndian, &h.maskKey)
+		if err != nil {
+			return header{}, err
+		}
+	}
+
+	return h, nil
+}
+
+// maxControlPayload is the maximum length of a control frame payload.
+// See https://tools.ietf.org/html/rfc6455#section-5.5.
+const maxControlPayload = 125
+
+// ParseClosePayload parses the bytes in p as a close frame payload,
+// returning the status code and reason.
+func ParseClosePayload(p []byte) (uint16, string, error) {
+	if len(p) < 2 {
+		return 0, "", fmt.Errorf("close payload %q too small, cannot even contain the 2 byte status code", p)
+	}
+
+	return binary.BigEndian.Uint16(p), string(p[2:]), nil
+}
+
+// WriteHeader writes the bytes of the header to w.
+// See https://tools.ietf.org/html/rfc6455#section-5.2
+func WriteHeader(h header, w *bufio.Writer) error {
+	err := writeHeader(h, w)
+	if err != nil {
+		return fmt.Errorf("failed to write frame header: %w", err)
+	}
+	return nil
+}
+
+func writeHeader(h header, w *bufio.Writer) error {
+	var b byte
+	if h.fin {
+		b |= 1 << 7
+	}
+	if h.rsv1 {
+		b |= 1 << 6
+	}
+	if h.rsv2 {
+		b |= 1 << 5
+	}
+	if h.rsv3 {
+		b |= 1 << 4
+	}
+
+	b |= byte(h.opcode)
+
+	err := w.WriteByte(b)
+	if err != nil {
+		return err
+	}
+
+	lengthByte := byte(0)
+	if h.masked {
+		lengthByte |= 1 << 7
+	}
+
+	switch {
+	case h.payloadLength > math.MaxUint16:
+		lengthByte |= 127
+	case h.payloadLength > 125:
+		lengthByte |= 126
+	case h.payloadLength >= 0:
+		lengthByte |= byte(h.payloadLength)
+	}
+	err = w.WriteByte(lengthByte)
+	if err != nil {
+		return err
+	}
+
+	switch {
+	case h.payloadLength > math.MaxUint16:
+		err = binary.Write(w, binary.BigEndian, h.payloadLength)
+	case h.payloadLength > 125:
+		err = binary.Write(w, binary.BigEndian, uint16(h.payloadLength))
+	}
+	if err != nil {
+		return err
+	}
+
+	if h.masked {
+		err = binary.Write(w, binary.LittleEndian, h.maskKey)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// mask applies the WebSocket masking algorithm to p
 // with the given key.
 // See https://tools.ietf.org/html/rfc6455#section-5.3
 //
@@ -16,7 +194,7 @@ import (
 // to be in little endian.
 //
 // See https://github.com/golang/go/issues/31586
-func Mask(key uint32, b []byte) uint32 {
+func mask(key uint32, b []byte) uint32 {
 	if len(b) >= 8 {
 		key64 := uint64(key)<<32 | uint64(key)
 

internal/wsframe/mask_test.go renamed to frame_test.go

@@ -1,24 +1,115 @@
-package wsframe_test
+// +build !js
+
+package websocket
 
 import (
-	"crypto/rand"
+	"bufio"
+	"bytes"
 	"encoding/binary"
-	"github.com/gobwas/ws"
-	"github.com/google/go-cmp/cmp"
 	"math/bits"
-	"nhooyr.io/websocket/internal/wsframe"
+	"nhooyr.io/websocket/internal/assert"
 	"strconv"
 	"testing"
+	"time"
 	_ "unsafe"
+
+	"github.com/gobwas/ws"
+	"github.com/google/go-cmp/cmp"
+
+	_ "github.com/gorilla/websocket"
+	"math/rand"
 )
 
+func init() {
+	rand.Seed(time.Now().UnixNano())
+}
+
+func TestHeader(t *testing.T) {
+	t.Parallel()
+
+	t.Run("readNegativeLength", func(t *testing.T) {
+		t.Parallel()
+
+		r := bufio.NewReader(bytes.NewReader([]byte{}))
+		_, err := ReadFrameHeader(r)
+		assert.Error(t, err)
+	})
+
+	t.Run("lengths", func(t *testing.T) {
+		t.Parallel()
+
+		lengths := []int{
+			124,
+			125,
+			126,
+			127,
+
+			65534,
+			65535,
+			65536,
+			65537,
+		}
+
+		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()
+
+		randBool := func() bool {
+			return rand.Intn(1) == 0
+		}
+
+		for i := 0; i < 10000; i++ {
+			h := header{
+				fin:    randBool(),
+				rsv1:   randBool(),
+				rsv2:   randBool(),
+				rsv3:   randBool(),
+				opcode: opcode(rand.Intn(16)),
+
+				masked:        randBool(),
+				maskKey:       rand.Uint32(),
+				payloadLength: rand.Int63(),
+			}
+
+			testHeader(t, h)
+		}
+	})
+}
+
+func testHeader(t *testing.T, h header) {
+	b := &bytes.Buffer{}
+	w := bufio.NewWriter(b)
+	r := bufio.NewReader(b)
+
+	err := WriteHeader(h, w)
+	assert.Success(t, err)
+	err = w.Flush()
+	assert.Success(t, err)
+
+	h2, err := ReadFrameHeader(r)
+	assert.Success(t, err)
+
+	assert.Equalf(t, h, h2, "written and read headers differ")
+}
+
 func Test_mask(t *testing.T) {
 	t.Parallel()
 
 	key := []byte{0xa, 0xb, 0xc, 0xff}
 	key32 := binary.LittleEndian.Uint32(key)
 	p := []byte{0xa, 0xb, 0xc, 0xf2, 0xc}
-	gotKey32 := wsframe.Mask(key32, p)
+	gotKey32 := mask(key32, p)
 
 	if exp := []byte{0, 0, 0, 0x0d, 0x6}; !cmp.Equal(exp, p) {
 		t.Fatalf("unexpected mask: %v", cmp.Diff(exp, p))
@@ -74,7 +165,7 @@ func Benchmark_mask(b *testing.B) {
 				b.ResetTimer()
 
 				for i := 0; i < b.N; i++ {
-					wsframe.Mask(key32, p)
+					mask(key32, p)
 				}
 			},
 		},