diff --git a/close.go b/close.go index fcc68065..da500624 100644 --- a/close.go +++ b/close.go @@ -96,7 +96,18 @@ func CloseStatus(err error) StatusCode { // // Close will unblock all goroutines interacting with the connection once // complete. -func (c *Conn) Close(code StatusCode, reason string) (err error) { +func (c *Conn) Close(code StatusCode, reason string) error { + firstCall := !c.closeCalled.Swap(true) + err := c.closeWithHandshake(code, reason) + if firstCall && errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + +// closeWithHandshake implements Close for library code, which must not set +// closeCalled. +func (c *Conn) closeWithHandshake(code StatusCode, reason string) (err error) { defer errd.Wrap(&err, "failed to close WebSocket") if c.casClosing() { @@ -129,7 +140,16 @@ func (c *Conn) Close(code StatusCode, reason string) (err error) { // CloseNow closes the WebSocket connection without attempting a close handshake. // Use when you do not want the overhead of the close handshake. -func (c *Conn) CloseNow() (err error) { +func (c *Conn) CloseNow() error { + firstCall := !c.closeCalled.Swap(true) + err := c.closeWithoutHandshake() + if firstCall && errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + +func (c *Conn) closeWithoutHandshake() (err error) { defer errd.Wrap(&err, "failed to immediately close WebSocket") if c.casClosing() { diff --git a/conn.go b/conn.go index 0569bd2c..b79ccce2 100644 --- a/conn.go +++ b/conn.go @@ -77,9 +77,13 @@ type Conn struct { closeReadCtx context.Context closeReadDone chan struct{} - closing atomic.Bool - closeMu sync.Mutex // Protects following. - closed chan struct{} + // closeCalled is set by the first Close or CloseNow call. Library code + // closes the connection without setting it, so the user's first Close + // does not report net.ErrClosed when the library closed first. + closeCalled atomic.Bool + closing atomic.Bool + closeMu sync.Mutex // Protects following. + closed chan struct{} pingCounter atomic.Int64 activePingsMu sync.Mutex diff --git a/conn_test.go b/conn_test.go index 01af8cd4..6568dc35 100644 --- a/conn_test.go +++ b/conn_test.go @@ -401,6 +401,71 @@ func TestConn(t *testing.T) { assert.ErrorIs(t, websocket.ErrClosed, err2) }) + t.Run("firstCloseAfterLibraryClose", func(t *testing.T) { + t.Parallel() + + closers := []struct { + name string + close func(*websocket.Conn) error + }{ + {"Close", func(c *websocket.Conn) error { return c.Close(websocket.StatusNormalClosure, "") }}, + {"CloseNow", (*websocket.Conn).CloseNow}, + } + for _, closer := range closers { + t.Run(closer.name, func(t *testing.T) { + t.Parallel() + + t.Run("peerClose", func(t *testing.T) { + tt, c1, c2 := newConnTest(t, nil, nil) + + readErr := xsync.Go(func() error { + _, _, err := c1.Read(tt.ctx) + return err + }) + assert.Success(t, c2.Close(websocket.StatusNormalClosure, "")) + assert.Success(t, assertCloseStatus(websocket.StatusNormalClosure, <-readErr)) + + assert.Success(t, closer.close(c1)) + }) + + t.Run("closeReadDataMessage", func(t *testing.T) { + tt, c1, c2 := newConnTest(t, nil, nil) + + closeReadCtx := c1.CloseRead(tt.ctx) + readErr := xsync.Go(func() error { + _, _, err := c2.Read(tt.ctx) + return err + }) + assert.Success(t, c2.Write(tt.ctx, websocket.MessageText, []byte("x"))) + assert.Success(t, assertCloseStatus(websocket.StatusPolicyViolation, <-readErr)) + <-closeReadCtx.Done() + + assert.Success(t, closer.close(c1)) + }) + + t.Run("netConnUnexpectedType", func(t *testing.T) { + tt, c1, c2 := newConnTest(t, nil, nil) + + nc := websocket.NetConn(tt.ctx, c1, websocket.MessageBinary) + readErr := xsync.Go(func() error { + _, _, err := c2.Read(tt.ctx) + return err + }) + // net.Pipe writes block until nc.Read reads them. + writeErr := xsync.Go(func() error { + return c2.Write(tt.ctx, websocket.MessageText, []byte("x")) + }) + _, err := nc.Read(make([]byte, 1)) + assert.Contains(t, err, "unexpected frame type") + assert.Success(t, <-writeErr) + assert.Success(t, assertCloseStatus(websocket.StatusUnsupportedData, <-readErr)) + + assert.Success(t, closer.close(c1)) + }) + }) + } + }) + t.Run("MidReadClose", func(t *testing.T) { tt, c1, c2 := newConnTest(t, nil, nil) @@ -449,6 +514,21 @@ func TestWasm(t *testing.T) { } s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Sec-WebSocket-Protocol") == "close" { + // Closes first so the client's Close runs after the browser + // closed the connection. + c, err := websocket.Accept(w, r, &websocket.AcceptOptions{ + Subprotocols: []string{"close"}, + InsecureSkipVerify: true, + }) + if err == nil { + err = c.Close(websocket.StatusNormalClosure, "") + } + if err != nil { + t.Error(err) + } + return + } err := echoServer(w, r, &websocket.AcceptOptions{ Subprotocols: []string{"echo"}, InsecureSkipVerify: true, diff --git a/netconn.go b/netconn.go index a854faa4..207393b6 100644 --- a/netconn.go +++ b/netconn.go @@ -174,7 +174,7 @@ func (nc *netConn) read(p []byte) (int, error) { } if typ != nc.msgType { err := fmt.Errorf("unexpected frame type read (expected %v): %v", nc.msgType, typ) - nc.c.Close(StatusUnsupportedData, err.Error()) + nc.c.closeWithHandshake(StatusUnsupportedData, err.Error()) return 0, err } nc.reader = r diff --git a/read.go b/read.go index 06672e10..e68ec9c8 100644 --- a/read.go +++ b/read.go @@ -79,7 +79,7 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context { defer c.close() _, _, err := c.Reader(ctx) if err == nil { - c.Close(StatusPolicyViolation, "unexpected data message") + c.closeWithHandshake(StatusPolicyViolation, "unexpected data message") } }() return ctx diff --git a/ws_js.go b/ws_js.go index 026b75fc..a37879ec 100644 --- a/ws_js.go +++ b/ws_js.go @@ -56,6 +56,10 @@ type Conn struct { closeErrOnce sync.Once closeErr error closeWasClean bool + // closeCalled is set by the first Close or CloseNow call. Library code + // and the browser close the connection without setting it, so the user's + // first Close does not report net.ErrClosed after either closed it. + closeCalled atomic.Bool releaseOnClose func() releaseOnError func() @@ -125,7 +129,7 @@ func (c *Conn) init() { } func (c *Conn) closeWithInternal() { - c.Close(StatusInternalError, "something went wrong") + c.closeWithHandshake(StatusInternalError, "something went wrong") } // Read attempts to read a message from the connection. @@ -145,7 +149,7 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) { readLimit := c.msgReadLimit.Load() if readLimit >= 0 && int64(len(p)) > readLimit { reason := fmt.Errorf("read limited at %d bytes", c.msgReadLimit.Load()) - c.Close(StatusMessageTooBig, reason.Error()) + c.closeWithHandshake(StatusMessageTooBig, reason.Error()) return 0, nil, fmt.Errorf("%w: %v", ErrMessageTooBig, reason) } return typ, p, nil @@ -154,7 +158,7 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) { func (c *Conn) read(ctx context.Context) (MessageType, []byte, error) { select { case <-ctx.Done(): - c.Close(StatusPolicyViolation, "read timed out") + c.closeWithHandshake(StatusPolicyViolation, "read timed out") return 0, nil, ctx.Err() case <-c.readSignal: case <-c.closed: @@ -229,6 +233,17 @@ func (c *Conn) write(typ MessageType, p []byte) error { // or the connection is closed. // It thus performs the full WebSocket close handshake. func (c *Conn) Close(code StatusCode, reason string) error { + firstCall := !c.closeCalled.Swap(true) + err := c.closeWithHandshake(code, reason) + if firstCall && errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + +// closeWithHandshake implements Close for library code, which must not set +// closeCalled. +func (c *Conn) closeWithHandshake(code StatusCode, reason string) error { err := c.exportedClose(code, reason) if err != nil { return fmt.Errorf("failed to close WebSocket: %w", err) @@ -403,10 +418,10 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context { go func() { defer cancel() - defer c.CloseNow() + defer c.closeWithHandshake(StatusGoingAway, "") _, _, err := c.read(ctx) if err != nil { - c.Close(StatusPolicyViolation, "unexpected data message") + c.closeWithHandshake(StatusPolicyViolation, "unexpected data message") } }() return ctx diff --git a/ws_js_test.go b/ws_js_test.go index 1fa242f4..7d8ff9c2 100644 --- a/ws_js_test.go +++ b/ws_js_test.go @@ -2,6 +2,7 @@ package websocket_test import ( "context" + "net" "net/http" "os" "testing" @@ -37,6 +38,25 @@ func TestWasm(t *testing.T) { assert.Success(t, err) } +func TestWasmFirstCloseAfterPeerClose(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + c, _, err := websocket.Dial(ctx, os.Getenv("WS_ECHO_SERVER_URL"), &websocket.DialOptions{ + Subprotocols: []string{"close"}, + }) + assert.Success(t, err) + + _, _, err = c.Read(ctx) + // Read returns once the browser has closed the connection. + assert.ErrorIs(t, net.ErrClosed, err) + + err = c.Close(websocket.StatusNormalClosure, "") + assert.Success(t, err) +} + func TestWasmDialTimeout(t *testing.T) { t.Parallel()