From e7ab6eb42d60837e263d4a16d288bd9569324984 Mon Sep 17 00:00:00 2001
From: Antoine Mercadal <amercadal@paloaltonetworks.com>
Date: Tue, 16 May 2023 15:14:49 -0700
Subject: [PATCH] new: add NetDialContext to config

This patch adds a new config to allow callers to pass a custom
net.Dialer. It also switch calling Dial for DialContext.
---
 config.go | 23 +++++++++++++----------
 wsc.go    |  3 ++-
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/config.go b/config.go
index 4539cdd..dbfbe77 100644
--- a/config.go
+++ b/config.go
@@ -12,21 +12,24 @@
 package wsc
 
 import (
+	"context"
 	"crypto/tls"
+	"net"
 	"net/http"
 	"time"
 )
 
 // Config contains configuration for the webbsocket.
 type Config struct {
-	WriteWait         time.Duration
-	PongWait          time.Duration
-	PingPeriod        time.Duration
-	TLSConfig         *tls.Config
-	ReadBufferSize    int
-	ReadChanSize      int
-	WriteBufferSize   int
-	WriteChanSize     int
-	EnableCompression bool
-	Headers           http.Header
+	TLSConfig          *tls.Config
+	Headers            http.Header
+	NetDialContextFunc func(ctx context.Context, network, addr string) (net.Conn, error)
+	WriteWait          time.Duration
+	PongWait           time.Duration
+	PingPeriod         time.Duration
+	ReadBufferSize     int
+	ReadChanSize       int
+	WriteBufferSize    int
+	WriteChanSize      int
+	EnableCompression  bool
 }
diff --git a/wsc.go b/wsc.go
index 70fe5a4..ab2141f 100644
--- a/wsc.go
+++ b/wsc.go
@@ -62,9 +62,10 @@ func Connect(ctx context.Context, url string, config Config) (Websocket, *http.R
 		ReadBufferSize:    config.ReadBufferSize,
 		WriteBufferSize:   config.WriteBufferSize,
 		EnableCompression: config.EnableCompression,
+		NetDialContext:    config.NetDialContextFunc,
 	}
 
-	conn, resp, err := dialer.Dial(url, config.Headers)
+	conn, resp, err := dialer.DialContext(ctx, url, config.Headers)
 	if err != nil {
 		return nil, resp, err
 	}