From ad86f963be26c1c78154303f8eec86864bc6f32b Mon Sep 17 00:00:00 2001
From: fulder <misad90@gmail.com>
Date: Thu, 18 Jun 2020 13:27:59 +0200
Subject: [PATCH] Add Proxy to WebsocketOptions

Signed-off-by: Michal Sadowski <misad90@gmail.com>
---
 websocket.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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,