From e919ac149da3a9eb17b75cae441e21a372168082 Mon Sep 17 00:00:00 2001 From: ADAMAYAMEOGO Date: Wed, 30 Sep 2026 01:32:34 +0000 Subject: [PATCH] socket: add timeout to websocket dial The dial used the context passed to Connect, which normally has no deadline, and neither net/http nor coder/websocket limit how long to wait for the upgrade response. A proxy or middlebox that accepts the connection but never answers the upgrade request made Connect block forever while holding the socket lock, so Disconnect blocked too and the client could only be recovered by restarting the process. Limit the dial to DialTimeout (20 seconds by default, zero disables it). The timeout only covers the dial, the established connection still uses the parent context. A timed out dial returns ErrDialFailed, so it is retried like any other connection failure. --- socket/framesocket.go | 15 +++- socket/framesocket_test.go | 140 +++++++++++++++++++++++++++++++++++++ 2 files changed, 154 insertions(+), 1 deletion(-) create mode 100644 socket/framesocket_test.go diff --git a/socket/framesocket.go b/socket/framesocket.go index 4d6fc6e62..f762312b0 100644 --- a/socket/framesocket.go +++ b/socket/framesocket.go @@ -13,12 +13,18 @@ import ( "net/http" "sync" "sync/atomic" + "time" "github.com/coder/websocket" waLog "go.mau.fi/whatsmeow/util/log" ) +// DialTimeout limits how long the websocket dial (including the HTTP upgrade) may take. +// Without a limit, a proxy or server that accepts the connection but never answers +// the upgrade request would block Connect forever. Set to zero to disable. +var DialTimeout = 20 * time.Second + type FrameSocket struct { parentCtx context.Context cancelCtx context.Context @@ -102,7 +108,14 @@ func (fs *FrameSocket) Connect(ctx context.Context) error { fs.cancelCtx, fs.cancel = context.WithCancel(ctx) fs.log.Debugf("Dialing %s", fs.URL) - conn, resp, err := websocket.Dial(ctx, fs.URL, fs.makeDialOptions()) + // The timeout only applies to the dial, the connection itself lives on ctx. + dialCtx := ctx + if DialTimeout > 0 { + var cancelDial context.CancelFunc + dialCtx, cancelDial = context.WithTimeout(ctx, DialTimeout) + defer cancelDial() + } + conn, resp, err := websocket.Dial(dialCtx, fs.URL, fs.makeDialOptions()) if err != nil { if resp != nil { err = ErrWithStatusCode{err, resp.StatusCode} diff --git a/socket/framesocket_test.go b/socket/framesocket_test.go new file mode 100644 index 000000000..e27022853 --- /dev/null +++ b/socket/framesocket_test.go @@ -0,0 +1,140 @@ +// Copyright (c) 2026 Tulir Asokan +// +// 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" + "errors" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/coder/websocket" + + waLog "go.mau.fi/whatsmeow/util/log" +) + +// silentServer accepts TCP connections but never answers the websocket upgrade request. +func silentServer(t *testing.T) string { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + var mu sync.Mutex + var conns []net.Conn + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + mu.Lock() + conns = append(conns, conn) + mu.Unlock() + go func() { + buf := make([]byte, 4096) + for { + if _, err := conn.Read(buf); err != nil { + return + } + } + }() + } + }() + t.Cleanup(func() { + _ = ln.Close() + mu.Lock() + defer mu.Unlock() + for _, conn := range conns { + _ = conn.Close() + } + }) + return ln.Addr().String() +} + +func setDialTimeout(t *testing.T, timeout time.Duration) { + t.Helper() + old := DialTimeout + DialTimeout = timeout + t.Cleanup(func() { DialTimeout = old }) +} + +func TestConnectTimesOutWhenUpgradeIsNotAnswered(t *testing.T) { + setDialTimeout(t, 300*time.Millisecond) + addr := silentServer(t) + + fs := NewFrameSocket(waLog.Noop, &http.Client{}) + fs.URL = "ws://" + addr + "/ws/chat" + result := make(chan error, 1) + go func() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result <- fs.Connect(ctx) + }() + + select { + case err := <-result: + if !errors.Is(err, ErrDialFailed) { + t.Fatalf("expected ErrDialFailed, got %v", err) + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected the dial timeout to be the cause, got %v", err) + } + if fs.IsConnected() { + t.Fatal("socket is connected after a failed dial") + } + case <-time.After(5 * time.Second): + t.Fatal("Connect did not return after the dial timeout") + } +} + +func TestDialTimeoutDoesNotCloseEstablishedSocket(t *testing.T) { + setDialTimeout(t, 200*time.Millisecond) + wait := 3 * DialTimeout + + payload := []byte("still alive after the dial timeout") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{OriginPatterns: []string{"web.whatsapp.com"}}) + if err != nil { + return + } + defer conn.CloseNow() + time.Sleep(wait) + frame := append([]byte{byte(len(payload) >> 16), byte(len(payload) >> 8), byte(len(payload))}, payload...) + if err := conn.Write(r.Context(), websocket.MessageBinary, frame); err != nil { + return + } + _, _, _ = conn.Read(r.Context()) + })) + defer srv.Close() + + fs := NewFrameSocket(waLog.Noop, srv.Client()) + fs.URL = "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/chat" + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if err := fs.Connect(ctx); err != nil { + t.Fatalf("Connect: %v", err) + } + defer fs.Close(websocket.StatusNormalClosure) + + select { + case frame := <-fs.Frames: + if string(frame) != string(payload) { + t.Fatalf("unexpected frame %q", frame) + } + case <-time.After(5 * time.Second): + t.Fatal("no frame received after the dial timeout") + } + if !fs.IsConnected() { + t.Fatal("socket was closed by the dial timeout") + } +}