thrift: refuse a length the reader cannot satisfy

thrift/binary.go

@@ -108,9 +108,24 @@ func (r *binaryReader) ReadLength() (int, error) {
 	if n > math.MaxInt32 {
 		return 0, fmt.Errorf("length out of range: %d", n)
 	}
+	if err := r.checkLength(int(n)); err != nil {
+		return 0, err
+	}
 	return int(n), nil
 }
 
+// checkLength rejects a length that the underlying reader cannot possibly
+// satisfy. The length is read from the wire, and the range check above still
+// admits values up to 2GiB, so without this a four byte header can make the
+// caller allocate gigabytes before io.ReadFull discovers there is no data
+// behind it. Readers that cannot report their remaining size are left alone.
+func (r *binaryReader) checkLength(n int) error {
+	if lr, ok := r.r.(interface{ Len() int }); ok && n > lr.Len() {
+		return fmt.Errorf("length %d exceeds the %d bytes remaining", n, lr.Len())
+	}
+	return nil
+}
+
 func (r *binaryReader) ReadMessage() (Message, error) {
 	m := Message{}
 
@@ -121,6 +136,9 @@ func (r *binaryReader) ReadMessage() (Message, error) {
 
 	if (b[0] >> 7) == 0 { // non-strict
 		n := int(binary.BigEndian.Uint32(b))
+		if err := r.checkLength(n); err != nil {
+			return m, err
+		}
 		s := make([]byte, n)
 		_, err := io.ReadFull(r.r, s)
 		if err != nil {

thrift/protocol_test.go

@@ -2,7 +2,9 @@ package thrift_test
 
 import (
 	"bytes"
+	"encoding/binary"
 	"reflect"
+	"runtime"
 	"strings"
 	"testing"
 
@@ -202,3 +204,63 @@ func testProtocolReadWriteValues(t *testing.T, p thrift.Protocol) {
 		})
 	}
 }
+
+// TestBinaryLengthBounds checks that a length larger than the data behind it does
+// not size an allocation before that data has arrived, while a length the data does
+// satisfy still round-trips. The assertion is on allocation volume, not merely on
+// getting an error: an overlong length errors either way once io.ReadFull runs out.
+func TestBinaryLengthBounds(t *testing.T) {
+	p := &thrift.BinaryProtocol{}
+
+	measure := func(f func() error) (uint64, error) {
+		var before, after runtime.MemStats
+		runtime.GC()
+		runtime.ReadMemStats(&before)
+		err := f()
+		runtime.ReadMemStats(&after)
+		return after.TotalAlloc - before.TotalAlloc, err
+	}
+
+	// A non-strict message header claiming a 64MiB name, with no name behind it.
+	hostile := make([]byte, 4)
+	binary.BigEndian.PutUint32(hostile, 1<<26)
+	hostile[0] &= 0x7f // clear the strict bit
+	alloc, err := measure(func() error {
+		_, err := p.NewReader(bytes.NewReader(hostile)).ReadMessage()
+		return err
+	})
+	if err == nil {
+		t.Fatal("expected an error for a name length with no data behind it")
+	}
+	if alloc > 1<<20 {
+		t.Fatalf("reading a message that claims a 64MiB name allocated %d bytes", alloc)
+	}
+
+	// ReadBytes goes through the same check.
+	hostileBytes := make([]byte, 4)
+	binary.BigEndian.PutUint32(hostileBytes, 1<<26)
+	alloc, err = measure(func() error {
+		_, err := p.NewReader(bytes.NewReader(hostileBytes)).ReadBytes()
+		return err
+	})
+	if err == nil {
+		t.Fatal("expected an error for a byte length with no data behind it")
+	}
+	if alloc > 1<<20 {
+		t.Fatalf("ReadBytes with a 64MiB length allocated %d bytes", alloc)
+	}
+
+	// A well formed message must still parse.
+	var good bytes.Buffer
+	w := p.NewWriter(&good)
+	if err := w.WriteMessage(thrift.Message{Type: thrift.Call, Name: "ping", SeqID: 1}); err != nil {
+		t.Fatal(err)
+	}
+	m, err := p.NewReader(bytes.NewReader(good.Bytes())).ReadMessage()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if m.Name != "ping" {
+		t.Fatalf("round trip changed the name: %q", m.Name)
+	}
+}