Add support for SOCKS proxies and only proxying login websocket

client.go

@@ -19,7 +19,9 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"go.mau.fi/util/random"
+	"golang.org/x/net/proxy"
 
 	"go.mau.fi/whatsmeow/appstate"
 	waBinary "go.mau.fi/whatsmeow/binary"
@@ -155,8 +157,10 @@ type Client struct {
 	uniqueID  string
 	idCounter atomic.Uint64
 
-	proxy socket.Proxy
-	http  *http.Client
+	proxy          Proxy
+	socksProxy     proxy.Dialer
+	proxyOnlyLogin bool
+	http           *http.Client
 
 	// This field changes the client to act like a Messenger client instead of a WhatsApp one.
 	//
@@ -249,19 +253,35 @@ func NewClient(deviceStore *store.Device, log waLog.Logger) *Client {
 	return cli
 }
 
-// SetProxyAddress is a helper method that parses a URL string and calls SetProxy.
+// SetProxyAddress is a helper method that parses a URL string and calls SetProxy or SetSOCKSProxy based on the URL scheme.
 //
 // Returns an error if url.Parse fails to parse the given address.
 func (cli *Client) SetProxyAddress(addr string) error {
+	if addr == "" {
+		cli.SetProxy(nil)
+		return nil
+	}
 	parsed, err := url.Parse(addr)
 	if err != nil {
 		return err
 	}
-	cli.SetProxy(http.ProxyURL(parsed))
+	if parsed.Scheme == "http" || parsed.Scheme == "https" {
+		cli.SetProxy(http.ProxyURL(parsed))
+	} else if parsed.Scheme == "socks5" {
+		px, err := proxy.FromURL(parsed, proxy.Direct)
+		if err != nil {
+			return err
+		}
+		cli.SetSOCKSProxy(px)
+	} else {
+		return fmt.Errorf("unsupported proxy scheme %q", parsed.Scheme)
+	}
 	return nil
 }
 
-// SetProxy sets the proxy to use for WhatsApp web websocket connections and media uploads/downloads.
+type Proxy = func(*http.Request) (*url.URL, error)
+
+// SetProxy sets a HTTP proxy to use for WhatsApp web websocket connections and media uploads/downloads.
 //
 // Must be called before Connect() to take effect in the websocket connection.
 // If you want to change the proxy after connecting, you must call Disconnect() and then Connect() again manually.
@@ -281,9 +301,51 @@ func (cli *Client) SetProxyAddress(addr string) error {
 //			return mediaProxyURL, nil
 //		}
 //	})
-func (cli *Client) SetProxy(proxy socket.Proxy) {
+func (cli *Client) SetProxy(proxy Proxy) {
 	cli.proxy = proxy
-	cli.http.Transport.(*http.Transport).Proxy = proxy
+	cli.socksProxy = nil
+	transport := cli.http.Transport.(*http.Transport)
+	transport.Proxy = proxy
+	transport.Dial = nil
+	transport.DialContext = nil
+}
+
+type SetProxyOptions struct {
+	// If NoWebsocket is true, the proxy won't be used for the websocket
+	NoWebsocket bool
+	// If NoMedia is true, the proxy won't be used for media uploads/downloads
+	NoMedia bool
+}
+
+// SetSOCKSProxy sets a SOCKS5 proxy to use for WhatsApp web websocket connections and media uploads/downloads.
+//
+// Same details as SetProxy apply, but using a different proxy for the websocket and media is not currently supported.
+func (cli *Client) SetSOCKSProxy(px proxy.Dialer, opts ...SetProxyOptions) {
+	var opt SetProxyOptions
+	if len(opts) > 0 {
+		opt = opts[0]
+	}
+	if !opt.NoWebsocket {
+		cli.socksProxy = px
+		cli.proxy = nil
+	}
+	if !opt.NoMedia {
+		transport := cli.http.Transport.(*http.Transport)
+		transport.Proxy = nil
+		transport.Dial = cli.socksProxy.Dial
+		contextDialer, ok := cli.socksProxy.(proxy.ContextDialer)
+		if ok {
+			transport.DialContext = contextDialer.DialContext
+		} else {
+			transport.DialContext = nil
+		}
+	}
+}
+
+// ToggleProxyOnlyForLogin changes whether the proxy set with SetProxy or related methods
+// is only used for the pre-login websocket and not authenticated websockets.
+func (cli *Client) ToggleProxyOnlyForLogin(only bool) {
+	cli.proxyOnlyLogin = only
 }
 
 func (cli *Client) getSocketWaitChan() <-chan struct{} {
@@ -339,7 +401,19 @@ func (cli *Client) Connect() error {
 	}
 
 	cli.resetExpectedDisconnect()
-	fs := socket.NewFrameSocket(cli.Log.Sub("Socket"), cli.proxy)
+	wsDialer := websocket.Dialer{}
+	if !cli.proxyOnlyLogin || cli.Store.ID == nil {
+		if cli.proxy != nil {
+			wsDialer.Proxy = cli.proxy
+		} else if cli.socksProxy != nil {
+			wsDialer.NetDial = cli.socksProxy.Dial
+			contextDialer, ok := cli.socksProxy.(proxy.ContextDialer)
+			if ok {
+				wsDialer.NetDialContext = contextDialer.DialContext
+			}
+		}
+	}
+	fs := socket.NewFrameSocket(cli.Log.Sub("Socket"), wsDialer)
 	if cli.MessengerConfig != nil {
 		fs.URL = "wss://web-chat-e2ee.facebook.com/ws/chat"
 		fs.HTTPHeaders.Set("Origin", cli.MessengerConfig.BaseURL)

go.mod

@@ -8,13 +8,14 @@ require (
 	github.com/rs/zerolog v1.32.0
 	go.mau.fi/libsignal v0.1.0
 	go.mau.fi/util v0.4.1
-	golang.org/x/crypto v0.21.0
+	golang.org/x/crypto v0.23.0
+	golang.org/x/net v0.25.0
 	google.golang.org/protobuf v1.33.0
 )
 
 require (
 	filippo.io/edwards25519 v1.0.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
-	golang.org/x/sys v0.18.0 // indirect
+	golang.org/x/sys v0.20.0 // indirect
 )

go.sum

@@ -27,13 +27,15 @@ go.mau.fi/libsignal v0.1.0 h1:vAKI/nJ5tMhdzke4cTK1fb0idJzz1JuEIpmjprueC+c=
 go.mau.fi/libsignal v0.1.0/go.mod h1:R8ovrTezxtUNzCQE5PH30StOQWWeBskBsWE55vMfY9I=
 go.mau.fi/util v0.4.1 h1:3EC9KxIXo5+h869zDGf5OOZklRd/FjeVnimTwtm3owg=
 go.mau.fi/util v0.4.1/go.mod h1:GjkTEBsehYZbSh2LlE6cWEn+6ZIZTGrTMM/5DMNlmFY=
-golang.org/x/crypto v0.21.0 h1:X31++rzVUdKhX5sWmSOFZxx8UW/ldWx55cbf08iNAMA=
-golang.org/x/crypto v0.21.0/go.mod h1:0BP7YvVV9gBbVKyeTG0Gyn+gZm94bibOW5BjDEYAOMs=
+golang.org/x/crypto v0.23.0 h1:dIJU/v2J8Mdglj/8rJ6UUOM3Zc9zLZxVZwwxMooUSAI=
+golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
+golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac=
+golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
-golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
+golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=

mdtest/go.mod

@@ -20,8 +20,9 @@ require (
 	github.com/rs/zerolog v1.32.0 // indirect
 	go.mau.fi/libsignal v0.1.0 // indirect
 	go.mau.fi/util v0.4.1 // indirect
-	golang.org/x/crypto v0.21.0 // indirect
-	golang.org/x/sys v0.18.0 // indirect
+	golang.org/x/crypto v0.23.0 // indirect
+	golang.org/x/net v0.25.0 // indirect
+	golang.org/x/sys v0.20.0 // indirect
 	rsc.io/qr v0.2.0 // indirect
 )
 

mdtest/go.sum

@@ -34,14 +34,16 @@ go.mau.fi/libsignal v0.1.0 h1:vAKI/nJ5tMhdzke4cTK1fb0idJzz1JuEIpmjprueC+c=
 go.mau.fi/libsignal v0.1.0/go.mod h1:R8ovrTezxtUNzCQE5PH30StOQWWeBskBsWE55vMfY9I=
 go.mau.fi/util v0.4.1 h1:3EC9KxIXo5+h869zDGf5OOZklRd/FjeVnimTwtm3owg=
 go.mau.fi/util v0.4.1/go.mod h1:GjkTEBsehYZbSh2LlE6cWEn+6ZIZTGrTMM/5DMNlmFY=
-golang.org/x/crypto v0.21.0 h1:X31++rzVUdKhX5sWmSOFZxx8UW/ldWx55cbf08iNAMA=
-golang.org/x/crypto v0.21.0/go.mod h1:0BP7YvVV9gBbVKyeTG0Gyn+gZm94bibOW5BjDEYAOMs=
+golang.org/x/crypto v0.23.0 h1:dIJU/v2J8Mdglj/8rJ6UUOM3Zc9zLZxVZwwxMooUSAI=
+golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
+golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac=
+golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
-golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
+golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=

socket/framesocket.go

@@ -11,7 +11,6 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
-	"net/url"
 	"sync"
 	"time"
 
@@ -20,8 +19,6 @@ import (
 	waLog "go.mau.fi/whatsmeow/util/log"
 )
 
-type Proxy = func(*http.Request) (*url.URL, error)
-
 type FrameSocket struct {
 	conn   *websocket.Conn
 	ctx    context.Context
@@ -37,15 +34,15 @@ type FrameSocket struct {
 	WriteTimeout time.Duration
 
 	Header []byte
-	Proxy  Proxy
+	Dialer websocket.Dialer
 
 	incomingLength int
 	receivedLength int
 	incoming       []byte
 	partialHeader  []byte
 }
 
-func NewFrameSocket(log waLog.Logger, proxy Proxy) *FrameSocket {
+func NewFrameSocket(log waLog.Logger, dialer websocket.Dialer) *FrameSocket {
 	return &FrameSocket{
 		conn:   nil,
 		log:    log,
@@ -55,7 +52,7 @@ func NewFrameSocket(log waLog.Logger, proxy Proxy) *FrameSocket {
 		URL:         URL,
 		HTTPHeaders: http.Header{"Origin": {Origin}},
 
-		Proxy: proxy,
+		Dialer: dialer,
 	}
 }
 
@@ -104,12 +101,9 @@ func (fs *FrameSocket) Connect() error {
 		return ErrSocketAlreadyOpen
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	dialer := websocket.Dialer{
-		Proxy: fs.Proxy,
-	}
 
 	fs.log.Debugf("Dialing %s", fs.URL)
-	conn, _, err := dialer.Dial(fs.URL, fs.HTTPHeaders)
+	conn, _, err := fs.Dialer.Dial(fs.URL, fs.HTTPHeaders)
 	if err != nil {
 		cancel()
 		return fmt.Errorf("couldn't dial whatsapp web websocket: %w", err)