From 6fb8486760bbdc87f614360e3b4042a90a9e37cc Mon Sep 17 00:00:00 2001 From: Rajeh Taher Date: Mon, 7 Sep 2026 00:23:15 +0300 Subject: [PATCH 1/4] socket: race both chat endpoints and close the loser the way the client does WhatsApp Web opens wss://web.whatsapp.com/ws/chat and wss://web.whatsapp.com:5222/ws/chat concurrently on every connect and closes the one that loses with code 1000 and the reason "loser socket" (WAWebOpenSocket.js:10, 44-52). A client that makes exactly one attempt per connect and never closes a second socket differs from the real one by nothing more than counting. - socket.ConnectRace opens every URL concurrently, keeps the first to connect, aborts the dials still outstanding and closes an already-open loser with the client's own code and reason. - FrameSocket.CloseWithReason carries an explicit close reason; Close keeps its signature and passes an empty one. - SocketConfig.RaceURLs opts a client in; empty keeps the single-socket path. - Client.DisablePostConnectPassiveIQ stops the library sending the passive/active IQ after connect, which the real client never sends: WhatsApp Web puts passive:false in the login payload itself (WAWebGetClientPayloadForLogin.js:14-19). --- client.go | 31 +++++- connectionevents.go | 13 ++- socket/framesocket.go | 13 ++- socket/race.go | 133 ++++++++++++++++++++++++++ socket/race_test.go | 216 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 400 insertions(+), 6 deletions(-) create mode 100644 socket/race.go create mode 100644 socket/race_test.go diff --git a/client.go b/client.go index f809407df..78e5ed894 100644 --- a/client.go +++ b/client.go @@ -262,6 +262,11 @@ type Client struct { // The user agent to use (for non-Messenger connections). UserAgent string WebSocketHeaders http.Header + // DisablePostConnectPassiveIQ stops the library sending the + // `passive`/`active` IQ after a successful connect. Set it when the + // client payload already carries `passive: false`, which is what + // WhatsApp Web sends; the real client has no such IQ. + DisablePostConnectPassiveIQ bool } type groupMetaCache struct { @@ -281,6 +286,19 @@ type SocketConfig struct { URL string Origin string NoiseCertificateAuthority *[32]byte + // RaceURLs, when it has more than one entry, makes Connect open every + // URL concurrently and keep the one that connects first, closing the + // others with code 1000 and the reason "loser socket". + // + // This is what WhatsApp Web does on every connect + // (WAWebOpenSocket.js:10, 44-52). A client that makes exactly one + // attempt per connect and never closes a second socket differs from the + // real one by nothing more than counting. Use socket.RaceURLs for the + // endpoints the client itself races. + // + // Empty (the default) keeps the single-socket behaviour and the URL + // field above. + RaceURLs []string } const handlerQueueSize = 256 @@ -610,6 +628,7 @@ func (cli *Client) unlockedConnect(ctx context.Context) error { fs.HTTPHeaders.Set("Origin", cli.MessengerConfig.BaseURL) } maps.Copy(fs.HTTPHeaders, cli.WebSocketHeaders) + var raceURLs []string if cli.SocketConfig != nil { if cli.SocketConfig.URL != "" { fs.URL = cli.SocketConfig.URL @@ -617,11 +636,19 @@ func (cli *Client) unlockedConnect(ctx context.Context) error { if cli.SocketConfig.Origin != "" { fs.HTTPHeaders.Set("Origin", cli.SocketConfig.Origin) } + raceURLs = cli.SocketConfig.RaceURLs } - if err := fs.Connect(ctx); err != nil { + if len(raceURLs) > 1 { + raced, err := socket.ConnectRace(ctx, cli.Log.Sub("Socket"), client, raceURLs, fs.HTTPHeaders) + if err != nil { + return err + } + fs = raced + } else if err := fs.Connect(ctx); err != nil { fs.Close(0) return err - } else if err = cli.doHandshake(ctx, fs, *keys.NewKeyPair()); err != nil { + } + if err := cli.doHandshake(ctx, fs, *keys.NewKeyPair()); err != nil { fs.Close(0) return fmt.Errorf("noise handshake failed: %w", err) } diff --git a/connectionevents.go b/connectionevents.go index 3e52853d3..1892b364a 100644 --- a/connectionevents.go +++ b/connectionevents.go @@ -197,9 +197,16 @@ func (cli *Client) handleConnectSuccess(ctx context.Context, node *waBinary.Node cli.Log.Debugf("Prekey count after upload: %d", sc) } } - err := cli.SetPassive(ctx, false) - if err != nil { - cli.Log.Warnf("Failed to send post-connect passive IQ: %v", err) + // WhatsApp Web sends `passive: false` in the login payload itself + // (WAWebGetClientPayloadForLogin.js:14-19) and never sends this IQ. + // A client that logs in passive and then immediately asks to become + // active does something no real page does, so a caller that already + // carries the right value in its payload turns this off. + if !cli.DisablePostConnectPassiveIQ { + err := cli.SetPassive(ctx, false) + if err != nil { + cli.Log.Warnf("Failed to send post-connect passive IQ: %v", err) + } } cli.dispatchEvent(&events.Connected{}) cli.closeSocketWaitChan() diff --git a/socket/framesocket.go b/socket/framesocket.go index bcf107b2f..fe4501db8 100644 --- a/socket/framesocket.go +++ b/socket/framesocket.go @@ -65,6 +65,17 @@ func (fs *FrameSocket) IsConnected() bool { } func (fs *FrameSocket) Close(code websocket.StatusCode) { + fs.CloseWithReason(code, "") +} + +// CloseWithReason closes the socket with an explicit close reason. +// +// WhatsApp Web races two websockets on every connect and closes the one that +// loses with code 1000 and the reason "loser socket" +// (WAWebOpenSocket.js:44-52). A close with an empty reason where the client +// sends one is observable, so the reason is part of the wire behaviour and not +// a log detail. +func (fs *FrameSocket) CloseWithReason(code websocket.StatusCode, reason string) { fs.lock.Lock() defer fs.lock.Unlock() @@ -75,7 +86,7 @@ func (fs *FrameSocket) Close(code websocket.StatusCode) { fs.closed.Store(true) if code > 0 { - err := conn.Close(code, "") + err := conn.Close(code, reason) if err != nil { fs.log.Warnf("Error sending close to websocket: %v", err) } diff --git a/socket/race.go b/socket/race.go new file mode 100644 index 000000000..857866265 --- /dev/null +++ b/socket/race.go @@ -0,0 +1,133 @@ +//go:build !js + +package socket + +import ( + "context" + "errors" + "net/http" + "sync" + + "github.com/coder/websocket" + + waLog "go.mau.fi/whatsmeow/util/log" +) + +// LoserSocketCloseReason is the close reason WhatsApp Web sends on the socket +// that loses the race (WAWebOpenSocket.js:47). The code is 1000. +const LoserSocketCloseReason = "loser socket" + +// RaceURLs are the two endpoints WhatsApp Web opens concurrently on every +// connect (WAWebOpenSocket.js:10). The second exists because some networks +// block or throttle 443 for long-lived upgrades, and which one wins is a +// property of the network the client is on. +var RaceURLs = []string{ + "wss://web.whatsapp.com/ws/chat", + "wss://web.whatsapp.com:5222/ws/chat", +} + +// ConnectRace opens every URL concurrently and returns the frame socket that +// connected first, closing the others the way the client closes them: code +// 1000 with the reason "loser socket". It mirrors +// openWebSocketsConcurrently (WAWebOpenSocket.js:44-64) — the first success +// wins, the outstanding dials are aborted, and the call fails only when every +// URL failed. +// +// A client that makes one attempt per connect and never closes a second socket +// is distinguishable from the real one by nothing more than counting, which is +// why this is worth having. +// +// The returned socket is connected and its read pump is running; the caller +// owns it exactly as it owns one from NewFrameSocket plus Connect. +func ConnectRace( + ctx context.Context, + log waLog.Logger, + httpClient *http.Client, + urls []string, + headers http.Header, +) (*FrameSocket, error) { + if len(urls) == 0 { + return nil, ErrDialFailed + } + if len(urls) == 1 { + fs := newRacer(log, httpClient, urls[0], headers) + if err := fs.Connect(ctx); err != nil { + fs.Close(0) + return nil, err + } + return fs, nil + } + + // The decision is taken inside the racer goroutines under one lock, so a + // socket that opens at the same instant as the winner is chosen either + // becomes the winner or closes itself as the loser — there is no window + // in which an already-open loser has its context cancelled from outside + // and is torn down before it can send the close frame. + var ( + mu sync.Mutex + winner *FrameSocket + errs []error + ) + cancels := make([]context.CancelFunc, len(urls)) + var wg sync.WaitGroup + for i, url := range urls { + racerCtx, cancel := context.WithCancel(ctx) + cancels[i] = cancel + wg.Add(1) + go func(i int, url string, racerCtx context.Context) { + defer wg.Done() + fs := newRacer(log, httpClient, url, headers) + if err := fs.Connect(racerCtx); err != nil { + fs.Close(0) + mu.Lock() + errs = append(errs, err) + mu.Unlock() + return + } + mu.Lock() + if winner == nil { + winner = fs + mu.Unlock() + log.Debugf("Opened socket with %s (race winner)", url) + // Abort every dial still outstanding, exactly as the client + // aborts its AbortController on the first success. The + // winner's own context stays live: it is the parent of the + // socket's context. + for j, cancel := range cancels { + if j != i { + cancel() + } + } + return + } + mu.Unlock() + // A socket that opened after the winner is closed with the + // client's own code and reason. + fs.CloseWithReason(websocket.StatusNormalClosure, LoserSocketCloseReason) + }(i, url, racerCtx) + } + wg.Wait() + + if winner == nil { + for _, cancel := range cancels { + cancel() + } + if len(errs) == 0 { + return nil, ErrDialFailed + } + return nil, errors.Join(errs...) + } + return winner, nil +} + +func newRacer(log waLog.Logger, httpClient *http.Client, url string, headers http.Header) *FrameSocket { + fs := NewFrameSocket(log, httpClient) + fs.URL = url + for name, values := range headers { + if len(values) == 0 { + continue + } + fs.HTTPHeaders.Set(name, values[0]) + } + return fs +} diff --git a/socket/race_test.go b/socket/race_test.go new file mode 100644 index 000000000..456949033 --- /dev/null +++ b/socket/race_test.go @@ -0,0 +1,216 @@ +//go:build !js + +package socket + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/coder/websocket" + + waLog "go.mau.fi/whatsmeow/util/log" +) + +// wsServer accepts websocket upgrades and records how each connection was +// closed, so a test can assert the loser got the client's own code and reason. +type wsServer struct { + *httptest.Server + delay time.Duration + + mu sync.Mutex + closes []closeRecord + opened int + handled sync.WaitGroup +} + +type closeRecord struct { + code websocket.StatusCode + reason string +} + +func newWSServer(t *testing.T, delay time.Duration) *wsServer { + t.Helper() + s := &wsServer{delay: delay} + s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if s.delay > 0 { + time.Sleep(s.delay) + } + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}}) + if err != nil { + return + } + s.mu.Lock() + s.opened++ + s.mu.Unlock() + s.handled.Add(1) + defer s.handled.Done() + // Block until the peer closes, then record how. + _, _, readErr := conn.Read(r.Context()) + status := websocket.CloseStatus(readErr) + reason := "" + var ce websocket.CloseError + if errors.As(readErr, &ce) { + reason = ce.Reason + } + s.mu.Lock() + s.closes = append(s.closes, closeRecord{code: status, reason: reason}) + s.mu.Unlock() + _ = conn.CloseNow() + })) + t.Cleanup(s.Close) + return s +} + +func (s *wsServer) wsURL() string { return "ws" + s.URL[len("http"):] } + +// awaitClose waits for the server to record one close and returns the record. +func (s *wsServer) awaitClose(t *testing.T) []closeRecord { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if got := s.recorded(); len(got) > 0 { + return got + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("the socket was never closed") + return nil +} + +func (s *wsServer) recorded() []closeRecord { + s.mu.Lock() + defer s.mu.Unlock() + return append([]closeRecord(nil), s.closes...) +} + +// The close frame a loser sends is the client's own: code 1000 with the +// reason "loser socket" (WAWebOpenSocket.js:47). An empty reason where the +// client sends one is observable on the wire, so this is the behaviour, not a +// log detail. +func TestCloseWithReasonSendsTheClientsCloseFrame(t *testing.T) { + server := newWSServer(t, 0) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + fs := NewFrameSocket(waLog.Noop, http.DefaultClient) + fs.URL = server.wsURL() + if err := fs.Connect(ctx); err != nil { + t.Fatalf("connect: %v", err) + } + fs.CloseWithReason(websocket.StatusNormalClosure, LoserSocketCloseReason) + + closes := server.awaitClose(t) + if closes[0].code != websocket.StatusNormalClosure { + t.Errorf("close code = %d, want %d", closes[0].code, websocket.StatusNormalClosure) + } + if closes[0].reason != LoserSocketCloseReason { + t.Errorf("close reason = %q, want %q", closes[0].reason, LoserSocketCloseReason) + } +} + +// The socket that connects first is the one kept. +func TestConnectRaceKeepsTheFirstToConnect(t *testing.T) { + fast := newWSServer(t, 0) + slow := newWSServer(t, 150*time.Millisecond) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + fs, err := ConnectRace(ctx, waLog.Noop, http.DefaultClient, + []string{fast.wsURL(), slow.wsURL()}, http.Header{"Origin": {Origin}}) + if err != nil { + t.Fatalf("race: %v", err) + } + defer fs.Close(websocket.StatusNormalClosure) + + if fs.URL != fast.wsURL() { + t.Fatalf("the slow endpoint won the race: %s", fs.URL) + } + if !fs.IsConnected() { + t.Fatal("the winner is not connected") + } +} + +// With both endpoints answering at once, exactly one socket survives, and a +// loser that did open is closed the client's way rather than dropped. +func TestConnectRaceLeavesOneSocket(t *testing.T) { + a := newWSServer(t, 0) + b := newWSServer(t, 0) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + fs, err := ConnectRace(ctx, waLog.Noop, http.DefaultClient, []string{a.wsURL(), b.wsURL()}, nil) + if err != nil { + t.Fatalf("race: %v", err) + } + defer fs.Close(websocket.StatusNormalClosure) + if !fs.IsConnected() { + t.Fatal("the winner is not connected") + } + if fs.URL != a.wsURL() && fs.URL != b.wsURL() { + t.Fatalf("the race returned an endpoint nobody offered: %s", fs.URL) + } + + // ConnectRace has already waited for every racer, so any loser that + // opened has been closed by now. + for _, rec := range append(a.recorded(), b.recorded()...) { + if rec.code != websocket.StatusNormalClosure || rec.reason != LoserSocketCloseReason { + t.Errorf("a loser was closed with (%d, %q), want (%d, %q)", + rec.code, rec.reason, websocket.StatusNormalClosure, LoserSocketCloseReason) + } + } +} + +// One URL is the old behaviour, unchanged. +func TestConnectRaceWithOneURL(t *testing.T) { + only := newWSServer(t, 0) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + fs, err := ConnectRace(ctx, waLog.Noop, http.DefaultClient, []string{only.wsURL()}, nil) + if err != nil { + t.Fatalf("race: %v", err) + } + if fs.URL != only.wsURL() || !fs.IsConnected() { + t.Fatal("the single-URL path did not return a connected socket") + } + fs.Close(websocket.StatusNormalClosure) +} + +// Every endpoint failing fails the call, and the errors are all reported. +func TestConnectRaceAllFail(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err := ConnectRace(ctx, waLog.Noop, http.DefaultClient, + []string{"ws://127.0.0.1:1/a", "ws://127.0.0.1:1/b"}, nil) + if err == nil { + t.Fatal("a race where every endpoint failed returned a socket") + } +} + +// A racer that loses before it ever opened is aborted rather than left dialing. +func TestConnectRaceAbortsOutstandingDials(t *testing.T) { + fast := newWSServer(t, 0) + slow := newWSServer(t, 3*time.Second) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + start := time.Now() + fs, err := ConnectRace(ctx, waLog.Noop, http.DefaultClient, + []string{fast.wsURL(), slow.wsURL()}, nil) + if err != nil { + t.Fatalf("race: %v", err) + } + defer fs.Close(websocket.StatusNormalClosure) + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("the race waited %v for the slow endpoint", elapsed) + } +} From 9c6bf2c91afeb7f098946613fc310ec89b1729e9 Mon Sep 17 00:00:00 2001 From: Rajeh Taher Date: Tue, 8 Sep 2026 04:14:19 +0300 Subject: [PATCH 2/4] fix(socket): initialize race cancellation before dialing --- socket/race.go | 10 ++++++---- socket/race_test.go | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/socket/race.go b/socket/race.go index 857866265..9775e5a63 100644 --- a/socket/race.go +++ b/socket/race.go @@ -10,7 +10,7 @@ import ( "github.com/coder/websocket" - waLog "go.mau.fi/whatsmeow/util/log" + waLog "github.com/polymorfa/hypermeow/util/log" ) // LoserSocketCloseReason is the close reason WhatsApp Web sends on the socket @@ -68,11 +68,13 @@ func ConnectRace( winner *FrameSocket errs []error ) + racerContexts := make([]context.Context, len(urls)) cancels := make([]context.CancelFunc, len(urls)) + for i := range urls { + racerContexts[i], cancels[i] = context.WithCancel(ctx) + } var wg sync.WaitGroup for i, url := range urls { - racerCtx, cancel := context.WithCancel(ctx) - cancels[i] = cancel wg.Add(1) go func(i int, url string, racerCtx context.Context) { defer wg.Done() @@ -104,7 +106,7 @@ func ConnectRace( // A socket that opened after the winner is closed with the // client's own code and reason. fs.CloseWithReason(websocket.StatusNormalClosure, LoserSocketCloseReason) - }(i, url, racerCtx) + }(i, url, racerContexts[i]) } wg.Wait() diff --git a/socket/race_test.go b/socket/race_test.go index 456949033..ee5e97455 100644 --- a/socket/race_test.go +++ b/socket/race_test.go @@ -13,7 +13,7 @@ import ( "github.com/coder/websocket" - waLog "go.mau.fi/whatsmeow/util/log" + waLog "github.com/polymorfa/hypermeow/util/log" ) // wsServer accepts websocket upgrades and records how each connection was From 639a4b8b05bc42fb433b30e8120da74c91f12dd2 Mon Sep 17 00:00:00 2001 From: Rajeh Taher Date: Tue, 8 Sep 2026 04:27:10 +0300 Subject: [PATCH 3/4] test(socket): distinguish canceled in-flight racers --- socket/race.go | 2 +- socket/race_test.go | 32 +++++++++++++++++++++++++------- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/socket/race.go b/socket/race.go index 9775e5a63..94397794e 100644 --- a/socket/race.go +++ b/socket/race.go @@ -129,7 +129,7 @@ func newRacer(log waLog.Logger, httpClient *http.Client, url string, headers htt if len(values) == 0 { continue } - fs.HTTPHeaders.Set(name, values[0]) + fs.HTTPHeaders[name] = append([]string(nil), values...) } return fs } diff --git a/socket/race_test.go b/socket/race_test.go index ee5e97455..1948fdf66 100644 --- a/socket/race_test.go +++ b/socket/race_test.go @@ -113,6 +113,17 @@ func TestCloseWithReasonSendsTheClientsCloseFrame(t *testing.T) { } } +func TestNewRacerPreservesAllHeaderValues(t *testing.T) { + headers := http.Header{"X-Test-Value": {"first", "second"}} + fs := newRacer(waLog.Noop, http.DefaultClient, "ws://example.invalid", headers) + headers["X-Test-Value"][0] = "changed" + + got := fs.HTTPHeaders.Values("X-Test-Value") + if len(got) != 2 || got[0] != "first" || got[1] != "second" { + t.Fatalf("copied header values = %q, want [first second]", got) + } +} + // The socket that connects first is the one kept. func TestConnectRaceKeepsTheFirstToConnect(t *testing.T) { fast := newWSServer(t, 0) @@ -136,8 +147,11 @@ func TestConnectRaceKeepsTheFirstToConnect(t *testing.T) { } } -// With both endpoints answering at once, exactly one socket survives, and a -// loser that did open is closed the client's way rather than dropped. +// With both endpoints answering at once, exactly one client-visible socket +// survives. A loser whose Dial returned is closed the client's way. The server +// may also accept an upgrade just before the winner cancels that still-in-flight +// Dial; coder/websocket then reports an abnormal close because no client socket +// existed yet on which ConnectRace could send a close frame. func TestConnectRaceLeavesOneSocket(t *testing.T) { a := newWSServer(t, 0) b := newWSServer(t, 0) @@ -157,13 +171,17 @@ func TestConnectRaceLeavesOneSocket(t *testing.T) { t.Fatalf("the race returned an endpoint nobody offered: %s", fs.URL) } - // ConnectRace has already waited for every racer, so any loser that - // opened has been closed by now. + // ConnectRace has already waited for every racer, so any client-visible + // loser has sent the explicit close frame by now. TestCloseWithReasonSendsTheClientsCloseFrame + // above pins that frame independently; this loop also permits cancellation + // of a server-accepted upgrade whose Dial had not returned to the racer. for _, rec := range append(a.recorded(), b.recorded()...) { - if rec.code != websocket.StatusNormalClosure || rec.reason != LoserSocketCloseReason { - t.Errorf("a loser was closed with (%d, %q), want (%d, %q)", - rec.code, rec.reason, websocket.StatusNormalClosure, LoserSocketCloseReason) + explicitLoser := rec.code == websocket.StatusNormalClosure && rec.reason == LoserSocketCloseReason + canceledDial := rec.code == -1 && rec.reason == "" + if explicitLoser || canceledDial { + continue } + t.Errorf("a loser was closed with unexpected code and reason (%d, %q)", rec.code, rec.reason) } } From ff99cde0ab2e60bc2c5a29f412e483c5eabf7341 Mon Sep 17 00:00:00 2001 From: Rajeh Taher Date: Tue, 8 Sep 2026 04:54:34 +0300 Subject: [PATCH 4/4] fix(socket): separate racer dial lifetimes --- socket/framesocket.go | 16 ++++++++++---- socket/race.go | 20 ++++++++++++++--- socket/race_test.go | 50 ++++++++++++++++++++++++++++--------------- 3 files changed, 62 insertions(+), 24 deletions(-) diff --git a/socket/framesocket.go b/socket/framesocket.go index fe4501db8..ba3978e5d 100644 --- a/socket/framesocket.go +++ b/socket/framesocket.go @@ -104,16 +104,24 @@ func (fs *FrameSocket) CloseWithReason(code websocket.StatusCode, reason string) } func (fs *FrameSocket) Connect(ctx context.Context) error { + return fs.connect(ctx, ctx) +} + +// connect separates the lifetime of the opened socket from the context used +// for the HTTP upgrade. Most callers use the same context for both through +// Connect. Socket racing uses a short-lived dial context so aborting another +// in-flight upgrade cannot cancel a connection that has already opened. +func (fs *FrameSocket) connect(parentCtx, dialCtx context.Context) error { fs.lock.Lock() defer fs.lock.Unlock() if fs.conn.Load() != nil { return ErrSocketAlreadyOpen } - fs.parentCtx = ctx - fs.cancelCtx, fs.cancel = context.WithCancel(ctx) + fs.parentCtx = parentCtx + fs.cancelCtx, fs.cancel = context.WithCancel(parentCtx) fs.log.Debugf("Dialing %s", fs.URL) - conn, resp, err := websocket.Dial(ctx, fs.URL, fs.makeDialOptions()) + conn, resp, err := websocket.Dial(dialCtx, fs.URL, fs.makeDialOptions()) if err != nil { if resp != nil { err = ErrWithStatusCode{err, resp.StatusCode} @@ -125,7 +133,7 @@ func (fs *FrameSocket) Connect(ctx context.Context) error { fs.conn.Store(conn) - go fs.readPump(conn, ctx) + go fs.readPump(conn, fs.cancelCtx) return nil } diff --git a/socket/race.go b/socket/race.go index 94397794e..9084758c4 100644 --- a/socket/race.go +++ b/socket/race.go @@ -45,6 +45,17 @@ func ConnectRace( httpClient *http.Client, urls []string, headers http.Header, +) (*FrameSocket, error) { + return connectRace(ctx, log, httpClient, urls, headers, nil) +} + +func connectRace( + ctx context.Context, + log waLog.Logger, + httpClient *http.Client, + urls []string, + headers http.Header, + afterConnect func(int), ) (*FrameSocket, error) { if len(urls) == 0 { return nil, ErrDialFailed @@ -78,14 +89,18 @@ func ConnectRace( wg.Add(1) go func(i int, url string, racerCtx context.Context) { defer wg.Done() + defer cancels[i]() fs := newRacer(log, httpClient, url, headers) - if err := fs.Connect(racerCtx); err != nil { + if err := fs.connect(ctx, racerCtx); err != nil { fs.Close(0) mu.Lock() errs = append(errs, err) mu.Unlock() return } + if afterConnect != nil { + afterConnect(i) + } mu.Lock() if winner == nil { winner = fs @@ -93,8 +108,7 @@ func ConnectRace( log.Debugf("Opened socket with %s (race winner)", url) // Abort every dial still outstanding, exactly as the client // aborts its AbortController on the first success. The - // winner's own context stays live: it is the parent of the - // socket's context. + // dial contexts do not own sockets after their upgrades complete. for j, cancel := range cancels { if j != i { cancel() diff --git a/socket/race_test.go b/socket/race_test.go index 1948fdf66..dbf2f1d9d 100644 --- a/socket/race_test.go +++ b/socket/race_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "sync" + "sync/atomic" "testing" "time" @@ -147,11 +148,9 @@ func TestConnectRaceKeepsTheFirstToConnect(t *testing.T) { } } -// With both endpoints answering at once, exactly one client-visible socket -// survives. A loser whose Dial returned is closed the client's way. The server -// may also accept an upgrade just before the winner cancels that still-in-flight -// Dial; coder/websocket then reports an abnormal close because no client socket -// existed yet on which ConnectRace could send a close frame. +// With both upgrades completed before winner selection, exactly one +// client-visible socket survives and the opened loser is closed the client's +// way. func TestConnectRaceLeavesOneSocket(t *testing.T) { a := newWSServer(t, 0) b := newWSServer(t, 0) @@ -159,7 +158,21 @@ func TestConnectRaceLeavesOneSocket(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - fs, err := ConnectRace(ctx, waLog.Noop, http.DefaultClient, []string{a.wsURL(), b.wsURL()}, nil) + // Hold both racers after their upgrades complete. This forces winner + // selection to happen with two opened sockets, so the loser must send the + // explicit close frame rather than being an in-flight canceled dial. + bothConnected := make(chan struct{}) + var connected atomic.Int32 + afterConnect := func(int) { + if connected.Add(1) == 2 { + close(bothConnected) + } + select { + case <-bothConnected: + case <-ctx.Done(): + } + } + fs, err := connectRace(ctx, waLog.Noop, http.DefaultClient, []string{a.wsURL(), b.wsURL()}, nil, afterConnect) if err != nil { t.Fatalf("race: %v", err) } @@ -167,21 +180,24 @@ func TestConnectRaceLeavesOneSocket(t *testing.T) { if !fs.IsConnected() { t.Fatal("the winner is not connected") } + if err := fs.Context().Err(); err != nil { + t.Fatalf("winner context was canceled with its completed dial: %v", err) + } if fs.URL != a.wsURL() && fs.URL != b.wsURL() { t.Fatalf("the race returned an endpoint nobody offered: %s", fs.URL) } - // ConnectRace has already waited for every racer, so any client-visible - // loser has sent the explicit close frame by now. TestCloseWithReasonSendsTheClientsCloseFrame - // above pins that frame independently; this loop also permits cancellation - // of a server-accepted upgrade whose Dial had not returned to the racer. - for _, rec := range append(a.recorded(), b.recorded()...) { - explicitLoser := rec.code == websocket.StatusNormalClosure && rec.reason == LoserSocketCloseReason - canceledDial := rec.code == -1 && rec.reason == "" - if explicitLoser || canceledDial { - continue - } - t.Errorf("a loser was closed with unexpected code and reason (%d, %q)", rec.code, rec.reason) + if connected.Load() != 2 { + t.Fatalf("completed upgrades = %d, want 2", connected.Load()) + } + loser := a + if fs.URL == a.wsURL() { + loser = b + } + rec := loser.awaitClose(t)[0] + if rec.code != websocket.StatusNormalClosure || rec.reason != LoserSocketCloseReason { + t.Fatalf("loser close = (%d, %q), want (%d, %q)", + rec.code, rec.reason, websocket.StatusNormalClosure, LoserSocketCloseReason) } }