Adds support for Proxy function specified by the user

ronakg · Mar 6, 2021 · 5133e49 · 5133e49
1 parent 318cdc6
commit 5133e49
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/core/transport/websocket_transport.go b/core/transport/websocket_transport.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/url"
 	"sync"
 	"time"
 
@@ -177,16 +178,12 @@ func NewWebsocketServerTransportWithAddr(addr string, path string, upgrader *web
 }
 
 // NewWebsocketClientTransport creates a new client-side transport.
-func NewWebsocketClientTransport(ctx context.Context, url string, config *tls.Config, header http.Header) (*Transport, error) {
+func NewWebsocketClientTransport(ctx context.Context, url string, config *tls.Config, header http.Header, proxy func(*http.Request) (*url.URL, error)) (*Transport, error) {
 	var dial *websocket.Dialer
-	if config == nil {
-		dial = websocket.DefaultDialer
-	} else {
-		dial = &websocket.Dialer{
-			Proxy:            http.ProxyFromEnvironment,
-			HandshakeTimeout: 45 * time.Second,
-			TLSClientConfig:  config,
-		}
+	dial = &websocket.Dialer{
+		Proxy:            proxy,
+		HandshakeTimeout: 45 * time.Second,
+		TLSClientConfig:  config,
 	}
 	conn, _, err := dial.DialContext(ctx, url, header)
 	if err != nil {

diff --git a/transporter_builder.go b/transporter_builder.go
@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"fmt"
 	"net/http"
+	"net/url"
 	"os"
 
 	"github.com/gorilla/websocket"
@@ -35,6 +36,7 @@ type WebsocketClientBuilder struct {
 	url    string
 	tlsCfg *tls.Config
 	header http.Header
+	proxy  func(*http.Request) (*url.URL, error)
 }
 
 // WebsocketServerBuilder provides builder which can be used to create a server-side Websocket transport easily.
@@ -173,10 +175,16 @@ func (wc *WebsocketClientBuilder) SetHeader(header http.Header) *WebsocketClient
 	return wc
 }
 
+// SetProxy sets proxy.
+func (wc *WebsocketClientBuilder) SetProxy(proxy func(*http.Request) (*url.URL, error)) *WebsocketClientBuilder {
+	wc.proxy = proxy
+	return wc
+}
+
 // Build builds and returns a new websocket ClientTransporter
 func (wc *WebsocketClientBuilder) Build() transport.ClientTransporter {
 	return func(ctx context.Context) (*transport.Transport, error) {
-		return transport.NewWebsocketClientTransport(ctx, wc.url, wc.tlsCfg, wc.header)
+		return transport.NewWebsocketClientTransport(ctx, wc.url, wc.tlsCfg, wc.header, wc.proxy)
 	}
 }
 
@@ -276,7 +284,8 @@ func TCPServer() *TCPServerBuilder {
 // WebsocketClient creates a new WebsocketClientBuilder.
 func WebsocketClient() *WebsocketClientBuilder {
 	return &WebsocketClientBuilder{
-		url: fmt.Sprintf("ws://127.0.0.1:%d", DefaultPort),
+		url:   fmt.Sprintf("ws://127.0.0.1:%d", DefaultPort),
+		proxy: http.ProxyFromEnvironment,
 	}
 }