fix: timeout should close connection also when body reading started

Author: HeavyPunk

helper/clickhouse/clickhouse.go

@@ -484,10 +484,7 @@ func sendRequestWithProgressCheck(request *http.Request, opts *Options) (*http.R
 		DisableKeepAlives: true,
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), opts.Timeout)
-	defer cancel()
-
-	return httpHelper.DoHTTPOverTCP(ctx, transport, request)
+	return httpHelper.DoHTTPOverTCP(request.Context(), transport, request, opts.Timeout)
 }
 
 func do(ctx context.Context, dsn string, query string, postBody io.Reader, encoding ContentEncoding, opts Options, extData *ExternalData) ([]byte, int64, int64, error) {

helper/http/live-http-client.go

@@ -5,13 +5,19 @@ import (
 	"bytes"
 	"context"
 	"io"
+	"net"
 	"net/http"
+	"time"
 )
 
 const TCPNetwork string = "tcp"
 
-func DoHTTPOverTCP(ctx context.Context, transport *http.Transport, req *http.Request) (*http.Response, error) {
-	conn, err := transport.DialContext(context.Background(), TCPNetwork, req.URL.Host)
+func DoHTTPOverTCP(ctx context.Context, transport *http.Transport, req *http.Request, timeout time.Duration) (*http.Response, error) {
+	conn, err := transport.DialContext(ctx, TCPNetwork, req.URL.Host)
+	if err != nil {
+		return nil, err
+	}
+	err = conn.SetDeadline(time.Now().Add(timeout))
 	if err != nil {
 		return nil, err
 	}
@@ -24,25 +30,21 @@ func DoHTTPOverTCP(ctx context.Context, transport *http.Transport, req *http.Req
 	var backup_buf bytes.Buffer
 	reader := bufio.NewReader(io.TeeReader(conn, &backup_buf))
 
-	HEADERS:
 	for {
-		select {
-		case <- ctx.Done():
-			fake_body_delimer := bytes.NewBuffer([]byte{'\r', '\n', '\r', '\n'})
-			resp, err := http.ReadResponse(bufio.NewReader(io.MultiReader(&backup_buf, fake_body_delimer)), nil)
-			if err != nil {
-				return nil, err
-			}
-			return resp, ctx.Err()
-		
-		default:
-			line, err := reader.ReadString('\n')
-			if err != nil {
-				return nil, err
-			}
-			if line == "\r\n" {
-				break HEADERS
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+				fake_body_delimer := bytes.NewBuffer([]byte{'\r', '\n', '\r', '\n'})
+				resp, err := http.ReadResponse(bufio.NewReader(io.MultiReader(&backup_buf, fake_body_delimer)), nil)
+				if err != nil {
+					return nil, err
+				}
+				return resp, netErr
 			}
+			return nil, err
+		}
+		if line == "\r\n" {
+			break
 		}
 	}