diff --git a/conn.go b/conn.go index 861b2390a9de905278ca10a75f2bc22d54ffc2b4..cbb7fa5646a8a0a8caa10e1b251015ad6583e029 100644 --- a/conn.go +++ b/conn.go @@ -42,11 +42,12 @@ type Conn struct { closer io.Closer client bool - closeOnce sync.Once - closeErrOnce sync.Once - closeErr error - closed chan struct{} - closing *atomicInt64 + closeOnce sync.Once + closeErrOnce sync.Once + closeErr error + closed chan struct{} + closing *atomicInt64 + closeReceived error // messageWriter state. // writeMsgLock is acquired to write a data message. @@ -339,10 +340,12 @@ func (c *Conn) handleControl(ctx context.Context, h header) error { if err != nil { err = fmt.Errorf("received invalid close payload: %w", err) c.exportedClose(StatusProtocolError, err.Error(), false) + c.closeReceived = err return err } err = fmt.Errorf("received close: %w", ce) + c.closeReceived = err c.writeClose(b, err, false) if ctx.Err() != nil { @@ -941,6 +944,12 @@ func (c *Conn) waitClose() error { return err } defer c.releaseLock(c.readLock) + + if c.closeReceived != nil { + // goroutine reading just received the close. + return c.closeReceived + } + c.readerShouldLock = false b := bpool.Get() diff --git a/conn_test.go b/conn_test.go index d924fd0aaf3c4ee68c5b3a9dc2eca744fd270a75..83f09dbf9b839ef24273592a5a1bbcd6d6e90d4e 100644 --- a/conn_test.go +++ b/conn_test.go @@ -868,6 +868,29 @@ func TestConn(t *testing.T) { return c.Close(websocket.StatusNormalClosure, "") }, }, + { + // Issue #164 + name: "closeHandshake_concurrentRead", + server: func(ctx context.Context, c *websocket.Conn) error { + _, _, err := c.Read(ctx) + return assertCloseStatus(err, websocket.StatusNormalClosure) + }, + client: func(ctx context.Context, c *websocket.Conn) error { + errc := make(chan error, 1) + go func() { + _, _, err := c.Read(ctx) + errc <- err + }() + + err := c.Close(websocket.StatusNormalClosure, "") + if err != nil { + return err + } + + err = <-errc + return assertCloseStatus(err, websocket.StatusNormalClosure) + }, + }, } for _, tc := range testCases { tc := tc