[thrift] branch master updated: THRIFT-5257: Fix Go THeaderTransport endOfFrame handling

[yu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

yuxuan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/thrift.git


The following commit(s) were added to refs/heads/master by this push:
     new 64f419b  THRIFT-5257: Fix Go THeaderTransport endOfFrame handling
64f419b is described below

commit 64f419b5ad40df233d34cc3715c68b8d85712699
Author: Yuxuan 'fishy' Wang <yu...@reddit.com>
AuthorDate: Wed Jul 29 10:12:53 2020 -0700

    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.
---
 lib/go/thrift/header_transport.go      |  6 ++-
 lib/go/thrift/header_transport_test.go | 73 +++++++++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/lib/go/thrift/header_transport.go b/lib/go/thrift/header_transport.go
index c622c0e..e208034 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 cee1cad..320fb2a 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)
+	}
+}