diff --git a/websocket.go b/websocket.go index 5370d2f0..dd2aaf2b 100644 --- a/websocket.go +++ b/websocket.go @@ -5,6 +5,7 @@ import ( "io" "net" "net/http" + "net/url" "sync" "time" @@ -15,8 +16,11 @@ import ( type WebsocketOptions struct { ReadBufferSize int WriteBufferSize int + Proxy ProxyFunction } +type ProxyFunction func(req *http.Request) (*url.URL, error) + // NewWebsocket returns a new websocket and returns a net.Conn compatible interface using the gorilla/websocket package func NewWebsocket(host string, tlsc *tls.Config, timeout time.Duration, requestHeader http.Header, options *WebsocketOptions) (net.Conn, error) { if timeout == 0 { @@ -27,9 +31,11 @@ func NewWebsocket(host string, tlsc *tls.Config, timeout time.Duration, requestH // Apply default options options = &WebsocketOptions{} } - + if options.Proxy == nil { + options.Proxy = http.ProxyFromEnvironment + } dialer := &websocket.Dialer{ - Proxy: http.ProxyFromEnvironment, + Proxy: options.Proxy, HandshakeTimeout: timeout, EnableCompression: false, TLSClientConfig: tlsc,