diff --git a/example_test.go b/example_test.go index 050af907e8c27eeacae40c32b9151bf6a4537495..0b59e6a0b724dc00b7dce65241083aae8d5f6fe3 100644 --- a/example_test.go +++ b/example_test.go @@ -74,11 +74,7 @@ func Example_writeOnly() { ctx, cancel := context.WithTimeout(r.Context(), time.Minute*10) defer cancel() - go func() { - defer cancel() - c.Reader(ctx) - c.Close(websocket.StatusPolicyViolation, "server doesn't accept data messages") - }() + ctx = c.CloseRead(ctx) t := time.NewTicker(time.Second * 30) defer t.Stop() diff --git a/websocket.go b/websocket.go index a2ff415d35e17c8ad18e28e390b0480505027130..91a6808f981155ed29b38842a53baec1591f239b 100644 --- a/websocket.go +++ b/websocket.go @@ -22,7 +22,7 @@ import ( // and SetReadLimit. // // You must always read from the connection. Otherwise control -// frames will not be handled. See the docs on Reader. +// frames will not be handled. See the docs on Reader and CloseRead. // // Please be sure to call Close on the connection when you // are finished with it to release the associated resources. @@ -319,10 +319,8 @@ func (c *Conn) handleControl(ctx context.Context, h header) error { // to be closed so you do not need to write your own error message. // This applies to the Read methods in the wsjson/wspb subpackages as well. // -// You must read from the connection for close frames to be read. -// If you do not expect any data messages from the peer, just call -// Reader in a separate goroutine and close the connection with StatusPolicyViolation -// when it returns. See the writeOnly example. +// You must read from the connection for control frames to be handled. +// If you do not expect any data messages from the peer, call CloseRead. // // Only one Reader may be open at a time. // @@ -388,6 +386,11 @@ func (c *Conn) reader(ctx context.Context) (MessageType, io.Reader, error) { return MessageType(h.opcode), r, nil } +// CloseRead will close the connection if any data message is received from the peer. +// Call this when you are done reading data messages from the connection but will still write +// to it. Since CloseRead is still reading from the connection, it will respond to ping, pong +// and close frames automatically. It will only close the connection on a data frame. The returned +// context will be cancelled when the connection is closed. func (c *Conn) CloseRead(ctx context.Context) context.Context { ctx, cancel := context.WithCancel(ctx) go func() {