Pass explicit timeout for custom dialers

emersion · Apr 25, 2024 · 68dd0ca · 68dd0ca
1 parent 8b2892f
commit 68dd0ca
Showing 1 changed file with 22 additions and 12 deletions.
diff --git a/client/client.go b/client/client.go
@@ -629,6 +629,10 @@ type Dialer interface {
 //
 // Among other uses, this allows to apply a dial timeout.
 func DialWithDialer(dialer Dialer, addr string) (*Client, error) {
+	return DialWithDialerAndTimeout(dialer, addr, 0)
+}
+
+func DialWithDialerAndTimeout(dialer Dialer, addr string, timeout time.Duration) (*Client, error) {
 	conn, err := dialer.Dial("tcp", addr)
 	if err != nil {
 		return nil, err
@@ -638,13 +642,14 @@ func DialWithDialer(dialer Dialer, addr string) (*Client, error) {
 	// there is no way to set the client's Timeout for that action. As a
 	// workaround, if the dialer has a timeout set, use that for the connection's
 	// deadline.
-	var timeout time.Duration
-	if netDialer, ok := dialer.(*net.Dialer); ok && netDialer.Timeout > 0 {
-		err := conn.SetDeadline(time.Now().Add(netDialer.Timeout))
-		if err != nil {
-			return nil, err
+	if timeout == 0 {
+		if netDialer, ok := dialer.(*net.Dialer); ok && netDialer.Timeout > 0 {
+			timeout = netDialer.Timeout
 		}
-		timeout = netDialer.Timeout
+	}
+	err = conn.SetDeadline(time.Now().Add(timeout))
+	if err != nil {
+		return nil, err
 	}
 
 	c, err := NewWithTimeout(conn, timeout)
@@ -666,6 +671,10 @@ func DialTLS(addr string, tlsConfig *tls.Config) (*Client, error) {
 //
 // Among other uses, this allows to apply a dial timeout.
 func DialWithDialerTLS(dialer Dialer, addr string, tlsConfig *tls.Config) (*Client, error) {
+	return DialWithDialerAndTimeoutTLS(dialer, addr, tlsConfig, 0)
+}
+
+func DialWithDialerAndTimeoutTLS(dialer Dialer, addr string, tlsConfig *tls.Config, timeout time.Duration) (*Client, error) {
 	conn, err := dialer.Dial("tcp", addr)
 	if err != nil {
 		return nil, err
@@ -685,13 +694,14 @@ func DialWithDialerTLS(dialer Dialer, addr string, tlsConfig *tls.Config) (*Clie
 	// there is no way to set the client's Timeout for that action. As a
 	// workaround, if the dialer has a timeout set, use that for the connection's
 	// deadline.
-	var timeout time.Duration
-	if netDialer, ok := dialer.(*net.Dialer); ok && netDialer.Timeout > 0 {
-		err := tlsConn.SetDeadline(time.Now().Add(netDialer.Timeout))
-		if err != nil {
-			return nil, err
+	if timeout == 0 {
+		if netDialer, ok := dialer.(*net.Dialer); ok && netDialer.Timeout > 0 {
+			timeout = netDialer.Timeout
 		}
-		timeout = netDialer.Timeout
+	}
+	err = conn.SetDeadline(time.Now().Add(timeout))
+	if err != nil {
+		return nil, err
 	}
 
 	c, err := NewWithTimeout(tlsConn, timeout)