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) + } +}