THRIFT-5257: Fix Go THeaderTransport endOfFrame handling

Client: go

In the current implementation, we only call endOfFrame when we hit EOF
when reading from the frameReader. The problem is in go stdlib the Read
call finished reading the remaining data from frameReader will not
return EOF, the next Read will. This caused us in most cases only call
endOfFrame at the beginning of the next frame, which could cause
troubles because we didn't read the beginning of the frame properly.

2 files changed, 77 insertions, 2 deletions

diff --git a/lib/go/thrift/header_transport.go b/lib/go/thrift/header_transport.go
index c622c0e4f..e2080344b 100644
--- a/lib/go/thrift/header_transport.go
+++ b/lib/go/thrift/header_transport.go
@@ -512,7 +512,11 @@ func (t *THeaderTransport) Read(p []byte) (read int, err error) {
 	}
 	if t.frameReader != nil {
 		read, err = t.frameReader.Read(p)
-		if err == io.EOF {
+		if err == nil && t.frameBuffer.Len() <= 0 {
+			// the last Read finished the frame, do endOfFrame
+			// handling here.
+			err = t.endOfFrame()
+		} else if err == io.EOF {
 			err = t.endOfFrame()
 			if err != nil {
 				return
diff --git a/lib/go/thrift/header_transport_test.go b/lib/go/thrift/header_transport_test.go
index cee1cadc3..320fb2a6f 100644
--- a/lib/go/thrift/header_transport_test.go
+++ b/lib/go/thrift/header_transport_test.go
@@ -143,7 +143,7 @@ func TestTHeaderTransportNoReadBeyondFrame(t *testing.T) {
 		return nil
 	}
 	f := func(content string) bool {
-		defer trans.Reset()
+		trans.Reset()
 		if len(content) == 0 {
 			return true
 		}
@@ -173,9 +173,80 @@ func TestTHeaderTransportNoReadBeyondFrame(t *testing.T) {
 				buf[:read],
 			)
 		}
+
+		// Check for endOfFrame handling
+		if !reader.needReadFrame() {
+			t.Error("Expected needReadFrame to be true after read the frame fully, got false")
+		}
 		return !t.Failed()
 	}
 	if err := quick.Check(f, nil); err != nil {
 		t.Error(err)
 	}
 }
+
+func TestTHeaderTransportEndOfFrameHandling(t *testing.T) {
+	trans := NewTMemoryBuffer()
+	writeContent := func(writer TTransport, content string) error {
+		if _, err := io.Copy(writer, strings.NewReader(content)); err != nil {
+			return err
+		}
+		if err := writer.Flush(context.Background()); err != nil {
+			return err
+		}
+		return nil
+	}
+
+	readFully := func(content string) bool {
+		trans.Reset()
+		if len(content) == 0 {
+			return true
+		}
+
+		reader := NewTHeaderTransport(trans)
+		writer := NewTHeaderTransport(trans)
+		// Write content
+		if err := writeContent(writer, content); err != nil {
+			t.Error(err)
+		}
+		buf := make([]byte, len(content))
+		_, err := reader.Read(buf)
+		if err != nil {
+			t.Error(err)
+		}
+		if !reader.needReadFrame() {
+			t.Error("Expected needReadFrame to be true after read the frame fully, got false")
+		}
+		return !t.Failed()
+	}
+	if err := quick.Check(readFully, nil); err != nil {
+		t.Error(err)
+	}
+
+	readPartially := func(content string) bool {
+		trans.Reset()
+		if len(content) < 1 {
+			return true
+		}
+
+		reader := NewTHeaderTransport(trans)
+		writer := NewTHeaderTransport(trans)
+		// Write content
+		if err := writeContent(writer, content); err != nil {
+			t.Error(err)
+		}
+		// Make the buf smaller so it can't read fully
+		buf := make([]byte, len(content)-1)
+		_, err := reader.Read(buf)
+		if err != nil {
+			t.Error(err)
+		}
+		if reader.needReadFrame() {
+			t.Error("Expected needReadFrame to be false before read the frame fully, got true")
+		}
+		return !t.Failed()
+	}
+	if err := quick.Check(readPartially, nil); err != nil {
+		t.Error(err)
+	}
+}