From 21b235d0330b37ecc4f80095751efcf67c6920cf Mon Sep 17 00:00:00 2001
From: Sniper91 <kevinzhao91@outlook.com>
Date: Wed, 19 Jun 2024 17:19:46 +0800
Subject: [PATCH] add timeout to proxy connection reading and writing (#1791)

Co-authored-by: kalmanzhao <kalmanzhao@tencent.com>
---
 fasthttpproxy/http.go      | 15 +++++++++++++++
 fasthttpproxy/proxy_env.go | 20 ++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/fasthttpproxy/http.go b/fasthttpproxy/http.go
index c09bb2c..a1a9e4c 100644
--- a/fasthttpproxy/http.go
+++ b/fasthttpproxy/http.go
@@ -42,6 +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
@@ -63,6 +64,13 @@ func FasthttpHTTPDialerTimeout(proxy string, timeout time.Duration) fasthttp.Dia
 			return nil, err
 		}
 
+		if timeout > 0 {
+			if err = conn.SetDeadline(start.Add(timeout)); err != nil {
+				conn.Close()
+				return nil, err
+			}
+		}
+
 		req := "CONNECT " + addr + " HTTP/1.1\r\nHost: " + addr + "\r\n"
 		if auth != "" {
 			req += "Proxy-Authorization: Basic " + auth + "\r\n"
@@ -70,6 +78,7 @@ func FasthttpHTTPDialerTimeout(proxy string, timeout time.Duration) fasthttp.Dia
 		req += "\r\n"
 
 		if _, err := conn.Write([]byte(req)); err != nil {
+			conn.Close()
 			return nil, err
 		}
 
@@ -86,6 +95,12 @@ func FasthttpHTTPDialerTimeout(proxy string, timeout time.Duration) fasthttp.Dia
 			conn.Close()
 			return nil, fmt.Errorf("could not connect to proxy: %s status code: %d", proxy, res.Header.StatusCode())
 		}
+		if timeout > 0 {
+			if err := conn.SetDeadline(time.Time{}); err != nil {
+				conn.Close()
+				return nil, err
+			}
+		}
 		return conn, nil
 	}
 }
diff --git a/fasthttpproxy/proxy_env.go b/fasthttpproxy/proxy_env.go
index 038f5cf..1448593 100644
--- 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,14 @@ func FasthttpProxyHTTPDialerTimeout(timeout time.Duration) fasthttp.DialFunc {
 			return nil, err
 		}
 
+		if timeout > 0 {
+			if err := conn.SetDeadline(start.Add(timeout)); err != nil {
+				if connErr := conn.Close(); connErr != nil {
+					return nil, fmt.Errorf("conn close err %v precede by set conn deadline %w", connErr, err)
+				}
+			}
+		}
+
 		req := "CONNECT " + addr + " HTTP/1.1\r\n"
 
 		if proxyURL.User != nil {
@@ -98,6 +107,9 @@ func FasthttpProxyHTTPDialerTimeout(timeout time.Duration) fasthttp.DialFunc {
 		req += "\r\n"
 
 		if _, err := conn.Write([]byte(req)); err != nil {
+			if connErr := conn.Close(); connErr != nil {
+				return nil, fmt.Errorf("conn close err %v precede by write conn err %w", connErr, err)
+			}
 			return nil, err
 		}
 
@@ -120,6 +132,14 @@ func FasthttpProxyHTTPDialerTimeout(timeout time.Duration) fasthttp.DialFunc {
 			}
 			return nil, fmt.Errorf("could not connect to proxy: code: %d body %q", res.StatusCode(), string(res.Body()))
 		}
+		if timeout > 0 {
+			if err := conn.SetDeadline(time.Time{}); err != nil {
+				if connErr := conn.Close(); connErr != nil {
+					return nil, fmt.Errorf("conn close err %v precede by clear conn deadline err %w", connErr, err)
+				}
+				return nil, err
+			}
+		}
 		return conn, nil
 	}
 }