diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 00eeae11d..8138665db 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -33,7 +33,8 @@ jobs: - name: Install goimports run: | - go install golang.org/x/tools/cmd/goimports@latest + # Pinned: x/tools @latest needs go >= 1.26, the 1.25 job runs with GOTOOLCHAIN=local. + go install golang.org/x/tools/cmd/goimports@v0.49.0 export PATH="$HOME/go/bin:$PATH" - name: Run pre-commit diff --git a/socket/framesocket.go b/socket/framesocket.go index bcf107b2f..ec2dff33a 100644 --- a/socket/framesocket.go +++ b/socket/framesocket.go @@ -31,7 +31,8 @@ type FrameSocket struct { HTTPHeaders http.Header HTTPClient *http.Client - Frames chan []byte + Frames chan []byte + // Use SetOnDisconnect to change the handler while the socket is in use. OnDisconnect func(ctx context.Context, remote bool) Header []byte @@ -64,6 +65,14 @@ func (fs *FrameSocket) IsConnected() bool { return fs.conn.Load() != nil } +// SetOnDisconnect replaces the disconnect handler under the socket lock. +// It does not cancel a callback already scheduled by Close. +func (fs *FrameSocket) SetOnDisconnect(handler func(ctx context.Context, remote bool)) { + fs.lock.Lock() + defer fs.lock.Unlock() + fs.OnDisconnect = handler +} + func (fs *FrameSocket) Close(code websocket.StatusCode) { fs.lock.Lock() defer fs.lock.Unlock() diff --git a/socket/noisesocket.go b/socket/noisesocket.go index 8175868ce..a4d16df0d 100644 --- a/socket/noisesocket.go +++ b/socket/noisesocket.go @@ -45,9 +45,9 @@ func newNoiseSocket( onFrame: frameHandler, stopConsumer: make(chan struct{}), } - fs.OnDisconnect = func(ctx context.Context, remote bool) { + fs.SetOnDisconnect(func(ctx context.Context, remote bool) { disconnectHandler(ctx, ns, remote) - } + }) go ns.consumeFrames(ctx, fs.Frames) return ns, nil } @@ -80,7 +80,7 @@ func (ns *NoiseSocket) Stop(disconnect, allowOnDisconnect bool) { if ns.destroyed.CompareAndSwap(false, true) { close(ns.stopConsumer) if !allowOnDisconnect { - ns.fs.OnDisconnect = nil + ns.fs.SetOnDisconnect(nil) } if disconnect { ns.fs.Close(websocket.StatusNormalClosure) diff --git a/socket/noisesocket_test.go b/socket/noisesocket_test.go new file mode 100644 index 000000000..73ee1241a --- /dev/null +++ b/socket/noisesocket_test.go @@ -0,0 +1,166 @@ +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. + +package socket + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/coder/websocket" + + waLog "github.com/polymorfa/hypermeow/util/log" +) + +func newTestFrameSocket(t *testing.T) *FrameSocket { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true}) + if err != nil { + t.Errorf("Failed to accept websocket: %v", err) + return + } + defer func() { _ = conn.CloseNow() }() + <-conn.CloseRead(r.Context()).Done() + })) + t.Cleanup(server.Close) + fs := NewFrameSocket(waLog.Noop, server.Client()) + fs.URL = server.URL + if err := fs.Connect(t.Context()); err != nil { + t.Fatalf("Failed to connect websocket: %v", err) + } + t.Cleanup(func() { fs.Close(0) }) + return fs +} + +func TestNoiseSocketStopConcurrentClose(t *testing.T) { + for range 100 { + fs := newTestFrameSocket(t) + ns, err := newNoiseSocket(t.Context(), fs, nil, nil, nil, func(context.Context, *NoiseSocket, bool) {}) + if err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + start := make(chan struct{}) + wg.Go(func() { + <-start + fs.Close(0) + }) + wg.Go(func() { + <-start + ns.Stop(true, false) + }) + close(start) + wg.Wait() + } +} + +func TestNewNoiseSocketConcurrentClose(t *testing.T) { + for range 100 { + fs := newTestFrameSocket(t) + var wg sync.WaitGroup + start := make(chan struct{}) + wg.Go(func() { + <-start + fs.Close(0) + }) + close(start) + ns, err := newNoiseSocket(t.Context(), fs, nil, nil, nil, func(context.Context, *NoiseSocket, bool) {}) + if err != nil { + t.Fatal(err) + } + wg.Wait() + ns.Stop(false, false) + } +} + +func TestNoiseSocketStop(t *testing.T) { + for _, tc := range []struct { + name string + disconnect bool + allowOnDisconnect bool + }{ + {name: "local suppressed", disconnect: true}, + {name: "local allowed", disconnect: true, allowOnDisconnect: true}, + {name: "remote suppressed"}, + {name: "remote allowed", allowOnDisconnect: true}, + } { + t.Run(tc.name, func(t *testing.T) { + fs := newTestFrameSocket(t) + called := make(chan bool, 2) + ns, err := newNoiseSocket(t.Context(), fs, nil, nil, nil, func(ctx context.Context, socket *NoiseSocket, remote bool) { + if ctx != t.Context() { + t.Error("Unexpected disconnect context") + } + socket.Stop(false, false) + fs.Close(0) + called <- remote + }) + if err != nil { + t.Fatal(err) + } + ns.Stop(tc.disconnect, tc.allowOnDisconnect) + if fs.IsConnected() == tc.disconnect { + t.Fatalf("Unexpected connection state after Stop: %t", fs.IsConnected()) + } + select { + case <-ns.stopConsumer: + default: + t.Fatal("Frame consumer was not stopped") + } + fs.lock.Lock() + hasHandler := fs.OnDisconnect != nil + fs.lock.Unlock() + if hasHandler != tc.allowOnDisconnect { + t.Fatalf("Unexpected disconnect handler after Stop: %t", hasHandler) + } + fs.Close(0) + fs.Close(0) + if tc.allowOnDisconnect { + select { + case remote := <-called: + if remote == tc.disconnect { + t.Errorf("Unexpected remote flag: %t", remote) + } + case <-time.After(5 * time.Second): + t.Fatal("Disconnect handler did not return") + } + } + select { + case <-called: + t.Fatal("Unexpected disconnect callback") + default: + } + }) + } +} + +func TestNoiseSocketRemoteDisconnect(t *testing.T) { + fs := newTestFrameSocket(t) + called := make(chan bool, 1) + ns, err := newNoiseSocket(t.Context(), fs, nil, nil, nil, func(ctx context.Context, socket *NoiseSocket, remote bool) { + socket.Stop(false, false) + fs.Close(0) + called <- remote + }) + if err != nil { + t.Fatal(err) + } + fs.Close(0) + select { + case remote := <-called: + if !remote { + t.Error("Expected remote disconnect") + } + case <-time.After(5 * time.Second): + t.Fatal("Disconnect handler did not return") + } + if !ns.destroyed.Load() { + t.Fatal("Disconnect handler did not stop the noise socket") + } +}