diff --git a/netconn.go b/netconn.go index 06cbc2f55e5fa2c4e773130aafdde385b0a29b30..25787703d97ec669930d898bd6b50f04a682468d 100644 --- a/netconn.go +++ b/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 } diff --git a/websocket.go b/websocket.go index 80d5511a10fbbe0a3cf0a98324938ba413fa1cf6..35ddb024d0a680c02c32c430856fa6fceb352669 100644 --- a/websocket.go +++ b/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. diff --git a/websocket_test.go b/websocket_test.go index 46f9c83396085d997fa3186a4f4f6a0f65848b4d..06e0fc6d732e5901029a378857f756d18d41eb3f 100644 --- a/websocket_test.go +++ b/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{})