diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 7ec1c1f2c..ee6ecc58a 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -32,6 +32,9 @@ jobs: - name: Build run: go build -v ./... + - name: Build JS/WASM + run: GOOS=js GOARCH=wasm go build ./... + - name: Test run: go test -v ./... diff --git a/CHANGELOG.md b/CHANGELOG.md index a4a8fe12b..4f474f90b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,11 @@ All notable HyperMeow changes are documented here. HyperMeow uses commit pseudo- ## [Unreleased] +### Added + +- Opt-in concurrent WebSocket connections through `SocketConfig.RaceURLs`, with asynchronous loser cleanup and support for native and JS/WASM builds. See [connection options](README.md#connection-options). +- `Client.DisablePostConnectPassiveIQ` to skip the post-connect passive/active IQ when the login payload already declares an active connection. + ### Fixed - Keep connection handler completion pending until an active handler exits, and process retired stream errors before automatic reconnect decisions. Manual disconnect retains its five-second wait bound. diff --git a/README.md b/README.md index eff8af791..3558d84ca 100644 --- a/README.md +++ b/README.md @@ -108,6 +108,45 @@ Custom `store.ChatSettingsStore` implementations must add `PutWASARootSecretID` and `GetWASARootSecretID` for bot root-secret state. `sqlstore` adds this field through schema version 19; its existing migrations remain unchanged. +## Connection options + +Connections use one WebSocket endpoint by default. To try both WhatsApp chat +endpoints concurrently, configure the client before connecting: + +```go +import ( + whatsmeow "github.com/polymorfa/hypermeow" + "github.com/polymorfa/hypermeow/socket" +) + +// client is an existing *whatsmeow.Client. +client.SocketConfig = &whatsmeow.SocketConfig{ + RaceURLs: append([]string(nil), socket.RaceURLs...), +} +``` + +`SocketConfig.RaceURLs` activates racing when it contains at least two URLs. +The first successful upgrade wins. Pending upgrades are canceled, and opened +losers close asynchronously with status 1000 and reason `loser socket`. +The winner proceeds without waiting for those close handshakes. If every +endpoint fails, the returned error joins the failures in configured URL order. +Initial background reconnect classification uses the first matching HTTP status +in configured URL order, ahead of network errors from other endpoints. When no +HTTP status is present, network and dial failures remain retryable. The existing +single-endpoint retry rules remain unchanged. + +With zero or one race URL, the client uses `SocketConfig.URL`, or the default +endpoint if `URL` is empty. With racing enabled, the list supplies the endpoints +and `URL` is ignored. `SocketConfig.Origin` and `Client.WebSocketHeaders` apply to +all racers. Repeated header values and empty entries that suppress defaults are +preserved. Browsers control request headers on JS/WASM builds, as they do for +single-endpoint connections. + +`Client.DisablePostConnectPassiveIQ` defaults to `false`. Set it to `true` only +when your login payload already declares `passive: false`; the client then skips +the post-connect passive/active IQ. This option does not change the login payload. +Configure both options before connecting rather than while a connection is active. + ## Discord Join the [WhatsApp Web ecosystem Discord](https://whiskey.so/discord) and visit diff --git a/client.go b/client.go index c44807e66..f5983e246 100644 --- a/client.go +++ b/client.go @@ -282,6 +282,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 { @@ -301,6 +306,15 @@ 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". + // + // Use socket.RaceURLs for the endpoints WhatsApp Web races + // (WAWebOpenSocket.js:10, 44-52). Loser cleanup runs asynchronously. + // + // Zero or one entry keeps the single-socket behavior and uses URL above. + RaceURLs []string } const handlerQueueSize = 256 @@ -563,10 +577,8 @@ func (cli *Client) Connect() error { } func isRetryableConnectError(err error) bool { - if exhttp.IsNetworkError(err) { - return true - } - + // HTTP upgrade rejections take precedence over network failures from + // other racers. Joined errors preserve the configured endpoint order. var statusErr socket.ErrWithStatusCode if errors.As(err, &statusErr) { switch statusErr.StatusCode { @@ -577,6 +589,9 @@ func isRetryableConnectError(err error) bool { } } + if exhttp.IsNetworkError(err) { + return true + } return errors.Is(err, socket.ErrDialFailed) } @@ -641,6 +656,7 @@ func (cli *Client) unlockedConnect(ctx context.Context) error { } var queue chan *waBinary.Node maps.Copy(fs.HTTPHeaders, cli.WebSocketHeaders) + var raceURLs []string if cli.SocketConfig != nil { if cli.SocketConfig.URL != "" { fs.URL = cli.SocketConfig.URL @@ -648,11 +664,20 @@ 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 queue, err = cli.doHandshake(fs, *keys.NewKeyPair()); err != nil { + } + var err error + if queue, err = cli.doHandshake(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 c16fcdf2c..54718ac2a 100644 --- a/socket/framesocket.go +++ b/socket/framesocket.go @@ -74,6 +74,17 @@ func (fs *FrameSocket) SetOnDisconnect(handler func(ctx context.Context, remote } 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() @@ -84,7 +95,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) } @@ -102,16 +113,24 @@ func (fs *FrameSocket) Close(code websocket.StatusCode) { } func (fs *FrameSocket) Connect(ctx context.Context) error { + return fs.dial(ctx, ctx, true) +} + +// dial 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) dial(parentCtx, dialCtx context.Context, startPump bool) 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} @@ -122,11 +141,20 @@ func (fs *FrameSocket) Connect(ctx context.Context) error { conn.SetReadLimit(FrameMaxSize) fs.conn.Store(conn) - - go fs.readPump(conn, ctx) + if startPump { + go fs.readPump(conn, fs.cancelCtx) + } return nil } +func (fs *FrameSocket) startReadPump() { + fs.lock.Lock() + defer fs.lock.Unlock() + if conn := fs.conn.Load(); conn != nil { + go fs.readPump(conn, fs.cancelCtx) + } +} + func (fs *FrameSocket) Context() context.Context { return fs.cancelCtx } @@ -163,13 +191,18 @@ func (fs *FrameSocket) SendFrame(data []byte) error { return conn.Write(fs.cancelCtx, websocket.MessageBinary, wholeFrame) } -func (fs *FrameSocket) frameComplete() { +func (fs *FrameSocket) frameComplete() bool { data := fs.incoming fs.incoming = nil fs.partialHeader = nil fs.incomingLength = 0 fs.receivedLength = 0 - fs.Frames <- data + select { + case fs.Frames <- data: + return true + case <-fs.cancelCtx.Done(): + return false + } } func (fs *FrameSocket) processData(msg []byte) { @@ -188,7 +221,9 @@ func (fs *FrameSocket) processData(msg []byte) { if len(msg) >= length { fs.incoming = msg[:length] msg = msg[length:] - fs.frameComplete() + if !fs.frameComplete() { + return + } } else { fs.incoming = make([]byte, length) copy(fs.incoming, msg) @@ -203,7 +238,9 @@ func (fs *FrameSocket) processData(msg []byte) { if fs.receivedLength+len(msg) >= fs.incomingLength { copy(fs.incoming[fs.receivedLength:], msg[:fs.incomingLength-fs.receivedLength]) msg = msg[fs.incomingLength-fs.receivedLength:] - fs.frameComplete() + if !fs.frameComplete() { + return + } } else { copy(fs.incoming[fs.receivedLength:], msg) fs.receivedLength += len(msg) diff --git a/socket/race.go b/socket/race.go new file mode 100644 index 000000000..37d64c7ad --- /dev/null +++ b/socket/race.go @@ -0,0 +1,123 @@ +package socket + +import ( + "context" + "errors" + "net/http" + + "github.com/coder/websocket" + + waLog "github.com/polymorfa/hypermeow/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. +// +// Loser cleanup runs asynchronously. If all endpoints fail, errors are joined +// in endpoint order. The returned socket is connected and its read pump is +// running; the caller owns it 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) { + 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 + } + 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 + } + + type dialResult struct { + index int + fs *FrameSocket + err error + } + results := make(chan dialResult, len(urls)) + cancels := make([]context.CancelFunc, len(urls)) + for i, url := range urls { + dialCtx, cancel := context.WithCancel(ctx) + cancels[i] = cancel + go func() { + fs := newRacer(log, httpClient, url, headers) + err := fs.dial(ctx, dialCtx, false) + cancel() + if err == nil && afterConnect != nil { + afterConnect(i) + } + results <- dialResult{i, fs, err} + }() + } + + // Preserve endpoint order in the joined error so retry classification is + // independent of the order in which unsuccessful dials finish. + errs := make([]error, len(urls)) + for remaining := len(urls); remaining > 0; remaining-- { + result := <-results + if result.err != nil { + errs[result.index] = result.err + continue + } + for _, cancel := range cancels { + cancel() + } + result.fs.startReadPump() + log.Debugf("Opened socket with %s (race winner)", result.fs.URL) + // Upgraded losers have no frame consumer. Never start their read pumps, + // and don't delay the winner while waiting for their close handshakes. + go func() { + for range remaining - 1 { + loser := <-results + if loser.err == nil { + go loser.fs.CloseWithReason(websocket.StatusNormalClosure, LoserSocketCloseReason) + } + } + }() + return result.fs, nil + } + return nil, errors.Join(errs...) +} + +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 { + fs.HTTPHeaders[name] = append([]string(nil), values...) + } + return fs +} diff --git a/socket/race_test.go b/socket/race_test.go new file mode 100644 index 000000000..7cdc5a2d0 --- /dev/null +++ b/socket/race_test.go @@ -0,0 +1,503 @@ +//go:build !js + +package socket + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/coder/websocket" + + waLog "github.com/polymorfa/hypermeow/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) + } +} + +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) + 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 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) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + // 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) + } + defer fs.Close(websocket.StatusNormalClosure) + 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) + } + + 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) + } +} + +// 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) + } +} + +func TestConnectRaceHeadersReachServer(t *testing.T) { + for _, origin := range [][]string{nil, {}} { + for _, endpoints := range []int{1, 2} { + t.Run(fmt.Sprintf("origin=%v/endpoints=%d", origin, endpoints), func(t *testing.T) { + seen := make(chan http.Header, endpoints) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen <- r.Header.Clone() + conn, err := websocket.Accept(w, r, nil) + if err == nil { + defer conn.CloseNow() + _, _, _ = conn.Read(r.Context()) + } + })) + defer server.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + urls := make([]string, endpoints) + for i := range urls { + urls[i] = "ws" + server.URL[len("http"):] + } + fs, err := ConnectRace(ctx, waLog.Noop, http.DefaultClient, urls, + http.Header{"Origin": origin, "X-Test-Value": {"first", "second"}}) + if err != nil { + t.Fatal(err) + } + defer fs.Close(websocket.StatusNormalClosure) + headers := <-seen + if got := headers.Values("Origin"); len(got) != 0 { + t.Errorf("suppressed Origin reached server: %q", got) + } + if got := headers.Values("X-Test-Value"); !slices.Equal(got, []string{"first", "second"}) { + t.Errorf("repeated header = %q", got) + } + }) + } + } +} + +func TestConnectRaceErrorsFollowEndpointOrder(t *testing.T) { + secondFinished := make(chan struct{}) + first := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-secondFinished: + case <-r.Context().Done(): + return + } + w.WriteHeader(http.StatusForbidden) + })) + defer first.Close() + second := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer second.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + // The first endpoint cannot fail until the second response was received. + client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) { + resp, err := http.DefaultTransport.RoundTrip(r) + if r.URL.Host == strings.TrimPrefix(second.URL, "http://") { + close(secondFinished) + } + return resp, err + })} + _, err := ConnectRace(ctx, waLog.Noop, client, + []string{"ws" + first.URL[4:], "ws" + second.URL[4:]}, nil) + var statusErr ErrWithStatusCode + if !errors.As(err, &statusErr) || statusErr.StatusCode != http.StatusForbidden { + t.Fatalf("first status = %v, want 403", err) + } + joined, ok := err.(interface{ Unwrap() []error }) + if !ok || len(joined.Unwrap()) != 2 || !errors.As(joined.Unwrap()[1], &statusErr) || statusErr.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("second endpoint error missing: %v", err) + } +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +type pumpLogger struct { + waLog.Logger + started chan struct{} + exited chan struct{} +} + +func (l *pumpLogger) Debugf(msg string, _ ...any) { + switch msg { + case "Frame websocket read pump starting %p": + l.started <- struct{}{} + case "Frame websocket read pump exiting %p": + l.exited <- struct{}{} + } +} + +func TestConnectRaceReturnsBeforeLoserCloseAndStartsOnlyWinnerPump(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + release := make(chan struct{}) + var releaseOnce sync.Once + releaseServers := func() { releaseOnce.Do(func() { close(release) }) } + defer releaseServers() + closes := make(chan closeRecord, 2) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}}) + if err != nil { + return + } + defer conn.CloseNow() + // Send a complete pre-handshake frame. An unconsumed loser must not + // retain a read pump, and closing the unconsumed winner must release it. + _ = conn.Write(ctx, websocket.MessageBinary, []byte{0, 0, 1, 42}) + select { + case <-release: + case <-ctx.Done(): + return + } + _, _, err = conn.Read(ctx) + var closeErr websocket.CloseError + errors.As(err, &closeErr) + closes <- closeRecord{websocket.CloseStatus(err), closeErr.Reason} + })) + defer server.Close() + logger := &pumpLogger{Logger: waLog.Noop, started: make(chan struct{}, 2), exited: make(chan struct{}, 2)} + bothConnected := make(chan struct{}) + var connected atomic.Int32 + afterConnect := func(int) { + if connected.Add(1) == 2 { + close(bothConnected) + } + select { + case <-bothConnected: + case <-ctx.Done(): + } + } + type outcome struct { + fs *FrameSocket + err error + } + returned := make(chan outcome, 1) + url := "ws" + server.URL[4:] + go func() { + fs, err := connectRace(ctx, logger, http.DefaultClient, []string{url, url}, nil, afterConnect) + returned <- outcome{fs, err} + }() + var winner *FrameSocket + select { + case result := <-returned: + if result.err != nil { + t.Fatal(result.err) + } + winner = result.fs + case <-time.After(time.Second): + releaseServers() + t.Fatal("winner waited for an unacknowledged loser close") + } + defer winner.Close(0) + select { + case <-logger.started: + case <-ctx.Done(): + t.Fatal("winner read pump did not start") + } + winner.Close(0) + select { + case <-logger.exited: + case <-time.After(time.Second): + t.Fatal("closing a socket with no frame consumer retained its pump") + } + releaseServers() + var explicitLosers int + for range 2 { + select { + case rec := <-closes: + if rec.code == websocket.StatusNormalClosure && rec.reason == LoserSocketCloseReason { + explicitLosers++ + } + case <-ctx.Done(): + t.Fatal("server did not observe socket cleanup") + } + } + if explicitLosers != 1 { + t.Fatalf("explicit loser closes = %d, want 1", explicitLosers) + } + select { + case <-logger.started: + t.Fatal("a loser read pump started") + default: + } +} + +func TestFrameDeliveryStopsOnCancellation(t *testing.T) { + for _, fragmented := range []bool{false, true} { + t.Run(fmt.Sprintf("fragmented=%v", fragmented), func(t *testing.T) { + fs := NewFrameSocket(waLog.Noop, nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + fs.cancelCtx = ctx + if fragmented { + fs.processData([]byte{0, 0, 2, 42}) + } + done := make(chan struct{}) + go func() { + if fragmented { + fs.processData([]byte{43}) + } else { + fs.processData([]byte{0, 0, 1, 42}) + } + close(done) + }() + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("frame delivery did not stop after cancellation") + } + }) + } +} + +func TestConnectRaceCancelsEveryPendingDial(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + started := make(chan struct{}, 2) + client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) { + started <- struct{}{} + <-r.Context().Done() + return nil, r.Context().Err() + })} + returned := make(chan error, 1) + go func() { + _, err := ConnectRace(ctx, waLog.Noop, client, []string{"ws://first.invalid", "ws://second.invalid"}, nil) + returned <- err + }() + for range 2 { + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("a pending dial did not start") + } + } + cancel() + select { + case err := <-returned: + if !errors.Is(err, context.Canceled) || !errors.Is(err, ErrDialFailed) { + t.Fatalf("canceled race error = %v", err) + } + case <-time.After(time.Second): + t.Fatal("pending race did not return on cancellation") + } +} diff --git a/socket_race_test.go b/socket_race_test.go new file mode 100644 index 000000000..2bc86501e --- /dev/null +++ b/socket_race_test.go @@ -0,0 +1,61 @@ +// Copyright (c) 2026 Rajeh Taher +// Licensed under the MIT License. See LICENSE-MIT for details. + +package whatsmeow + +import ( + "context" + "errors" + "net" + "net/http" + "net/http/httptest" + "syscall" + "testing" + "time" + + "github.com/polymorfa/hypermeow/socket" + waLog "github.com/polymorfa/hypermeow/util/log" +) + +func connectStatusError(t *testing.T, status int) error { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(status) + })) + defer server.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _, err := socket.ConnectRace(ctx, waLog.Noop, http.DefaultClient, []string{"ws" + server.URL[4:]}, nil) + var statusErr socket.ErrWithStatusCode + if !errors.As(err, &statusErr) || statusErr.StatusCode != status { + t.Fatalf("connect error = %v, want HTTP %d", err, status) + } + return err +} + +func TestConnectRetryClassificationPreservesHTTPStatusPrecedence(t *testing.T) { + forbidden := connectStatusError(t, http.StatusForbidden) + unavailable := connectStatusError(t, http.StatusServiceUnavailable) + network := &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED} + for _, tc := range []struct { + name string + err error + want bool + }{ + {"single rejection", forbidden, false}, + {"single server error", unavailable, true}, + {"single network error", network, true}, + {"rejection before network", errors.Join(forbidden, network), false}, + {"network before rejection", errors.Join(network, forbidden), false}, + {"rejection before server error", errors.Join(forbidden, unavailable), false}, + {"server error before rejection", errors.Join(unavailable, forbidden), true}, + {"all rejected", errors.Join(forbidden, forbidden), false}, + {"no error", nil, false}, + } { + t.Run(tc.name, func(t *testing.T) { + if got := isRetryableConnectError(tc.err); got != tc.want { + t.Fatalf("retryable = %v, want %v: %v", got, tc.want, tc.err) + } + }) + } +}