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.
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 @@
}
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 @@
return nil
}
f := func(content string) bool {
- defer trans.Reset()
+ trans.Reset()
if len(content) == 0 {
return true
}
@@ -173,9 +173,80 @@
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)
+ }
+}