set deadline for proxy connection

valyala · Jun 12, 2024 · fb78ace · fb78ace
1 parent f56f9e4
commit fb78ace
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/fasthttpproxy/http.go b/fasthttpproxy/http.go
@@ -42,7 +42,7 @@ func FasthttpHTTPDialerTimeout(proxy string, timeout time.Duration) fasthttp.Dia
 	return func(addr string) (net.Conn, error) {
 		var conn net.Conn
 		var err error
-
+		start := time.Now()
 		if strings.HasPrefix(proxy, "[") {
 			// ipv6
 			if timeout == 0 {
@@ -63,6 +63,12 @@ func FasthttpHTTPDialerTimeout(proxy string, timeout time.Duration) fasthttp.Dia
 			return nil, err
 		}
 
+		if timeout > 0 {
+			if err = conn.SetDeadline(start.Add(timeout)); err != nil {
+				return nil, err
+			}
+		}
+
 		req := "CONNECT " + addr + " HTTP/1.1\r\nHost: " + addr + "\r\n"
 		if auth != "" {
 			req += "Proxy-Authorization: Basic " + auth + "\r\n"

diff --git a/fasthttpproxy/proxy_env.go b/fasthttpproxy/proxy_env.go
@@ -47,6 +47,7 @@ func FasthttpProxyHTTPDialerTimeout(timeout time.Duration) fasthttp.DialFunc {
 	authHTTPSStorage := &atomic.Value{}
 
 	return func(addr string) (net.Conn, error) {
+		start := time.Now()
 		port, _, err := net.SplitHostPort(addr)
 		if err != nil {
 			return nil, fmt.Errorf("unexpected addr format: %w", err)
@@ -78,6 +79,12 @@ func FasthttpProxyHTTPDialerTimeout(timeout time.Duration) fasthttp.DialFunc {
 			return nil, err
 		}
 
+		if timeout > 0 {
+			if err = conn.SetDeadline(start.Add(timeout)); err != nil {
+				return nil, err
+			}
+		}
+
 		req := "CONNECT " + addr + " HTTP/1.1\r\n"
 
 		if proxyURL.User != nil {