Rename DataType to MessageType

- Make messageReader a value type to avoid allocation
- Add a bunch of important TODOs

Author: nhooyr

accept.go

@@ -41,6 +41,7 @@ func (o acceptOrigins) acceptOption() {}
 // Use this option with caution to avoid exposing your WebSocket
 // server to a CSRF attack.
 // See https://stackoverflow.com/a/37837709/4283659
+// TODO remove in favour of AcceptInsecureOrigin
 func AcceptOrigins(origins ...string) AcceptOption {
 	return acceptOrigins(origins)
 }

datatype.go

@@ -14,13 +14,18 @@ import (
 
 func ExampleAccept_echo() {
 	fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		c, err := websocket.Accept(w, r)
+		c, err := websocket.Accept(w, r, websocket.AcceptSubprotocols("echo"))
 		if err != nil {
 			log.Printf("server handshake failed: %v", err)
 			return
 		}
 		defer c.Close(websocket.StatusInternalError, "")
 
+		if c.Subprotocol() == "" {
+			c.Close(websocket.StatusPolicyViolation, "cannot communicate with the default protocol")
+			return
+		}
+
 		echo := func() error {
 			ctx, cancel := context.WithTimeout(r.Context(), time.Minute)
 			defer cancel()

datatype_string.go

@@ -28,7 +28,7 @@ func (jc *JSONConn) read(ctx context.Context, v interface{}) error {
 		return err
 	}
 
-	if typ != DataText {
+	if typ != MessageText {
 		return xerrors.Errorf("unexpected frame type for json (expected DataText): %v", typ)
 	}
 
@@ -39,6 +39,7 @@ func (jc *JSONConn) read(ctx context.Context, v interface{}) error {
 	if err != nil {
 		return xerrors.Errorf("failed to decode json: %w", err)
 	}
+
 	return nil
 }
 
@@ -52,7 +53,7 @@ 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, DataText)
+	w := jc.Conn.Write(ctx, MessageText)
 
 	e := json.NewEncoder(w)
 	err := e.Encode(v)

example_test.go

@@ -0,0 +1,12 @@
+package websocket
+
+// MessageType represents the Opcode of a WebSocket data frame.
+type MessageType int
+
+//go:generate go run golang.org/x/tools/cmd/stringer -type=MessageType
+
+// MessageType constants.
+const (
+	MessageText   MessageType = MessageType(opText)
+	MessageBinary MessageType = MessageType(opBinary)
+)

json.go

@@ -23,8 +23,7 @@ type control struct {
 type Conn struct {
 	subprotocol string
 	br          *bufio.Reader
-	// TODO switch to []byte for write buffering because for messages larger than buffers, there will always be 3 writes. One for the frame, one for the message, one for the fin.
-	// Also will help for compression.
+	// TODO switch to []byte for write buffering for predicting compression in memory maybe
 	bw     *bufio.Writer
 	closer io.Closer
 	client bool
@@ -36,7 +35,7 @@ type Conn struct {
 	// Writers should send on write to begin sending
 	// a message and then follow that up with some data
 	// on writeBytes.
-	write      chan DataType
+	write      chan MessageType
 	control    chan control
 	writeBytes chan []byte
 	writeDone  chan struct{}
@@ -45,9 +44,10 @@ type Conn struct {
 	// Then send a byte slice to readBytes to read into it.
 	// The n of bytes read will be sent on readDone once the read into a slice is complete.
 	// readDone will receive 0 when EOF is reached.
-	read      chan opcode
-	readBytes chan []byte
-	readDone  chan int
+	read       chan opcode
+	readBytes  chan []byte
+	readDone   chan int
+	readerDone chan struct{}
 }
 
 func (c *Conn) close(err error) {
@@ -77,12 +77,13 @@ func (c *Conn) Subprotocol() string {
 
 func (c *Conn) init() {
 	c.closed = make(chan struct{})
-	c.write = make(chan DataType)
+	c.write = make(chan MessageType)
 	c.control = make(chan control)
 	c.writeDone = make(chan struct{})
 	c.read = make(chan opcode)
 	c.readDone = make(chan int)
 	c.readBytes = make(chan []byte)
+	c.readerDone = make(chan struct{})
 
 	runtime.SetFinalizer(c, func(c *Conn) {
 		c.Close(StatusInternalError, "websocket: connection ended up being garbage collected")
@@ -120,7 +121,7 @@ messageLoop:
 	for {
 		c.writeBytes = make(chan []byte)
 
-		var dataType DataType
+		var dataType MessageType
 		select {
 		case <-c.closed:
 			return
@@ -321,7 +322,7 @@ func (c *Conn) readLoop() {
 			select {
 			case <-c.closed:
 				return
-			case c.readDone <- 0:
+			case c.readerDone <- struct{}{}:
 			}
 		}
 	}
@@ -407,7 +408,8 @@ func (c *Conn) writeControl(ctx context.Context, opcode opcode, p []byte) error
 // a WebSocket data frame of type dataType to the connection.
 // Ensure you close the messageWriter once you have written to entire message.
 // Concurrent calls to messageWriter are ok.
-func (c *Conn) Write(ctx context.Context, dataType DataType) io.WriteCloser {
+func (c *Conn) Write(ctx context.Context, dataType MessageType) io.WriteCloser {
+	// TODO acquire write here, move state into Conn and make messageWriter allocation free.
 	return &messageWriter{
 		c:        c,
 		ctx:      ctx,
@@ -418,7 +420,7 @@ func (c *Conn) Write(ctx context.Context, dataType DataType) io.WriteCloser {
 // messageWriter enables writing to a WebSocket connection.
 // Ensure you close the messageWriter once you have written to entire message.
 type messageWriter struct {
-	datatype     DataType
+	datatype     MessageType
 	ctx          context.Context
 	c            *Conn
 	acquiredLock bool
@@ -489,12 +491,16 @@ func (w *messageWriter) Close() error {
 // Please use SetContext on the reader to bound the read operation.
 // Your application must keep reading messages for the Conn to automatically respond to ping
 // and close frames.
-func (c *Conn) Read(ctx context.Context) (DataType, io.Reader, error) {
+func (c *Conn) Read(ctx context.Context) (MessageType, io.Reader, error) {
+	// TODO error if the reader is not done
 	select {
+	case <-c.readerDone:
+		// The previous reader just hit a io.EOF, we handle it for users
+		return c.Read(ctx)
 	case <-c.closed:
 		return 0, nil, xerrors.Errorf("failed to read message: %w", c.closeErr)
 	case opcode := <-c.read:
-		return DataType(opcode), &messageReader{
+		return MessageType(opcode), messageReader{
 			ctx: ctx,
 			c:   c,
 		}, nil
@@ -510,13 +516,26 @@ type messageReader struct {
 }
 
 // Read reads as many bytes as possible into p.
-func (r *messageReader) Read(p []byte) (n int, err error) {
+func (r messageReader) Read(p []byte) (int, error) {
+	n, err := r.read(p)
+	if err != nil {
+		// Have to return io.EOF directly for now.
+		if err == io.EOF {
+			return 0, io.EOF
+		}
+		return n, xerrors.Errorf("failed to read: %w", err)
+	}
+	return n, nil
+}
+
+func (r messageReader) read(p []byte) (int, error) {
 	select {
 	case <-r.c.closed:
 		return 0, r.c.closeErr
-	case <-r.c.readDone:
+	case <-r.c.readerDone:
 		return 0, io.EOF
 	case r.c.readBytes <- p:
+		// TODO this is potentially racey as if we return if the context is cancelled, or the conn is closed we don't know if the p is ok to use. we must close the connection and also ensure the readLoop is done before returning, likewise with writes.
 		select {
 		case <-r.c.closed:
 			return 0, r.c.closeErr