fix: avoid allocating declared lengths before reading data

core/declared_length_test.go

@@ -0,0 +1,148 @@
+package core
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/hdt3213/rdb/model"
+)
+
+// rdbWith builds a tiny RDB file with one object of the given type code and key "k".
+func rdbWith(typeCode byte, body ...[]byte) []byte {
+	buf := []byte("REDIS0009")
+	buf = append(buf, typeCode, 0x01, 'k')
+	for _, b := range body {
+		buf = append(buf, b...)
+	}
+	return append(buf, 0xff)
+}
+
+func length64(n uint64) []byte {
+	b := make([]byte, 9)
+	b[0] = len64Bit
+	binary.BigEndian.PutUint64(b[1:], n)
+	return b
+}
+
+func le32(n uint32) []byte {
+	b := make([]byte, 4)
+	binary.LittleEndian.PutUint32(b, n)
+	return b
+}
+
+func TestParseDeclaredLengthExceedsData(t *testing.T) {
+	const huge = 1 << 40
+	// intset: intSize 2, cardinality 1000, but no entries
+	intset := append(le32(2), le32(1000)...)
+	cases := []struct {
+		name string
+		data []byte
+	}{
+		{"string", rdbWith(0, length64(huge))},
+		{"lzf", rdbWith(0, []byte{encodeLZFPrefix, 0x01}, length64(huge), []byte{0x00})},
+		{"list", rdbWith(1, length64(huge))},
+		{"set", rdbWith(2, length64(huge))},
+		{"zset", rdbWith(3, length64(huge))},
+		{"intset", rdbWith(11, []byte{byte(len(intset))}, intset)},
+	}
+	for _, c := range cases {
+		c := c
+		t.Run(c.name, func(t *testing.T) {
+			dec := NewDecoder(bytes.NewReader(c.data))
+			err := dec.Parse(func(o model.RedisObject) bool { return true })
+			if err == nil {
+				t.Fatal("expect an error")
+			}
+		})
+	}
+}
+
+func TestParseValkeySlotImportHugeRanges(t *testing.T) {
+	data := []byte("VALKEY080")
+	data = append(data, opCodeSlotImport, 0x01, 'j')
+	data = append(data, length64(1<<40)...)
+	data = append(data, 0xff)
+	dec := NewDecoder(bytes.NewReader(data))
+	if err := dec.Parse(func(o model.RedisObject) bool { return true }); err == nil {
+		t.Fatal("expect an error")
+	}
+}
+
+func TestCheckCount(t *testing.T) {
+	cases := []struct {
+		name  string
+		count int64
+		size  int
+		ok    bool
+	}{
+		{"zero", 0, 0, true},
+		{"fits", 5, 5, true},
+		{"too many", 6, 5, false},
+		{"negative", -1, 5, false},
+		{"huge", 1 << 62, 5, false},
+	}
+	for _, c := range cases {
+		c := c
+		t.Run(c.name, func(t *testing.T) {
+			err := checkCount(c.count, c.size)
+			if (err == nil) != c.ok {
+				t.Fatalf("checkCount(%d, %d) = %v", c.count, c.size, err)
+			}
+		})
+	}
+}
+
+func TestCapHint(t *testing.T) {
+	if capHint(5) != 5 || capHint(1024) != 1024 || capHint(1<<40) != 1024 {
+		t.Fatal("unexpected capHint")
+	}
+}
+
+func TestReadStreamEntryContentHugeCount(t *testing.T) {
+	// listpack entries: count = 100000, deleted = 0 ; the buffer holds only a few bytes
+	dec := NewDecoder(bytes.NewReader(nil))
+	// entries: count (24-bit int 100000), deleted 0, master field number 0, end flag 0
+	buf := []byte{0xf2, 0xa0, 0x86, 0x01, 0x05, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01}
+	cursor := 0
+	_, err := dec.readStreamEntryContent(buf, &cursor, &model.StreamId{})
+	if err == nil || !strings.Contains(err.Error(), "does not fit") {
+		t.Fatalf("expect a count error, got %v", err)
+	}
+}
+
+func TestReadBytesAboveChunk(t *testing.T) {
+	n := readChunk + 10
+	data := make([]byte, n)
+	for i := range data {
+		data[i] = byte(i % 251)
+	}
+	t.Run("full data", func(t *testing.T) {
+		dec := NewDecoder(bytes.NewReader(data))
+		got, err := dec.readBytes(uint64(n))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !bytes.Equal(got, data) || cap(got) != n || dec.readCount != n {
+			t.Fatalf("bad result: len %d cap %d count %d", len(got), cap(got), dec.readCount)
+		}
+	})
+	truncated := []struct {
+		name string
+		size int
+	}{
+		{"partial data", readChunk + 5},
+		{"ends at chunk boundary", readChunk},
+	}
+	for _, c := range truncated {
+		t.Run(c.name, func(t *testing.T) {
+			dec := NewDecoder(bytes.NewReader(data[:c.size]))
+			if _, err := dec.readBytes(uint64(n)); !errors.Is(err, io.ErrUnexpectedEOF) {
+				t.Fatalf("expect io.ErrUnexpectedEOF, got %v", err)
+			}
+		})
+	}
+}

core/decoder.go

@@ -534,7 +534,7 @@ func (dec *Decoder) parse(cb func(object model.RedisObject) bool) error {
 					return err
 				}
 				var slot_from, slot_to uint64
-				ranges := make([]string, num_slot_ranges)
+				ranges := make([]string, 0, capHint(num_slot_ranges))
 				for i := uint64(0); i < num_slot_ranges; i++ {
 					slot_from, _, err = dec.readLength()
 					if err == nil {
@@ -543,7 +543,7 @@ func (dec *Decoder) parse(cb func(object model.RedisObject) bool) error {
 					if err != nil {
 						return err
 					}
-					ranges[i] = fmt.Sprintf("%d-%d", slot_from, slot_to)
+					ranges = append(ranges, fmt.Sprintf("%d-%d", slot_from, slot_to))
 				}
 				_, _ = job, ranges // safe to skip
 			} else {

core/list.go

@@ -26,9 +26,8 @@ func (dec *Decoder) readList() ([][]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	size := int(size64)
-	values := make([][]byte, 0, size)
-	for i := 0; i < size; i++ {
+	values := make([][]byte, 0, capHint(size64))
+	for i := uint64(0); i < size64; i++ {
 		val, err := dec.readString()
 		if err != nil {
 			return nil, err

core/set.go

@@ -15,9 +15,8 @@ func (dec *Decoder) readSet() ([][]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	size := int(size64)
-	values := make([][]byte, 0, size)
-	for i := 0; i < size; i++ {
+	values := make([][]byte, 0, capHint(size64))
+	for i := uint64(0); i < size64; i++ {
 		val, err := dec.readString()
 		if err != nil {
 			return nil, err
@@ -33,13 +32,19 @@ func (dec *Decoder) readIntSet() (result [][]byte, detail *model.IntsetDetail, e
 	if err != nil {
 		return nil, nil, err
 	}
+	if len(buf) < 8 {
+		return nil, nil, fmt.Errorf("intset is too short: %d bytes", len(buf))
+	}
 	sizeBytes := buf[0:4]
 	intSize := int(binary.LittleEndian.Uint32(sizeBytes))
 	if intSize != 2 && intSize != 4 && intSize != 8 {
 		return nil, nil, fmt.Errorf("unknown intset encoding: %d", intSize)
 	}
 	lenBytes := buf[4:8]
 	cardinality := binary.LittleEndian.Uint32(lenBytes)
+	if uint64(cardinality)*uint64(intSize) > uint64(len(buf)-8) {
+		return nil, nil, fmt.Errorf("intset cardinality %d does not fit in %d bytes", cardinality, len(buf))
+	}
 	cursor := 8
 	result = make([][]byte, 0, cardinality)
 	for i := uint32(0); i < cardinality; i++ {

core/stream.go

@@ -139,6 +139,9 @@ func (dec *Decoder) readStreamEntryContent(buf []byte, cursor *int, firstId *mod
 	if err != nil {
 		return nil, fmt.Errorf("read stream field number failed: %v", err)
 	}
+	if err = checkCount(fieldNum0, len(buf)); err != nil {
+		return nil, fmt.Errorf("read stream field number failed: %v", err)
+	}
 	masterFieldNum := int(fieldNum0)
 	masterFieldNames := make([]string, masterFieldNum)
 	for i := 0; i < masterFieldNum; i++ {
@@ -154,6 +157,15 @@ func (dec *Decoder) readStreamEntryContent(buf []byte, cursor *int, firstId *mod
 	}
 
 	total := count + deleted
+	if err = checkCount(count, len(buf)); err != nil {
+		return nil, fmt.Errorf("read stream entry count failed: %v", err)
+	}
+	if err = checkCount(deleted, len(buf)); err != nil {
+		return nil, fmt.Errorf("read stream entry deleted count failed: %v", err)
+	}
+	if err = checkCount(total, len(buf)); err != nil {
+		return nil, fmt.Errorf("read stream entry total count failed: %v", err)
+	}
 	msgs := make([]*model.StreamMessage, 0, total)
 	for i := int64(0); i < total; i++ {
 		flag, err := dec.readListPackEntryAsInt(buf, cursor)
@@ -179,6 +191,9 @@ func (dec *Decoder) readStreamEntryContent(buf []byte, cursor *int, firstId *mod
 			if err != nil {
 				return nil, fmt.Errorf("read stream item field number failed: %v", err)
 			}
+			if err = checkCount(fieldNum0, len(buf)); err != nil {
+				return nil, fmt.Errorf("read stream item field number failed: %v", err)
+			}
 			fieldNum = int(fieldNum0)
 		}
 		msg := &model.StreamMessage{
@@ -221,7 +236,7 @@ func (dec *Decoder) readStreamGroups(version uint) ([]*model.StreamGroup, error)
 	if err != nil {
 		return nil, err
 	}
-	groups := make([]*model.StreamGroup, 0, int(groupCount))
+	groups := make([]*model.StreamGroup, 0, capHint(groupCount))
 	for i := uint64(0); i < groupCount; i++ {
 		name, _ := dec.readString()
 		if err != nil {
@@ -245,7 +260,7 @@ func (dec *Decoder) readStreamGroups(version uint) ([]*model.StreamGroup, error)
 		if err != nil {
 			return nil, err
 		}
-		pending := make([]*model.StreamNAck, 0, int(pendingCount))
+		pending := make([]*model.StreamNAck, 0, capHint(pendingCount))
 		for j := uint64(0); j < pendingCount; j++ {
 			if err := dec.readFull(dec.buffer); err != nil {
 				return nil, err
@@ -279,7 +294,7 @@ func (dec *Decoder) readStreamGroups(version uint) ([]*model.StreamGroup, error)
 		if err != nil {
 			return nil, err
 		}
-		consumers := make([]*model.StreamConsumer, 0, int(consumerCount))
+		consumers := make([]*model.StreamConsumer, 0, capHint(consumerCount))
 		for j := uint64(0); j < consumerCount; j++ {
 			consumerName, err := dec.readString()
 			if err != nil {
@@ -300,7 +315,7 @@ func (dec *Decoder) readStreamGroups(version uint) ([]*model.StreamGroup, error)
 			if err != nil {
 				return nil, err
 			}
-			consumerPending := make([]*model.StreamId, 0, int(consumerPendingCount))
+			consumerPending := make([]*model.StreamId, 0, capHint(consumerPendingCount))
 			for k := uint64(0); k < consumerPendingCount; k++ {
 				if err := dec.readFull(dec.buffer); err != nil {
 					return nil, err

core/string.go

@@ -102,9 +102,7 @@ func (dec *Decoder) readString() ([]byte, error) {
 		}
 	}
 
-	res := make([]byte, length)
-	err = dec.readFull(res)
-	return res, err
+	return dec.readBytes(length)
 }
 
 func (dec *Decoder) readInt16() (int16, error) {
@@ -188,12 +186,14 @@ func (dec *Decoder) readLZF() ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	val := make([]byte, inLen)
-	err = dec.readFull(val)
+	val, err := dec.readBytes(inLen)
 	if err != nil {
 		return nil, err
 	}
-	return lzf.Decompress(val, int(inLen), int(outLen))
+	if outLen > math.MaxInt {
+		return nil, fmt.Errorf("lzf output length %d is too large", outLen)
+	}
+	return lzf.Decompress(val, len(val), int(outLen))
 }
 
 func (enc *Encoder) writeLength(value uint64) error {

core/utils.go

@@ -3,7 +3,9 @@ package core
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"io"
+	"math"
 	"math/rand"
 	"unsafe"
 )
@@ -60,6 +62,69 @@ func (dec *Decoder) readFull(buf []byte) error {
 	return nil
 }
 
+// readChunk is the largest length that readBytes allocates before it reads data.
+// A bigger declared length may be false, so readBytes grows the buffer as data arrives.
+const readChunk = 1 << 20
+
+// readBytes reads exactly n bytes from the input.
+// A length up to readChunk is allocated at once. A longer length grows the
+// buffer by doubling, so memory stays near twice the bytes that arrived.
+func (dec *Decoder) readBytes(n uint64) ([]byte, error) {
+	if n <= readChunk {
+		buf := make([]byte, n)
+		if err := dec.readFull(buf); err != nil {
+			return nil, err
+		}
+		return buf, nil
+	}
+	if n > math.MaxInt {
+		return nil, fmt.Errorf("declared length %d is too large", n)
+	}
+	size := int(n)
+	buf := make([]byte, readChunk)
+	if err := dec.readFull(buf); err != nil {
+		return nil, err
+	}
+	for len(buf) < size {
+		next := 2 * len(buf)
+		if next > size || next < 0 {
+			next = size
+		}
+		grown := make([]byte, next)
+		copy(grown, buf)
+		if err := dec.readFull(grown[len(buf):]); err != nil {
+			// Part of the value already arrived, so the end of the input here is a truncated value.
+			if err == io.EOF {
+				err = io.ErrUnexpectedEOF
+			}
+			return nil, err
+		}
+		buf = grown
+	}
+	return buf, nil
+}
+
+// maxCapHint limits the initial capacity taken from a count in the file.
+const maxCapHint = 1024
+
+// capHint returns an initial capacity for a count read from the file.
+// The count is not trusted, so append grows the slice when real data arrives.
+func capHint(n uint64) int {
+	if n > maxCapHint {
+		return maxCapHint
+	}
+	return int(n)
+}
+
+// checkCount rejects a count that cannot fit in a buffer of size bytes.
+// Each counted item uses at least one byte.
+func checkCount(count int64, size int) error {
+	if count < 0 || count > int64(size) {
+		return fmt.Errorf("count %d does not fit in a buffer of %d bytes", count, size)
+	}
+	return nil
+}
+
 var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
 
 // RandString create a random string no longer than n