From 3de0f5dfda9e615ef4ad04aec8cb2c6691e2977e Mon Sep 17 00:00:00 2001
From: Edoardo Spadolini <edoardo.spadolini@goteleport.com>
Date: Fri, 13 May 2022 14:40:59 +0200
Subject: [PATCH] Client timeout fixes (#12557)

* Respect timeout in `(directDial).DialTimeout`

* Add a Timeout to the api/client/webclient calls

* Respect timeout in (proxyDial).DialTimeout

* Applied suggestion

Co-authored-by: Alan Parra <alan.parra@goteleport.com>

Co-authored-by: Alan Parra <alan.parra@goteleport.com>
---
 api/client/webclient/webclient.go |  8 +++++++-
 lib/utils/proxy/proxy.go          | 11 +++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/api/client/webclient/webclient.go b/api/client/webclient/webclient.go
index 5c7e413a29d..bb33154001e 100644
--- a/api/client/webclient/webclient.go
+++ b/api/client/webclient/webclient.go
@@ -30,6 +30,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/gravitational/teleport/api/client/proxy"
 	"github.com/gravitational/teleport/api/constants"
@@ -61,6 +62,8 @@ type Config struct {
 	ExtraHeaders map[string]string
 	// IgnoreHTTPProxy disables support for HTTP proxying when true.
 	IgnoreHTTPProxy bool
+	// Timeout is a timeout for requests.
+	Timeout time.Duration
 }
 
 // CheckAndSetDefaults checks and sets defaults
@@ -72,7 +75,9 @@ func (c *Config) CheckAndSetDefaults() error {
 	if c.ProxyAddr == "" && os.Getenv(defaults.TunnelPublicAddrEnvar) == "" {
 		return trace.BadParameter(message, "missing parameter ProxyAddr")
 	}
-
+	if c.Timeout == 0 {
+		c.Timeout = defaults.DefaultDialTimeout
+	}
 	return nil
 }
 
@@ -94,6 +99,7 @@ func newWebClient(cfg *Config) (*http.Client, error) {
 	}
 	return &http.Client{
 		Transport: otelhttp.NewTransport(proxy.NewHTTPFallbackRoundTripper(&transport, cfg.Insecure)),
+		Timeout:   cfg.Timeout,
 	}, nil
 }
 
diff --git a/lib/utils/proxy/proxy.go b/lib/utils/proxy/proxy.go
index 12731056dcc..4cb9b940c72 100644
--- a/lib/utils/proxy/proxy.go
+++ b/lib/utils/proxy/proxy.go
@@ -127,7 +127,9 @@ func (d directDial) DialTimeout(network, address string, timeout time.Duration)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
-		tlsConn, err := tls.Dial("tcp", address, conf)
+		tlsConn, err := tls.DialWithDialer(&net.Dialer{
+			Timeout: timeout,
+		}, "tcp", address, conf)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
@@ -184,7 +186,12 @@ func (d proxyDial) DialTimeout(network, address string, timeout time.Duration) (
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
-		conn = tls.Client(conn, conf)
+		tlsConn := tls.Client(conn, conf)
+		if err = tlsConn.HandshakeContext(ctx); err != nil {
+			conn.Close()
+			return nil, trace.Wrap(err)
+		}
+		conn = tlsConn
 	}
 	return conn, nil
 }