add timeout to proxy connection reading and writing

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,13 +64,25 @@ func FasthttpHTTPDialerTimeout(proxy string, timeout time.Duration) fasthttp.Dia
 			return nil, err
 		}
 
+		unsetDeadline := func(net.Conn) error { return nil }
+		if timeout > 0 {
+			if err = conn.SetDeadline(start.Add(timeout)); err != nil {
+				conn.Close()
+				return nil, err
+			}
+			unsetDeadline = func(conn net.Conn) error {
+				return conn.SetDeadline(time.Time{})
+			}
+		}
+
 		req := "CONNECT " + addr + " HTTP/1.1\r\nHost: " + addr + "\r\n"
 		if auth != "" {
 			req += "Proxy-Authorization: Basic " + auth + "\r\n"
 		}
 		req += "\r\n"
 
 		if _, err := conn.Write([]byte(req)); err != nil {
+			conn.Close()
 			return nil, err
 		}
 
@@ -86,6 +99,10 @@ 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 err := unsetDeadline(conn); err != nil {
+			conn.Close()
+			return nil, err
+		}
 		return conn, nil
 	}
 }

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,18 @@ func FasthttpProxyHTTPDialerTimeout(timeout time.Duration) fasthttp.DialFunc {
 			return nil, err
 		}
 
+		unsetDeadline := func(net.Conn) error { return nil }
+		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)
+				}
+			}
+			unsetDeadline = func(conn net.Conn) error {
+				return conn.SetDeadline(time.Time{})
+			}
+		}
+
 		req := "CONNECT " + addr + " HTTP/1.1\r\n"
 
 		if proxyURL.User != nil {
@@ -98,6 +111,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 +136,12 @@ 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 err := unsetDeadline(conn); 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
 	}
 }