THRIFT-5183: Don't try to read next frame in THeaderTransport.Read

lib/go/thrift/header_transport.go

@@ -499,10 +499,17 @@ func (t *THeaderTransport) Read(p []byte) (read int, err error) {
 			if err != nil {
 				return
 			}
-			if read < len(p) {
-				var nextRead int
-				nextRead, err = t.Read(p[read:])
-				read += nextRead
+			if read == 0 {
+				// Try to read the next frame when we hit EOF
+				// (end of frame) immediately.
+				// When we got here, it means the last read
+				// finished the previous frame, but didn't
+				// do endOfFrame handling yet.
+				// We have to read the next frame here,
+				// as otherwise we would return 0 and nil,
+				// which is a case not handled well by most
+				// protocol implementations.
+				return t.Read(p)
 			}
 		}
 		return

lib/go/thrift/header_transport_test.go

@@ -23,7 +23,9 @@ import (
 	"context"
 	"io"
 	"io/ioutil"
+	"strings"
 	"testing"
+	"testing/quick"
 )
 
 func TestTHeaderHeadersReadWrite(t *testing.T) {
@@ -128,3 +130,52 @@ func TestTHeaderTransportNoDoubleWrapping(t *testing.T) {
 		t.Errorf("NewTHeaderTransport double wrapped THeaderTransport")
 	}
 }
+
+func TestTHeaderTransportNoReadBeyondFrame(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
+	}
+	f := func(content string) bool {
+		defer trans.Reset()
+		if len(content) == 0 {
+			return true
+		}
+
+		reader := NewTHeaderTransport(trans)
+		writer := NewTHeaderTransport(trans)
+		// Write content twice
+		if err := writeContent(writer, content); err != nil {
+			t.Error(err)
+		}
+		if err := writeContent(writer, content); err != nil {
+			t.Error(err)
+		}
+		// buf is big enough to read both content out,
+		// but it shouldn't read beyond the first one in a single Read call.
+		buf := make([]byte, len(content)*3)
+		read, err := reader.Read(buf)
+		if err != nil {
+			t.Error(err)
+		}
+		if read == 0 || read > len(content) {
+			t.Errorf(
+				"Expected read in no more than %d:%q, got %d:%q",
+				len(content),
+				content,
+				read,
+				buf[:read],
+			)
+		}
+		return !t.Failed()
+	}
+	if err := quick.Check(f, nil); err != nil {
+		t.Error(err)
+	}
+}