Apr 7, 2024 · Dec 18, 2023 · Apr 5, 2024 · Apr 5, 2024 · Apr 5, 2024 · Apr 5, 2024
diff --git a/accept_test.go b/accept_test.go
 "net/http"
 "net/http/httptest"
 "strings"
 "sync"
 "testing"

 "nhooyr.io/websocket/internal/test/assert"
 _, err := Accept(w, r, nil)
 assert.Contains(t, err, `failed to hijack connection`)
 })
 t.Run("closeRace", func(t *testing.T) {
 t.Parallel()

 server, _ := net.Pipe()

 rw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
 newResponseWriter := func() http.ResponseWriter {
 return mockHijacker{
 ResponseWriter: httptest.NewRecorder(),
 hijack: func() (net.Conn, *bufio.ReadWriter, error) {
 return server, rw, nil
 },
 }
 }
 w := newResponseWriter()

 r := httptest.NewRequest("GET", "/", nil)
 r.Header.Set("Connection", "Upgrade")
 r.Header.Set("Upgrade", "websocket")
 r.Header.Set("Sec-WebSocket-Version", "13")
 r.Header.Set("Sec-WebSocket-Key", xrand.Base64(16))

 c, err := Accept(w, r, nil)
 wg := &sync.WaitGroup{}
 wg.Add(2)
 go func() {
 c.Close(StatusInternalError, "the sky is falling")
 wg.Done()
 }()
 go func() {
 c.CloseNow()
 wg.Done()
 }()
 wg.Wait()
 assert.Success(t, err)
 })
 }

 func Test_verifyClientHandshake(t *testing.T) {
diff --git a/close.go b/close.go
 //
 // Close will unblock all goroutines interacting with the connection once
 // complete.
 func (c *Conn) Close(code StatusCode, reason string) error {
 defer c.wg.Wait()
 return c.closeHandshake(code, reason)
 func (c *Conn) Close(code StatusCode, reason string) (err error) {
 defer errd.Wrap(&err, "failed to close WebSocket")

 if !c.casClosing() {
 err = c.waitGoroutines()
 if err != nil {
 return err
 }
 return net.ErrClosed
 }
 defer func() {
 if errors.Is(err, net.ErrClosed) {
 err = nil
 }
 }()

 err = c.closeHandshake(code, reason)

 err2 := c.close()
 if err == nil && err2 != nil {
 err = err2
 }

 err2 = c.waitGoroutines()
 if err == nil && err2 != nil {
 err = err2
 }

 return err
 }

 // 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) {
 defer c.wg.Wait()
 defer errd.Wrap(&err, "failed to close WebSocket")

 if c.isClosed() {
 if !c.casClosing() {
 err = c.waitGoroutines()
 if err != nil {
 return err
 }
 return net.ErrClosed
 }
 defer func() {
 if errors.Is(err, net.ErrClosed) {
 err = nil
 }
 }()

 c.close(nil)
 return c.closeErr
 }

 func (c *Conn) closeHandshake(code StatusCode, reason string) (err error) {
 defer errd.Wrap(&err, "failed to close WebSocket")

 writeErr := c.writeClose(code, reason)
 closeHandshakeErr := c.waitCloseHandshake()
 err = c.close()

 if writeErr != nil {
 return writeErr
 err2 := c.waitGoroutines()
 if err == nil && err2 != nil {
 err = err2
 }
 return err
 }

 if CloseStatus(closeHandshakeErr) == -1 && !errors.Is(net.ErrClosed, closeHandshakeErr) {
 return closeHandshakeErr
 func (c *Conn) closeHandshake(code StatusCode, reason string) error {
 err := c.writeClose(code, reason)
 if err != nil {
 return err
 }

 err = c.waitCloseHandshake()
 if CloseStatus(err) != code {
 return err
 }
 return nil
 }

 func (c *Conn) writeClose(code StatusCode, reason string) error {
 c.closeMu.Lock()
 wroteClose := c.wroteClose
 c.wroteClose = true
 c.closeMu.Unlock()
 if wroteClose {
 return net.ErrClosed
 }

 ce := CloseError{
 Code:   code,
 Reason: reason,
 }

 var p []byte
 varmarshalErr error
 varerr error
 if ce.Code != StatusNoStatusRcvd {
 p, marshalErr = ce.bytes()
 }

 writeErr := c.writeControl(context.Background(), opClose, p)
 if CloseStatus(writeErr) != -1 {
 // Not a real error if it's due to a close frame being received.
 writeErr = nil
 p, err = ce.bytes()
 if err != nil {
 return err
 }
 }

 // We do this after in case there was an error writing the close frame.
 c.setCloseErr(fmt.Errorf("sent close frame: %w", ce))
 ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 defer cancel()

 if marshalErr != nil {
 return marshalErr
 err = c.writeControl(ctx, opClose, p)
 // If the connection closed as we're writing we ignore the error as we might
 // have written the close frame, the peer responded and then someone else read it
 // and closed the connection.
 if err != nil && !errors.Is(err, net.ErrClosed) {
 return err
 }
 returnwriteErr
 returnnil
 }

 func (c *Conn) waitCloseHandshake() error {
 defer c.close(nil)

 ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 defer cancel()

 }
 defer c.readMu.unlock()

 if c.readCloseFrameErr != nil {
 return c.readCloseFrameErr
 }

 for i := int64(0); i < c.msgReader.payloadLength; i++ {
 _, err := c.br.ReadByte()
 if err != nil {
 }
 }

 func (c *Conn) waitGoroutines() error {
 t := time.NewTimer(time.Second * 15)
 defer t.Stop()

 select {
 case <-c.timeoutLoopDone:
 case <-t.C:
 return errors.New("failed to wait for timeoutLoop goroutine to exit")
 }

 c.closeReadMu.Lock()
 closeRead := c.closeReadCtx != nil
 c.closeReadMu.Unlock()
 if closeRead {
 select {
 case <-c.closeReadDone:
 case <-t.C:
 return errors.New("failed to wait for close read goroutine to exit")
 }
 }

 select {
 case <-c.closed:
 case <-t.C:
 return errors.New("failed to wait for connection to be closed")
 }

 return nil
 }

 func parseClosePayload(p []byte) (CloseError, error) {
 if len(p) == 0 {
 return CloseError{
 return buf, nil
 }

 func (c *Conn)setCloseErr(err error) {
 func (c *Conn)casClosing() bool {
 c.closeMu.Lock()
 c.setCloseErrLocked(err)
 c.closeMu.Unlock()
 }

 func (c *Conn) setCloseErrLocked(err error) {
 if c.closeErr == nil && err != nil {
 c.closeErr = fmt.Errorf("WebSocket closed: %w", err)
 defer c.closeMu.Unlock()
 if !c.closing {
 c.closing = true
 return true
 }
 return false
 }

 func (c *Conn) isClosed() bool {
Original file line number	Diff line number	Diff line change
Expand Up		@@ -10,6 +10,7 @@ import (
		"net/http"
		"net/http/httptest"
		"strings"
		"sync"
		"testing"

		"nhooyr.io/websocket/internal/test/assert"
Expand DownExpand Up		@@ -142,6 +143,42 @@ func TestAccept(t *testing.T) {
		_, err := Accept(w, r, nil)
		assert.Contains(t, err, `failed to hijack connection`)
		})
		t.Run("closeRace", func(t *testing.T) {
		t.Parallel()

		server, _ := net.Pipe()

		rw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
		newResponseWriter := func() http.ResponseWriter {
		return mockHijacker{
		ResponseWriter: httptest.NewRecorder(),
		hijack: func() (net.Conn, *bufio.ReadWriter, error) {
		return server, rw, nil
		},
		}
		}
		w := newResponseWriter()

		r := httptest.NewRequest("GET", "/", nil)
		r.Header.Set("Connection", "Upgrade")
		r.Header.Set("Upgrade", "websocket")
		r.Header.Set("Sec-WebSocket-Version", "13")
		r.Header.Set("Sec-WebSocket-Key", xrand.Base64(16))

		c, err := Accept(w, r, nil)
		wg := &sync.WaitGroup{}
		wg.Add(2)
		go func() {
		c.Close(StatusInternalError, "the sky is falling")
		wg.Done()
		}()
		go func() {
		c.CloseNow()
		wg.Done()
		}()
		wg.Wait()
		assert.Success(t, err)
		})
		}

		func Test_verifyClientHandshake(t *testing.T) {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -97,80 +97,106 @@ func CloseStatus(err error) StatusCode {
		//
		// Close will unblock all goroutines interacting with the connection once
		// complete.
		func (c *Conn) Close(code StatusCode, reason string) error {
		defer c.wg.Wait()
		return c.closeHandshake(code, reason)
		func (c *Conn) Close(code StatusCode, reason string) (err error) {
		defer errd.Wrap(&err, "failed to close WebSocket")

		if !c.casClosing() {
		err = c.waitGoroutines()
		if err != nil {
		return err
		}
		return net.ErrClosed
		}
		defer func() {
		if errors.Is(err, net.ErrClosed) {
		err = nil
		}
		}()

		err = c.closeHandshake(code, reason)

		err2 := c.close()
		if err == nil && err2 != nil {
		err = err2
		}

		err2 = c.waitGoroutines()
		if err == nil && err2 != nil {
		err = err2
		}

		return err
		}

		// 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) {
		defer c.wg.Wait()
		defer errd.Wrap(&err, "failed to close WebSocket")

		if c.isClosed() {
		if !c.casClosing() {
		err = c.waitGoroutines()
		if err != nil {
		return err
		}
		return net.ErrClosed
		}
		defer func() {
		if errors.Is(err, net.ErrClosed) {
		err = nil
		}
		}()

		c.close(nil)
		return c.closeErr
		}

		func (c *Conn) closeHandshake(code StatusCode, reason string) (err error) {
		defer errd.Wrap(&err, "failed to close WebSocket")

		writeErr := c.writeClose(code, reason)
		closeHandshakeErr := c.waitCloseHandshake()
		err = c.close()

		if writeErr != nil {
		return writeErr
		err2 := c.waitGoroutines()
		if err == nil && err2 != nil {
		err = err2
		}
		return err
		}

		if CloseStatus(closeHandshakeErr) == -1 && !errors.Is(net.ErrClosed, closeHandshakeErr) {
		return closeHandshakeErr
		func (c *Conn) closeHandshake(code StatusCode, reason string) error {
		err := c.writeClose(code, reason)
		if err != nil {
		return err
		}

		err = c.waitCloseHandshake()
		if CloseStatus(err) != code {
		return err
		}
		return nil
		}

		func (c *Conn) writeClose(code StatusCode, reason string) error {
		c.closeMu.Lock()
		wroteClose := c.wroteClose
		c.wroteClose = true
		c.closeMu.Unlock()
		if wroteClose {
		return net.ErrClosed
		}

		ce := CloseError{
		Code: code,
		Reason: reason,
		}

		var p []byte
		varmarshalErr error
		varerr error
		if ce.Code != StatusNoStatusRcvd {
		p, marshalErr = ce.bytes()
		}

		writeErr := c.writeControl(context.Background(), opClose, p)
		if CloseStatus(writeErr) != -1 {
		// Not a real error if it's due to a close frame being received.
		writeErr = nil
		p, err = ce.bytes()
		if err != nil {
		return err
		}
		}

		// We do this after in case there was an error writing the close frame.
		c.setCloseErr(fmt.Errorf("sent close frame: %w", ce))
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

		if marshalErr != nil {
		return marshalErr
		err = c.writeControl(ctx, opClose, p)
		// If the connection closed as we're writing we ignore the error as we might
		// have written the close frame, the peer responded and then someone else read it
		// and closed the connection.
		if err != nil && !errors.Is(err, net.ErrClosed) {
		return err
		}
		returnwriteErr
		returnnil
		}

		func (c *Conn) waitCloseHandshake() error {
		defer c.close(nil)

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

Expand All		@@ -180,10 +206,6 @@ func (c *Conn) waitCloseHandshake() error {
		}
		defer c.readMu.unlock()

		if c.readCloseFrameErr != nil {
		return c.readCloseFrameErr
		}

		for i := int64(0); i < c.msgReader.payloadLength; i++ {
		_, err := c.br.ReadByte()
		if err != nil {
Expand All		@@ -206,6 +228,36 @@ func (c *Conn) waitCloseHandshake() error {
		}
		}

		func (c *Conn) waitGoroutines() error {
		t := time.NewTimer(time.Second * 15)
		defer t.Stop()

		select {
		case <-c.timeoutLoopDone:
		case <-t.C:
		return errors.New("failed to wait for timeoutLoop goroutine to exit")
		}

		c.closeReadMu.Lock()
		closeRead := c.closeReadCtx != nil
		c.closeReadMu.Unlock()
		if closeRead {
		select {
		case <-c.closeReadDone:
		case <-t.C:
		return errors.New("failed to wait for close read goroutine to exit")
		}
		}

		select {
		case <-c.closed:
		case <-t.C:
		return errors.New("failed to wait for connection to be closed")
		}

		return nil
		}

		func parseClosePayload(p []byte) (CloseError, error) {
		if len(p) == 0 {
		return CloseError{
Expand DownExpand Up		@@ -276,16 +328,14 @@ func (ce CloseError) bytesErr() ([]byte, error) {
		return buf, nil
		}

		func (c *Conn)setCloseErr(err error) {
		func (c *Conn)casClosing() bool {
		c.closeMu.Lock()
		c.setCloseErrLocked(err)
		c.closeMu.Unlock()
		}

		func (c *Conn) setCloseErrLocked(err error) {
		if c.closeErr == nil && err != nil {
		c.closeErr = fmt.Errorf("WebSocket closed: %w", err)
		defer c.closeMu.Unlock()
		if !c.closing {
		c.closing = true
		return true
		}
		return false
		}

		func (c *Conn) isClosed() bool {
Expand Down