From 2542d11a6201c0f2bbe3b1d93031fc1e83201056 Mon Sep 17 00:00:00 2001 From: fatedier Date: Sat, 10 Oct 2026 19:58:05 +0800 Subject: [PATCH] fix(net): coordinate websocket handoff with listener shutdown (#5563) --- pkg/util/net/websocket.go | 39 ++++--- pkg/util/net/websocket_test.go | 182 +++++++++++++++++++++++++++------ 2 files changed, 171 insertions(+), 50 deletions(-) diff --git a/pkg/util/net/websocket.go b/pkg/util/net/websocket.go index 0124fc83..f3636ec0 100644 --- a/pkg/util/net/websocket.go +++ b/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 - closed bool + closeOnce sync.Once } // 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 - // over, so the send has to tolerate a closed channel. A nil error - // means the connection was accepted and is owned by the caller. - if err := gerr.PanicToError(func() { - wl.acceptCh <- conn - }); err != nil { + select { + case <-wl.done: 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 - if !ok { + select { + case <-p.done: return nil, ErrWebsocketListenerClosed + default: + } + select { + case <-p.done: + return nil, ErrWebsocketListenerClosed + case c := <-p.acceptCh: + return c, nil } - return c, nil } func (p *WebsocketListener) Close() error { - p.mu.Lock() - 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() + p.closeOnce.Do(func() { close(p.done) }) return p.server.Close() } diff --git a/pkg/util/net/websocket_test.go b/pkg/util/net/websocket_test.go index f0f6b9d9..baa095ba 100644 --- a/pkg/util/net/websocket_test.go +++ b/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 { - case <-l.closed: - default: - close(l.closed) - } - return nil + l.once.Do(func() { close(l.closed) }) + return l.err } func (l *closeStubListener) Addr() net.Addr { @@ -46,28 +51,147 @@ func (l *closeStubListener) Addr() net.Addr { } func TestWebsocketListenerAcceptUnblocksOnClose(t *testing.T) { - wl := NewWebsocketListener(&closeStubListener{closed: make(chan struct{})}) + synctest.Test(t, func(t *testing.T) { + wl := NewWebsocketListener(&closeStubListener{closed: make(chan struct{})}) + t.Cleanup(func() { wl.Close() }) + errCh := make(chan error, 3) + for range 3 { + go func() { + c, err := wl.Accept() + if c != nil { + t.Error("Accept returned a connection without a sender") + } + errCh <- err + }() + } + synctest.Wait() // All Accept calls are blocked before Close. + require.Empty(t, errCh) - errCh := make(chan error, 1) - go func() { - _, err := wl.Accept() - errCh <- err - }() - - select { - case err := <-errCh: - t.Fatalf("Accept returned before the listener was closed: %v", err) - case <-time.After(100 * time.Millisecond): - } - - require.NoError(t, wl.Close()) - // Closing twice must stay safe and keep reporting the closed error. - require.NoError(t, wl.Close()) - - select { - 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") - } + var wg sync.WaitGroup + for range 3 { + wg.Go(func() { + if err := wl.Close(); err != nil { + t.Errorf("Close: %v", err) + } + }) + } + 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) + }) }