diff --git a/thrift/binary.go b/thrift/binary.go index 18d95d9a..962d8c65 100644 --- a/thrift/binary.go +++ b/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 { diff --git a/thrift/protocol_test.go b/thrift/protocol_test.go index 8aac085c..da2c544a 100644 --- a/thrift/protocol_test.go +++ b/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) + } +}