test conn close/state

Author: mccutchen

websocket.go

@@ -380,8 +380,12 @@ func (ws *Websocket) doCloseHandshake(closeFrame *Frame, cause error) error {
 	// also ensure no one read or write can exceed our new close deadline,
 	// since we may do multiple reads below while waiting for our close
 	// frame to be ACK'd
-	ws.writeTimeout = ws.closeTimeout
-	ws.readTimeout = ws.closeTimeout
+	if ws.readTimeout > 0 {
+		ws.readTimeout = ws.closeTimeout
+	}
+	if ws.writeTimeout > 0 {
+		ws.writeTimeout = ws.closeTimeout
+	}
 
 	ws.hooks.OnCloseHandshakeStart(ws.clientKey, 0, cause) // TODO: close code
 	if err := ws.writeFrame(closeFrame); err != nil && !errors.Is(err, net.ErrClosed) {

websocket_test.go

@@ -747,6 +747,40 @@ func TestNew(t *testing.T) {
 	})
 }
 
+func TestClose(t *testing.T) {
+	// simulate a websocket client and server connaction
+	clientConn, serverConn := net.Pipe()
+	ws := websocket.New(serverConn, websocket.ClientKey("test-client-key"), websocket.ServerMode, websocket.Options{})
+
+	// close the server connection
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		assert.NilError(t, ws.Close())
+	}()
+
+	// confirm that the client got the close frame as expected
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		mustReadCloseFrame(t, clientConn, websocket.StatusNormalClosure, nil)
+		mustWriteFrame(t, clientConn, true, websocket.NewCloseFrame(websocket.StatusNormalClosure, ""))
+	}()
+	wg.Wait()
+
+	// make sure any reads or writes after closing the connection are rejected
+	{
+		msg, err := ws.ReadMessage(t.Context())
+		assert.Equal(t, msg, nil, "msg should be nil")
+		assert.Error(t, err, websocket.ErrConnectionClosed)
+	}
+	{
+		err := ws.WriteMessage(t.Context(), &websocket.Message{})
+		assert.Error(t, err, websocket.ErrConnectionClosed)
+	}
+}
+
 func mustReadFrame(t testing.TB, src io.Reader, maxPayloadLen int) *websocket.Frame {
 	t.Helper()
 	frame, err := websocket.ReadFrame(src, websocket.ClientMode, maxPayloadLen)
@@ -994,27 +1028,34 @@ var (
 // provide a separate reader and writer, and ensures that the conn can't be
 // used after it's closed.
 type dummyConn struct {
+	mu     sync.Mutex
 	in     io.Reader
 	out    io.Writer
-	closed atomic.Bool
+	closed bool
 }
 
 func (c *dummyConn) Read(p []byte) (int, error) {
-	if c.closed.Load() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed {
 		return 0, errors.New("reader closed")
 	}
 	return c.in.Read(p)
 }
 
 func (c *dummyConn) Write(p []byte) (int, error) {
-	if c.closed.Load() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed {
 		return 0, errors.New("writer closed")
 	}
 	return c.out.Write(p)
 }
 
 func (c *dummyConn) Close() error {
-	c.closed.Swap(true)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.closed = true
 	return nil
 }