Support exit node in ws client

This commit is contained in:
Zoltán Papp 2024-06-28 11:44:50 +02:00
parent 9ae03046e7
commit faeae52329

View File

@ -4,9 +4,12 @@ import (
"context" "context"
"fmt" "fmt"
"net" "net"
"net/http"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"nhooyr.io/websocket" "nhooyr.io/websocket"
nbnet "github.com/netbirdio/netbird/util/net"
) )
func Dial(address string) (net.Conn, error) { func Dial(address string) (net.Conn, error) {
@ -21,7 +24,8 @@ func Dial(address string) (net.Conn, error) {
url := fmt.Sprintf("ws://%s:%d", addr.IP.String(), addr.Port) url := fmt.Sprintf("ws://%s:%d", addr.IP.String(), addr.Port)
opts := &websocket.DialOptions{ opts := &websocket.DialOptions{
Host: hostName, Host: hostName,
HTTPClient: httpClientNbDialer(),
} }
wsConn, _, err := websocket.Dial(context.Background(), url, opts) wsConn, _, err := websocket.Dial(context.Background(), url, opts)
@ -34,3 +38,17 @@ func Dial(address string) (net.Conn, error) {
return conn, nil return conn, nil
} }
func httpClientNbDialer() *http.Client {
customDialer := nbnet.NewDialer()
customTransport := &http.Transport{
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return customDialer.DialContext(ctx, network, addr)
},
}
return &http.Client{
Transport: customTransport,
}
}