Merge pull request #114 from nhooyr/netconn

Add msgType parameter to NetConn adapter

Author: nhooyr

netconn.go

@@ -2,6 +2,7 @@ package websocket
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"math"
 	"net"
@@ -17,8 +18,11 @@ import (
 // correctly and so provided in the library.
 // See https://github.com/nhooyr/websocket/issues/100.
 //
-// Every Write to the net.Conn will correspond to a binary message
-// write on *webscoket.Conn.
+// Every Write to the net.Conn will correspond to a message write of
+// the given type on *websocket.Conn.
+//
+// If a message is read that is not of the correct type, an error
+// will be thrown.
 //
 // Close will close the *websocket.Conn with StatusNormalClosure.
 //
@@ -30,9 +34,10 @@ import (
 // and "websocket/unknown-addr" for String.
 //
 // A received StatusNormalClosure close frame will be translated to EOF when reading.
-func NetConn(c *Conn) net.Conn {
+func NetConn(c *Conn, msgType MessageType) net.Conn {
 	nc := &netConn{
-		c: c,
+		c:       c,
+		msgType: msgType,
 	}
 
 	var cancel context.CancelFunc
@@ -52,7 +57,8 @@ func NetConn(c *Conn) net.Conn {
 }
 
 type netConn struct {
-	c *Conn
+	c       *Conn
+	msgType MessageType
 
 	writeTimer   *time.Timer
 	writeContext context.Context
@@ -71,7 +77,7 @@ func (c *netConn) Close() error {
 }
 
 func (c *netConn) Write(p []byte) (int, error) {
-	err := c.c.Write(c.writeContext, MessageBinary, p)
+	err := c.c.Write(c.writeContext, c.msgType, p)
 	if err != nil {
 		return 0, err
 	}
@@ -93,9 +99,9 @@ func (c *netConn) Read(p []byte) (int, error) {
 			}
 			return 0, err
 		}
-		if typ != MessageBinary {
-			c.c.Close(StatusUnsupportedData, "can only accept binary messages")
-			return 0, xerrors.Errorf("unexpected frame type read for net conn adapter (expected %v): %v", MessageBinary, typ)
+		if typ != c.msgType {
+			c.c.Close(StatusUnsupportedData, fmt.Sprintf("can only accept %v messages", c.msgType))
+			return 0, xerrors.Errorf("unexpected frame type read for net conn adapter (expected %v): %v", c.msgType, typ)
 		}
 		c.reader = r
 	}

websocket.go

@@ -333,6 +333,10 @@ func (c *Conn) handleControl(ctx context.Context, h header) error {
 // This applies to the Read methods in the wsjson/wspb subpackages as well.
 //
 // You must read from the connection for control frames to be handled.
+// Thus if you expect messages to take a long time to be responded to,
+// you should handle such messages async to reading from the connection
+// to ensure control frames are promptly handled.
+//
 // If you do not expect any data messages from the peer, call CloseRead.
 //
 // Only one Reader may be open at a time.

websocket_test.go

@@ -127,7 +127,7 @@ func TestHandshake(t *testing.T) {
 				}
 				defer c.Close(websocket.StatusInternalError, "")
 
-				nc := websocket.NetConn(c)
+				nc := websocket.NetConn(c, websocket.MessageBinary)
 				defer nc.Close()
 
 				nc.SetWriteDeadline(time.Time{})
@@ -152,7 +152,7 @@ func TestHandshake(t *testing.T) {
 				}
 				defer c.Close(websocket.StatusInternalError, "")
 
-				nc := websocket.NetConn(c)
+				nc := websocket.NetConn(c, websocket.MessageBinary)
 				defer nc.Close()
 
 				nc.SetReadDeadline(time.Time{})