fix(net): coordinate websocket handoff with listener shutdown (#5563)

pkg/util/net/websocket.go

@@ -7,7 +7,6 @@ import (
 	"sync"
 	"time"
 
-	gerr "github.com/fatedier/golib/errors"
 	"golang.org/x/net/websocket"
 )
 
@@ -20,11 +19,11 @@ const (
 type WebsocketListener struct {
 	ln       net.Listener
 	acceptCh chan net.Conn
+	done     chan struct{}
 
 	server *http.Server
 
-	mu     sync.Mutex
+	closeOnce sync.Once
-	closed bool
 }
 
 // NewWebsocketListener to handle websocket connections
@@ -33,6 +32,7 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
 	wl = &WebsocketListener{
 		ln:       ln,
 		acceptCh: make(chan net.Conn),
+		done:     make(chan struct{}),
 	}
 
 	muxer := http.NewServeMux()
@@ -46,14 +46,13 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
 		conn := WrapCloseNotifyConn(c, func(_ error) {
 			close(notifyCh)
 		})
-		// The listener may be closed while this connection is being handed
+		select {
-		// over, so the send has to tolerate a closed channel. A nil error
+		case <-wl.done:
-		// means the connection was accepted and is owned by the caller.
-		if err := gerr.PanicToError(func() {
-			wl.acceptCh <- conn
-		}); err != nil {
 			conn.Close()
 			return
+		case wl.acceptCh <- conn:
+			// A ready handoff may win when Close races with it. Once delivered,
+			// the connection is owned by the caller of Accept.
 		}
 		<-notifyCh
 	}))
@@ -71,23 +70,21 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
 }
 
 func (p *WebsocketListener) Accept() (net.Conn, error) {
-	c, ok := <-p.acceptCh
+	select {
-	if !ok {
+	case <-p.done:
 		return nil, ErrWebsocketListenerClosed
+	default:
 	}
+	select {
+	case <-p.done:
+		return nil, ErrWebsocketListenerClosed
+	case c := <-p.acceptCh:
 		return c, nil
+	}
 }
 
 func (p *WebsocketListener) Close() error {
-	p.mu.Lock()
+	p.closeOnce.Do(func() { close(p.done) })
-	if !p.closed {
-		p.closed = true
-		// Closing acceptCh releases a pending Accept, matching the other
-		// listeners in this package. Otherwise it blocks forever and the
-		// caller never learns that the listener is closed.
-		close(p.acceptCh)
-	}
-	p.mu.Unlock()
 	return p.server.Close()
 }
 

pkg/util/net/websocket_test.go

@@ -15,16 +15,25 @@
 package net
 
 import (
+	"bufio"
 	"errors"
+	"io"
 	"net"
+	"net/http"
+	"net/http/httptest"
+	"sync"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"golang.org/x/net/websocket"
 )
 
 type closeStubListener struct {
 	closed chan struct{}
+	once   sync.Once
+	err    error
 }
 
 func (l *closeStubListener) Accept() (net.Conn, error) {
@@ -33,12 +42,8 @@ func (l *closeStubListener) Accept() (net.Conn, error) {
 }
 
 func (l *closeStubListener) Close() error {
-	select {
+	l.once.Do(func() { close(l.closed) })
-	case <-l.closed:
+	return l.err
-	default:
-		close(l.closed)
-	}
-	return nil
 }
 
 func (l *closeStubListener) Addr() net.Addr {
@@ -46,28 +51,147 @@ func (l *closeStubListener) Addr() net.Addr {
 }
 
 func TestWebsocketListenerAcceptUnblocksOnClose(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
 		wl := NewWebsocketListener(&closeStubListener{closed: make(chan struct{})})
-
+		t.Cleanup(func() { wl.Close() })
-	errCh := make(chan error, 1)
+		errCh := make(chan error, 3)
+		for range 3 {
 			go func() {
-		_, err := wl.Accept()
+				c, err := wl.Accept()
+				if c != nil {
+					t.Error("Accept returned a connection without a sender")
+				}
 				errCh <- err
 			}()
-
-	select {
-	case err := <-errCh:
-		t.Fatalf("Accept returned before the listener was closed: %v", err)
-	case <-time.After(100 * time.Millisecond):
 		}
+		synctest.Wait() // All Accept calls are blocked before Close.
+		require.Empty(t, errCh)
 
-	require.NoError(t, wl.Close())
+		var wg sync.WaitGroup
-	// Closing twice must stay safe and keep reporting the closed error.
+		for range 3 {
-	require.NoError(t, wl.Close())
+			wg.Go(func() {
-
+				if err := wl.Close(); err != nil {
-	select {
+					t.Errorf("Close: %v", err)
-	case err := <-errCh:
-		require.True(t, errors.Is(err, ErrWebsocketListenerClosed), "got %v", err)
-	case <-time.After(time.Second):
-		t.Fatal("Accept did not return after the listener was closed")
 				}
+			})
+		}
+		wg.Wait()
+		synctest.Wait()
+		require.Len(t, errCh, 3)
+		for range 3 {
+			require.ErrorIs(t, <-errCh, ErrWebsocketListenerClosed)
+		}
+		c, err := wl.Accept()
+		require.Nil(t, c)
+		require.ErrorIs(t, err, ErrWebsocketListenerClosed)
+	})
+}
+
+func TestWebsocketListenerCloseError(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		closeErr := errors.New("listener close failed")
+		wl := NewWebsocketListener(&closeStubListener{closed: make(chan struct{}), err: closeErr})
+		synctest.Wait() // Let Serve register its listener before Close.
+		require.ErrorIs(t, wl.Close(), closeErr)
+		require.NoError(t, wl.Close())
+	})
+}
+
+type websocketPipeHijacker struct {
+	*httptest.ResponseRecorder
+	conn net.Conn
+	rw   *bufio.ReadWriter
+}
+
+func (h *websocketPipeHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return h.conn, h.rw, nil
+}
+
+func newWebsocketPipe(t *testing.T) (*WebsocketListener, *websocket.Conn, net.Conn, <-chan struct{}) {
+	t.Helper()
+	wl := NewWebsocketListener(&closeStubListener{closed: make(chan struct{})})
+	serverConn, clientConn := net.Pipe()
+	t.Cleanup(func() {
+		wl.Close()
+		serverConn.Close()
+		clientConn.Close()
+	})
+	require.NoError(t, serverConn.SetDeadline(time.Now().Add(time.Second)))
+	require.NoError(t, clientConn.SetDeadline(time.Now().Add(time.Second)))
+	handlerDone := make(chan struct{})
+	go func() {
+		defer close(handlerDone)
+		rw := bufio.NewReadWriter(bufio.NewReader(serverConn), bufio.NewWriter(serverConn))
+		req, err := http.ReadRequest(rw.Reader)
+		if err != nil {
+			t.Errorf("ReadRequest: %v", err)
+			return
+		}
+		wl.server.Handler.ServeHTTP(&websocketPipeHijacker{httptest.NewRecorder(), serverConn, rw}, req)
+	}()
+	cfg, err := websocket.NewConfig("ws://localhost"+FrpWebsocketPath, "http://localhost")
+	require.NoError(t, err)
+	peer, err := websocket.NewClient(cfg, clientConn)
+	require.NoError(t, err)
+	return wl, peer, serverConn, handlerDone
+}
+
+func TestWebsocketListenerUnacceptedConnectionClose(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		wl, peer, raw, handlerDone := newWebsocketPipe(t)
+		peerErr := make(chan error, 1)
+		go func() {
+			_, err := peer.Read(make([]byte, 1))
+			peerErr <- err
+		}()
+		// No Accept consumer exists. Do not Wait here: it would synchronize
+		// the handler's send with Close and hide the baseline send/close race.
+		require.NoError(t, wl.Close())
+		require.ErrorIs(t, <-peerErr, io.EOF)
+		<-handlerDone
+		_, err := raw.Read(make([]byte, 1))
+		require.ErrorIs(t, err, io.ErrClosedPipe)
+	})
+}
+
+func TestWebsocketListenerAcceptedConnectionOwnership(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		wl, peer, raw, handlerDone := newWebsocketPipe(t)
+		owner, err := wl.Accept()
+		require.NoError(t, err)
+		t.Cleanup(func() {
+			raw.Close()
+			owner.Close()
+		})
+		require.NoError(t, wl.Close())
+
+		payload := []byte{0, 255, 128, 42}
+		reply := []byte("reply")
+		gotReply := make([]byte, len(reply))
+		peerErr := make(chan error, 2)
+		go func() {
+			if _, err := peer.Write(payload); err != nil {
+				peerErr <- err
+				return
+			}
+			_, err := io.ReadFull(peer, gotReply)
+			peerErr <- err
+			_, err = peer.Read(make([]byte, 1))
+			peerErr <- err
+		}()
+		got := make([]byte, len(payload))
+		_, err = io.ReadFull(owner, got)
+		require.NoError(t, err)
+		require.Equal(t, payload, got)
+		_, err = owner.Write(reply)
+		require.NoError(t, err)
+		require.NoError(t, <-peerErr)
+		require.Equal(t, reply, gotReply)
+		require.NoError(t, owner.Close())
+		require.NoError(t, owner.Close())
+		require.ErrorIs(t, <-peerErr, io.EOF)
+		<-handlerDone
+		_, err = raw.Read(make([]byte, 1))
+		require.ErrorIs(t, err, io.ErrClosedPipe)
+	})
 }