Add more options to verify

buffer.go

@@ -36,6 +36,10 @@ func newMemoryBackedBuffer(buf []byte) *memoryBackedBuffer {
 
 func (b *memoryBackedBuffer) Get(n int) ([]byte, error) {
 	if b.len < b.pos+n {
+		if b.len == b.pos {
+			return nil, io.EOF
+		}
+
 		return nil, io.ErrUnexpectedEOF
 	}
 
@@ -158,6 +162,10 @@ func newFileBackedBuffer(file *os.File, fileLen int, bufCap int) *fileBackedBuff
 func (b *fileBackedBuffer) Get(n int) ([]byte, error) {
 	if b.fileLen < b.filePos+n {
 		// we use the file pos as the source of truth
+		if b.fileLen == b.filePos {
+			return nil, io.EOF
+		}
+
 		return nil, io.ErrUnexpectedEOF
 	}
 

file_reader.go

@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"strconv"
 	"time"
@@ -115,26 +116,27 @@ func readFile(buf buffer, handler FileHandler, maxLz77StrLen uint64) error {
 
 		switch t {
 		case typeOpCodeEOF:
-			if !endsWithCRC {
-				return nil
-			}
-
-			buf.DoNotCalcCrc()
-			crcFooter, err := buf.Get(crcLen)
-			if err != nil {
-				return err
-			}
+			if endsWithCRC {
+				buf.DoNotCalcCrc()
+				crcFooter, err := buf.Get(crcLen)
+				if err != nil {
+					return err
+				}
 
-			crc := binary.LittleEndian.Uint64(crcFooter)
-			if crc == 0 {
-				// crc calculation can be disabled by the redis config.
-				// if it is disabled, the crc bytes are still there but
-				// it is equal to 0.
-				return nil
+				crc := binary.LittleEndian.Uint64(crcFooter)
+				if crc != 0 && buf.Crc() != crc {
+					// crc calculation can be disabled by the redis config.
+					// if it is disabled, the crc bytes are still there but
+					// it is equal to 0.
+					return errors.New("wrong CRC at the end of the RDB file")
+				}
 			}
 
-			if buf.Crc() != crc {
-				return errors.New("wrong CRC at the end of the RDB file")
+			if handler.RequireStrictEOF() {
+				_, err = buf.Get(1)
+				if err != io.EOF {
+					return errors.New("required file to end after eof opcode")
+				}
 			}
 
 			return nil

file_reader_test.go

@@ -49,6 +49,10 @@ func (db *dummyDB) AllowPartialRead() bool {
 	return db.partialRead
 }
 
+func (db *dummyDB) RequireStrictEOF() bool {
+	return false
+}
+
 func (db *dummyDB) HandleString(key, value string) error {
 	db.strings[key] = value
 	return nil

handler.go

@@ -7,6 +7,10 @@ type ValueHandler interface {
 	// whether the handler can skip known but not yet supported types or not.
 	AllowPartialRead() bool
 
+	// whether the handler expects the file to end with the eof opcode, i.e,
+	// has no more bytes to the right of it.
+	RequireStrictEOF() bool
+
 	// called when a string value is read for the key.
 	HandleString(key, value string) error
 
@@ -62,6 +66,10 @@ func (nopHandler) AllowPartialRead() bool {
 	return true
 }
 
+func (nopHandler) RequireStrictEOF() bool {
+	return false
+}
+
 func (nopHandler) HandleString(key, value string) error {
 	return nil
 }

verify.go

@@ -16,10 +16,12 @@ var defaultMaxStreamPELSize = 1000
 const maxStreamStrSize = math.MaxUint32
 
 type VerifyFileOptions struct {
-	MaxDataSize      int
-	MaxEntrySize     int
-	MaxKeySize       int
-	MaxStreamPELSize int
+	MaxDataSize        int
+	MaxEntrySize       int
+	MaxKeySize         int
+	MaxStreamPELSize   int
+	AllowPartialVerify bool
+	RequireStrictEOF   bool
 }
 
 func (o *VerifyFileOptions) maybeSetDefaults() {
@@ -45,10 +47,12 @@ func (o *VerifyFileOptions) maybeSetDefaults() {
 func VerifyFile(path string, opts VerifyFileOptions) error {
 	opts.maybeSetDefaults()
 	v := &verifier{
-		maxDataSize:      opts.MaxDataSize,
-		maxEntrySize:     opts.MaxEntrySize,
-		maxKeySize:       opts.MaxKeySize,
-		maxStreamPELSize: opts.MaxStreamPELSize,
+		maxDataSize:        opts.MaxDataSize,
+		maxEntrySize:       opts.MaxEntrySize,
+		maxKeySize:         opts.MaxKeySize,
+		maxStreamPELSize:   opts.MaxStreamPELSize,
+		allowPartialVerify: opts.AllowPartialVerify,
+		requireStrictEOF:   opts.RequireStrictEOF,
 	}
 
 	file, err := os.Open(path)
@@ -69,10 +73,12 @@ func VerifyFile(path string, opts VerifyFileOptions) error {
 }
 
 type VerifyReaderOptions struct {
-	MaxDataSize      int
-	MaxEntrySize     int
-	MaxKeySize       int
-	MaxStreamPELSize int
+	MaxDataSize        int
+	MaxEntrySize       int
+	MaxKeySize         int
+	MaxStreamPELSize   int
+	AllowPartialVerify bool
+	RequireStrictEOF   bool
 }
 
 func (o *VerifyReaderOptions) maybeSetDefaults() {
@@ -98,10 +104,12 @@ func (o *VerifyReaderOptions) maybeSetDefaults() {
 func VerifyReader(r io.Reader, opts VerifyReaderOptions) error {
 	opts.maybeSetDefaults()
 	v := &verifier{
-		maxDataSize:      opts.MaxDataSize,
-		maxEntrySize:     opts.MaxEntrySize,
-		maxKeySize:       opts.MaxKeySize,
-		maxStreamPELSize: opts.MaxStreamPELSize,
+		maxDataSize:        opts.MaxDataSize,
+		maxEntrySize:       opts.MaxEntrySize,
+		maxKeySize:         opts.MaxKeySize,
+		maxStreamPELSize:   opts.MaxStreamPELSize,
+		allowPartialVerify: opts.AllowPartialVerify,
+		requireStrictEOF:   opts.RequireStrictEOF,
 	}
 
 	buf := newForwardOnlyBuffer(r)
@@ -147,11 +155,13 @@ var errMaxStreamPELSizeExceeded = errors.New("max stream pel size is exceeded")
 var errMaxStreamStrSizeExceeded = errors.New("max stream string item size is exceeded")
 
 type verifier struct {
-	maxDataSize      int
-	maxEntrySize     int
-	maxKeySize       int
-	maxStreamPELSize int
-	dataSize         int
+	maxDataSize        int
+	maxEntrySize       int
+	maxKeySize         int
+	maxStreamPELSize   int
+	allowPartialVerify bool
+	requireStrictEOF   bool
+	dataSize           int
 }
 
 func (v *verifier) HashWithExpEntryHandler(key string) func(field string, value string, exp time.Time) error {
@@ -425,7 +435,11 @@ func (v *verifier) StreamGroupHandler(key string) func(group StreamConsumerGroup
 }
 
 func (v *verifier) AllowPartialRead() bool {
-	return true
+	return v.allowPartialVerify
+}
+
+func (v *verifier) RequireStrictEOF() bool {
+	return v.requireStrictEOF
 }
 
 func (v *verifier) HandleExpireTime(key string, expireTime time.Duration) {

verify_test.go

@@ -13,6 +13,8 @@ var streamWithPELRDBPath = filepath.Join(dumpsPath, "stream-with-pel.rdb")
 var badCrcRDBPath = filepath.Join(dumpsPath, "bad-crc.rdb")
 var stringRDBValuePath = filepath.Join(valueDumpsPath, "string.bin")
 var streamWithPELRDBValuePath = filepath.Join(valueDumpsPath, "stream-listpacks3.bin")
+var multiDBRDBPath = filepath.Join(dumpsPath, "multi-db.rdb")
+var withPaddingRDBPath = filepath.Join(dumpsPath, "with-padding.rdb")
 
 func TestVerifyFile(t *testing.T) {
 	err := VerifyFile(allTypesRDBPath, VerifyFileOptions{})
@@ -24,6 +26,30 @@ func TestVerifyFile_withPEL(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestVerifyFile_AllowPartialRead(t *testing.T) {
+	err := VerifyFile(multiDBRDBPath, VerifyFileOptions{
+		AllowPartialVerify: true,
+	})
+	require.NoError(t, err)
+
+	err = VerifyFile(multiDBRDBPath, VerifyFileOptions{
+		AllowPartialVerify: false,
+	})
+	require.ErrorContains(t, err, "partial restore")
+}
+
+func TestVerifyFile_RequireStrictEOF(t *testing.T) {
+	err := VerifyFile(withPaddingRDBPath, VerifyFileOptions{
+		RequireStrictEOF: false,
+	})
+	require.NoError(t, err)
+
+	err = VerifyFile(withPaddingRDBPath, VerifyFileOptions{
+		RequireStrictEOF: true,
+	})
+	require.ErrorContains(t, err, "eof")
+}
+
 func TestVerifyFile_BadCrc(t *testing.T) {
 	err := VerifyFile(badCrcRDBPath, VerifyFileOptions{})
 	require.ErrorContains(t, err, "CRC")
@@ -175,3 +201,45 @@ func TestVerifyReader_BadCrc(t *testing.T) {
 	err = VerifyReader(file, VerifyReaderOptions{})
 	require.ErrorContains(t, err, "CRC")
 }
+
+func TestVerifyReader_AllowPartialRead(t *testing.T) {
+	file, err := os.Open(multiDBRDBPath)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		file.Close()
+	})
+
+	err = VerifyReader(file, VerifyReaderOptions{
+		AllowPartialVerify: true,
+	})
+	require.NoError(t, err)
+
+	_, err = file.Seek(0, 0)
+	require.NoError(t, err)
+
+	err = VerifyReader(file, VerifyReaderOptions{
+		AllowPartialVerify: false,
+	})
+	require.ErrorContains(t, err, "partial restore")
+}
+
+func TestVerifyReader_RequireStrictEOF(t *testing.T) {
+	file, err := os.Open(withPaddingRDBPath)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		file.Close()
+	})
+
+	err = VerifyReader(file, VerifyReaderOptions{
+		RequireStrictEOF: false,
+	})
+	require.NoError(t, err)
+
+	_, err = file.Seek(0, 0)
+	require.NoError(t, err)
+
+	err = VerifyReader(file, VerifyReaderOptions{
+		RequireStrictEOF: true,
+	})
+	require.ErrorContains(t, err, "eof")
+}