extension/websocket: support dial with customized dialer instance

lesismal · Dec 31, 2023 · 03c172d · 03c172d
1 parent e32daca
commit 03c172d
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/extension/protocol/websocket/websocket.go b/extension/protocol/websocket/websocket.go
@@ -150,8 +150,15 @@ func Listen(addr string, upgrader *websocket.Upgrader) (net.Listener, error) {
 }
 
 // Dial wraps websocket dial
-func Dial(url string) (net.Conn, error) {
-	c, _, err := websocket.DefaultDialer.Dial(url, nil)
+func Dial(url string, args ...interface{}) (net.Conn, error) {
+	dialer := websocket.DefaultDialer
+	if len(args) > 0 {
+		d, ok := args[0].(*websocket.Dialer)
+		if ok {
+			dialer = d
+		}
+	}
+	c, _, err := dialer.Dial(url, nil)
 	if err != nil {
 		return nil, err
 	}