From c1cd3e85a792d621510de2a28a7e552228981845 Mon Sep 17 00:00:00 2001
From: Artem Fetishev <149964189+rtm0@users.noreply.github.com>
Date: Thu, 3 Oct 2024 18:27:15 +0200
Subject: [PATCH] lib/promscrape: Fix TestClientProxyReadOk flaky test (#7173)

This PR fixes #7062

For hijacked connections, one has to read from the connection buffer,
but still write directly to the connection. Otherwise, when reading
directly from such connections, the first byte may be lost. This, in
turn corrupts the ClientHello TLS handshake message and when the backend
server receives it, it closes the connection and reports the following
error in the log:

```
http: TLS handshake error from 127.0.0.1:33150: tls: first record does not look
like a TLS handshake
```

The first byte may be lost because underlying HTTP request handler may
read it from the connection and put it into the buffer. As the result,
subsequent connection reads won't see that byte.

-   See: https://github.com/golang/go/issues/27408
-   The fix is taken from : https://github.com/k3s-io/k3s/pull/6216

### Checklist

The following checks are **mandatory**:

- [x] My change adheres [VictoriaMetrics contributing
guidelines](https://docs.victoriametrics.com/contributing/).

Signed-off-by: Artem Fetishev <rtm@victoriametrics.com>
---
 lib/promscrape/client_test.go | 44 ++++++++++++++++++++++++-----------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/lib/promscrape/client_test.go b/lib/promscrape/client_test.go
index 8e53d9ad3..5fb942ee0 100644
--- a/lib/promscrape/client_test.go
+++ b/lib/promscrape/client_test.go
@@ -1,6 +1,21 @@
 package promscrape
 
-/*
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/bytesutil"
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/promauth"
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/proxy"
+)
+
 func copyHeader(dst, src http.Header) {
 	for k, vv := range src {
 		for _, v := range vv {
@@ -9,13 +24,18 @@ func copyHeader(dst, src http.Header) {
 	}
 }
 
+type connReadWriteCloser struct {
+	io.Reader
+	io.WriteCloser
+}
+
 func proxyTunnel(w http.ResponseWriter, r *http.Request) {
 	transfer := func(src io.ReadCloser, dst io.WriteCloser) {
 		defer dst.Close()
 		defer src.Close()
 		io.Copy(dst, src) //nolint
 	}
-	destConn, err := net.DialTimeout("tcp", r.Host, 10*time.Second)
+	server, err := net.DialTimeout("tcp", r.Host, 10*time.Second)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusServiceUnavailable)
 		return
@@ -26,12 +46,16 @@ func proxyTunnel(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
 		return
 	}
-	clientConn, _, err := hijacker.Hijack()
+	clientConn, clientBuf, err := hijacker.Hijack()
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusServiceUnavailable)
 	}
-	go transfer(clientConn, destConn)
-	transfer(destConn, clientConn)
+	// For hijacked connections, one has to read from the connection buffer, but
+	// still write directly to the connection.
+	client := &connReadWriteCloser{clientBuf, clientConn}
+
+	go transfer(client, server)
+	transfer(server, client)
 }
 
 type testProxyServer struct {
@@ -133,14 +157,7 @@ func TestClientProxyReadOk(t *testing.T) {
 		}
 
 		var bb bytesutil.ByteBuffer
-		err = c.ReadData(&bb)
-		if errors.Is(err, io.EOF) {
-			bb.Reset()
-			// EOF could occur in slow envs, like CI
-			err = c.ReadData(&bb)
-		}
-
-		if err != nil {
+		if err = c.ReadData(&bb); err != nil {
 			t.Fatalf("unexpected error at ReadData: %s", err)
 		}
 		got, err := io.ReadAll(bb.NewReader())
@@ -183,4 +200,3 @@ func TestClientProxyReadOk(t *testing.T) {
 	// backend tls and proxy auth
 	f(true, false, nil, &promauth.BasicAuthConfig{Username: "proxy-test", Password: promauth.NewSecret("1234")})
 }
-*/