http client: add Dial configuration

lesismal · Mar 22, 2024 · 1d12ed0 · 1d12ed0
1 parent 88d61a5
commit 1d12ed0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/nbhttp/client.go b/nbhttp/client.go
@@ -74,6 +74,7 @@ func (hcs *hostConns) getConn() (*hostConns, *ClientConn, error) {
 				Timeout:         c.Timeout,
 				IdleConnTimeout: c.IdleConnTimeout,
 				TLSClientConfig: c.TLSClientConfig,
+				Dial:            c.Dial,
 				Proxy:           c.Proxy,
 				CheckRedirect:   c.CheckRedirect,
 			}
@@ -122,6 +123,8 @@ type Client struct {
 
 	TLSClientConfig *tls.Config
 
+	Dial func(network, addr string) (net.Conn, error)
+
 	Proxy func(*http.Request) (*url.URL, error)
 
 	CheckRedirect func(req *http.Request, via []*http.Request) error

diff --git a/nbhttp/client_conn.go b/nbhttp/client_conn.go
@@ -43,6 +43,8 @@ type ClientConn struct {
 
 	TLSClientConfig *tls.Config
 
+	Dial func(network, addr string) (net.Conn, error)
+
 	Proxy func(*http.Request) (*url.URL, error)
 
 	CheckRedirect func(req *http.Request, via []*http.Request) error
@@ -213,14 +215,23 @@ func (c *ClientConn) Do(req *http.Request, handler func(res *http.Response, conn
 		}
 		addr := host + ":" + port
 
+		var dialer = c.Dial
 		var netDial netDialerFunc
 		if confTimeout <= 0 {
+			if dialer == nil {
+				dialer = net.Dial
+			}
 			netDial = func(network, addr string) (net.Conn, error) {
-				return net.Dial(network, addr)
+				return dialer(network, addr)
 			}
 		} else {
+			if dialer == nil {
+				dialer = func(network, addr string) (net.Conn, error) {
+					return net.DialTimeout(network, addr, timeout)
+				}
+			}
 			netDial = func(network, addr string) (net.Conn, error) {
-				conn, err := net.DialTimeout(network, addr, timeout)
+				conn, err := dialer(network, addr)
 				if err == nil {
 					conn.SetReadDeadline(deadline)
 				}