From a273eef084755efb792b839ab9ff1617f1a0da79 Mon Sep 17 00:00:00 2001 From: Zsolt Gaspar <17578237+zsltg@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:02:42 +0200 Subject: [PATCH] fix: avoid allocating declared lengths before reading data The decoder read a length or count from the file and then allocated a buffer of that size before it read the data. A file of 22 bytes could declare a string of 2^40 bytes. The Go runtime then stopped with "fatal error: out of memory". Parse cannot recover from that error. Strings and LZF input up to 1 MiB are read as before. A longer declared length grows the buffer as the data arrives. A short input still returns the same EOF error as before. lzf.Decompress rejects an output length that is more than 88 times the input length, because valid data never exceeds that ratio. Counts that come from a buffer in memory (intset cardinality and stream listpack counts) are now checked against the buffer size. Counts that come from the file only set a capped initial capacity. The slice grows with append, and the loops still run for the real count. The Valkey slot import branch now uses append as well. Fixes #75 Signed-off-by: Zsolt Gaspar <17578237+zsltg@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- core/declared_length_test.go | 148 +++++++++++++++++++++++++++++++++++ core/decoder.go | 4 +- core/list.go | 5 +- core/set.go | 11 ++- core/stream.go | 23 +++++- core/string.go | 12 +-- core/utils.go | 65 +++++++++++++++ core/zset.go | 2 +- lzf/lzf.go | 10 +++ lzf/lzf_test.go | 22 ++++++ 10 files changed, 283 insertions(+), 19 deletions(-) create mode 100644 core/declared_length_test.go diff --git a/core/declared_length_test.go b/core/declared_length_test.go new file mode 100644 index 0000000..36af561 --- /dev/null +++ b/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) + } + }) + } +} diff --git a/core/decoder.go b/core/decoder.go index 2ae8604..237bf3c 100644 --- a/core/decoder.go +++ b/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 { diff --git a/core/list.go b/core/list.go index 57044ff..22d79fc 100644 --- a/core/list.go +++ b/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 diff --git a/core/set.go b/core/set.go index 62c9a69..c732042 100644 --- a/core/set.go +++ b/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,6 +32,9 @@ 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 { @@ -40,6 +42,9 @@ func (dec *Decoder) readIntSet() (result [][]byte, detail *model.IntsetDetail, e } 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++ { diff --git a/core/stream.go b/core/stream.go index 26793b6..0d1b2df 100644 --- a/core/stream.go +++ b/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 diff --git a/core/string.go b/core/string.go index 68c78d3..1191918 100644 --- a/core/string.go +++ b/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 { diff --git a/core/utils.go b/core/utils.go index f16af08..4a93125 100644 --- a/core/utils.go +++ b/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 diff --git a/core/zset.go b/core/zset.go index d0164c7..56a4f1f 100644 --- a/core/zset.go +++ b/core/zset.go @@ -11,7 +11,7 @@ func (dec *Decoder) readZSet(zset2 bool) ([]*model.ZSetEntry, error) { if err != nil { return nil, err } - entries := make([]*model.ZSetEntry, 0, int(length)) + entries := make([]*model.ZSetEntry, 0, capHint(length)) for i := uint64(0); i < length; i++ { member, err := dec.readString() if err != nil { diff --git a/lzf/lzf.go b/lzf/lzf.go index c11e4e8..4a00b1e 100644 --- a/lzf/lzf.go +++ b/lzf/lzf.go @@ -8,17 +8,27 @@ const ( maxLit = 1 << 5 maxOff = 1 << 13 maxRef = (1 << 8) + (1 << 3) + maxRatio = maxRef / 3 ) var ( errInsufficientBuffer = errors.New("insufficient buffer") errDataCorruption = errors.New("data corruption") + errInvalidOutLen = errors.New("output length is impossible for the input") ) // using https://github.com/zhuyie/golzf according to MIT license // Decompress decompress lzf compressed data func Decompress(input []byte, inLen int, outLen int) ([]byte, error) { input = input[:inLen] + if len(input) == 0 { + return nil, nil + } + // A back reference emits at most 264 bytes from 3 input bytes, and a literal + // run emits fewer bytes than it consumes, so valid data never exceeds 88x. + if outLen < 0 || int64(outLen) > int64(maxRatio)*int64(len(input)) { + return nil, errInvalidOutLen + } output := make([]byte, outLen) var inputIndex, outputIndex int diff --git a/lzf/lzf_test.go b/lzf/lzf_test.go index c00e965..11847b8 100644 --- a/lzf/lzf_test.go +++ b/lzf/lzf_test.go @@ -1,6 +1,7 @@ package lzf import ( + "math" "math/rand" "strings" "testing" @@ -36,3 +37,24 @@ func TestLzf(t *testing.T) { } } } + +func TestDecompressRejectsImpossibleOutLen(t *testing.T) { + // one literal byte cannot expand to math.MaxInt bytes + if _, err := Decompress([]byte{0x00}, 1, math.MaxInt); err == nil { + t.Fatal("expect an error") + } + // the largest valid ratio is accepted + if _, err := Decompress([]byte{0x00, 'a'}, 2, 88*2); err != nil { + t.Fatal(err) + } + if _, err := Decompress([]byte{0x00, 'a'}, 2, 88*2+1); err == nil { + t.Fatal("expect an error") + } +} + +func TestDecompressEmptyInput(t *testing.T) { + out, err := Decompress(nil, 0, 0) + if err != nil || out != nil { + t.Fatalf("got %v, %v", out, err) + } +}