From a931507dd67b1143e58d475414ca91191bea4cf0 Mon Sep 17 00:00:00 2001
From: Dmitry Anderson <58323684+4nd3r5on@users.noreply.github.com>
Date: Mon, 16 Sep 2024 14:33:03 +0200
Subject: [PATCH] SplitHTTP: Read and validate HTTP/1.1 responses (#3797)

---
 .gitignore                              |  3 +++
 transport/internet/splithttp/client.go  | 27 ++++++++++++++++++++-----
 transport/internet/splithttp/dialer.go  |  1 -
 transport/internet/splithttp/h1_conn.go | 19 +++++++++++++++++
 4 files changed, 44 insertions(+), 6 deletions(-)
 create mode 100644 transport/internet/splithttp/h1_conn.go

diff --git a/.gitignore b/.gitignore
index c77bc579..22469b88 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,3 +28,6 @@ errorgen
 *.dat
 .vscode
 /build_assets
+
+# Output from dlv test
+**/debug.*
diff --git a/transport/internet/splithttp/client.go b/transport/internet/splithttp/client.go
index e491ef3e..7fc26945 100644
--- a/transport/internet/splithttp/client.go
+++ b/transport/internet/splithttp/client.go
@@ -3,6 +3,7 @@ package splithttp
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
 	gonet "net"
 	"net/http"
@@ -152,23 +153,39 @@ func (c *DefaultDialerClient) SendUploadRequest(ctx context.Context, url string,
 		// safely retried. if instead req.Write is called multiple
 		// times, the body is already drained after the first
 		// request
-		requestBytes := new(bytes.Buffer)
-		common.Must(req.Write(requestBytes))
+		requestBuff := new(bytes.Buffer)
+		common.Must(req.Write(requestBuff))
 
 		var uploadConn any
+		var h1UploadConn *H1Conn
 
 		for {
 			uploadConn = c.uploadRawPool.Get()
 			newConnection := uploadConn == nil
 			if newConnection {
-				uploadConn, err = c.dialUploadConn(context.WithoutCancel(ctx))
+				newConn, err := c.dialUploadConn(context.WithoutCancel(ctx))
 				if err != nil {
 					return err
 				}
+				h1UploadConn = NewH1Conn(newConn)
+				uploadConn = h1UploadConn
+			} else {
+				h1UploadConn = uploadConn.(*H1Conn)
+
+				// TODO: Replace 0 here with a config value later
+				// Or add some other condition for optimization purposes
+				if h1UploadConn.UnreadedResponsesCount > 0 {
+					resp, err := http.ReadResponse(h1UploadConn.RespBufReader, req)
+					if err != nil {
+						return fmt.Errorf("error while reading response: %s", err.Error())
+					}
+					if resp.StatusCode != 200 {
+						return fmt.Errorf("got non-200 error response code: %d", resp.StatusCode)
+					}
+				}
 			}
 
-			_, err = uploadConn.(net.Conn).Write(requestBytes.Bytes())
-
+			_, err := h1UploadConn.Write(requestBuff.Bytes())
 			// if the write failed, we try another connection from
 			// the pool, until the write on a new connection fails.
 			// failed writes to a pooled connection are normal when
diff --git a/transport/internet/splithttp/dialer.go b/transport/internet/splithttp/dialer.go
index 69a30870..ff1501a8 100644
--- a/transport/internet/splithttp/dialer.go
+++ b/transport/internet/splithttp/dialer.go
@@ -267,7 +267,6 @@ func Dial(ctx context.Context, dest net.Destination, streamSettings *internet.Me
 					&buf.MultiBufferContainer{MultiBuffer: chunk},
 					int64(chunk.Len()),
 				)
-
 				if err != nil {
 					errors.LogInfoInner(ctx, err, "failed to send upload")
 					uploadPipeReader.Interrupt()
diff --git a/transport/internet/splithttp/h1_conn.go b/transport/internet/splithttp/h1_conn.go
new file mode 100644
index 00000000..f89f2a66
--- /dev/null
+++ b/transport/internet/splithttp/h1_conn.go
@@ -0,0 +1,19 @@
+package splithttp
+
+import (
+	"bufio"
+	"net"
+)
+
+type H1Conn struct {
+	UnreadedResponsesCount int
+	RespBufReader          *bufio.Reader
+	net.Conn
+}
+
+func NewH1Conn(conn net.Conn) *H1Conn {
+	return &H1Conn{
+		RespBufReader: bufio.NewReader(conn),
+		Conn:          conn,
+	}
+}