THRIFT-5324: reset http client buffer after flush

THttpClient did not reset its internal buffer when HTTP client returned
an error, leaving the whole or partially read message in the buffer.
Now we reset the buffer in defer.

Client: go

2 files changed, 63 insertions, 5 deletions

diff --git a/lib/go/thrift/http_client.go b/lib/go/thrift/http_client.go
index 1924a1ae2..19c63a985 100644
--- a/lib/go/thrift/http_client.go
+++ b/lib/go/thrift/http_client.go
@@ -197,6 +197,14 @@ func (p *THttpClient) Flush(ctx context.Context) error {
 	// Close any previous response body to avoid leaking connections.
 	p.closeResponse()
 
+	// Request might not have been fully read by http client.
+	// Reset so we don't send the remains on next call.
+	defer func() {
+		if p.requestBuffer != nil {
+			p.requestBuffer.Reset()
+		}
+	}()
+
 	req, err := http.NewRequest("POST", p.url.String(), p.requestBuffer)
 	if err != nil {
 		return NewTTransportExceptionFromError(err)
diff --git a/lib/go/thrift/http_client_test.go b/lib/go/thrift/http_client_test.go
index a7977a385..eba366815 100644
--- a/lib/go/thrift/http_client_test.go
+++ b/lib/go/thrift/http_client_test.go
@@ -20,6 +20,8 @@
 package thrift
 
 import (
+	"bytes"
+	"context"
 	"net/http"
 	"testing"
 )
@@ -32,14 +34,14 @@ func TestHttpClient(t *testing.T) {
 	trans, err := NewTHttpPostClient("http://" + addr.String())
 	if err != nil {
 		l.Close()
-		t.Fatalf("Unable to connect to %s: %s", addr.String(), err)
+		t.Fatalf("Unable to connect to %s: %v", addr.String(), err)
 	}
 	TransportTest(t, trans, trans)
 
 	t.Run("nilBuffer", func(t *testing.T) {
 		_ = trans.Close()
 		if _, err = trans.Write([]byte{1, 2, 3, 4}); err == nil {
-			t.Fatalf("writing to a closed transport did not result in an error")
+			t.Fatal("writing to a closed transport did not result in an error")
 		}
 	})
 }
@@ -52,7 +54,7 @@ func TestHttpClientHeaders(t *testing.T) {
 	trans, err := NewTHttpPostClient("http://" + addr.String())
 	if err != nil {
 		l.Close()
-		t.Fatalf("Unable to connect to %s: %s", addr.String(), err)
+		t.Fatalf("Unable to connect to %s: %v", addr.String(), err)
 	}
 	TransportHeaderTest(t, trans, trans)
 }
@@ -72,7 +74,7 @@ func TestHttpCustomClient(t *testing.T) {
 	})
 	if err != nil {
 		l.Close()
-		t.Fatalf("Unable to connect to %s: %s", addr.String(), err)
+		t.Fatalf("Unable to connect to %s: %v", addr.String(), err)
 	}
 	TransportHeaderTest(t, trans, trans)
 
@@ -94,7 +96,7 @@ func TestHttpCustomClientPackageScope(t *testing.T) {
 	trans, err := NewTHttpPostClient("http://" + addr.String())
 	if err != nil {
 		l.Close()
-		t.Fatalf("Unable to connect to %s: %s", addr.String(), err)
+		t.Fatalf("Unable to connect to %s: %v", addr.String(), err)
 	}
 	TransportHeaderTest(t, trans, trans)
 
@@ -103,6 +105,54 @@ func TestHttpCustomClientPackageScope(t *testing.T) {
 	}
 }
 
+func TestHTTPClientFlushesRequestBufferOnErrors(t *testing.T) {
+	var (
+		write1 = []byte("write 1")
+		write2 = []byte("write 2")
+	)
+
+	l, addr := HttpClientSetupForTest(t)
+	if l != nil {
+		defer l.Close()
+	}
+	trans, err := NewTHttpPostClient("http://" + addr.String())
+	if err != nil {
+		t.Fatalf("Unable to connect to %s: %v", addr.String(), err)
+	}
+	defer trans.Close()
+
+	_, err = trans.Write(write1)
+	if err != nil {
+		t.Fatalf("Failed to write to transport: %v", err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = trans.Flush(ctx)
+	if err == nil {
+		t.Fatal("Expected flush error")
+	}
+
+	_, err = trans.Write(write2)
+	if err != nil {
+		t.Fatalf("Failed to write to transport: %v", err)
+	}
+	err = trans.Flush(context.Background())
+	if err != nil {
+		t.Fatalf("Failed to flush: %v", err)
+	}
+
+	data := make([]byte, 1024)
+	n, err := trans.Read(data)
+	if err != nil {
+		t.Fatalf("Failed to read: %v", err)
+	}
+
+	data = data[:n]
+	if !bytes.Equal(data, write2) {
+		t.Fatalf("Received unexpected data: %q, expected: %q", data, write2)
+	}
+}
+
 type customHttpTransport struct {
 	hit bool
 }