Address review feedback: snapshot safety and correctness fixes

Co-authored-by: bootjp <1306365+bootjp@users.noreply.github.com>

Author: Copilot

kv/snapshot_test.go

@@ -58,6 +58,7 @@ func TestSnapshot(t *testing.T) {
 	var buf bytes.Buffer
 	_, err = kvFSMSnap.snapshot.WriteTo(&buf)
 	assert.NoError(t, err)
+	kvFSMSnap.Release()
 	err = fsm2.Restore(io.NopCloser(bytes.NewReader(buf.Bytes())))
 	assert.NoError(t, err)
 

store/mvcc_store.go

@@ -26,8 +26,11 @@ type VersionedValue struct {
 }
 
 const (
-	checksumSize        = 4
-	mvccSnapshotVersion = uint32(1)
+	checksumSize              = 4
+	mvccSnapshotVersion       = uint32(1)
+	maxSnapshotKeySize        = 1 << 20        // 1 MiB per key
+	maxSnapshotVersionCount   = 1 << 20        // 1M versions per key
+	maxSnapshotValueSize      = 64 << 20       // 64 MiB per value
 )
 
 var mvccSnapshotMagic = [8]byte{'E', 'K', 'V', 'M', 'V', 'C', 'C', '2'}
@@ -606,7 +609,8 @@ func writeMVCCSnapshotHeader(f *os.File) (int64, error) {
 
 func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) {
 	hash := crc32.NewIEEE()
-	w := io.MultiWriter(f, hash)
+	bw := bufio.NewWriter(f)
+	w := io.MultiWriter(bw, hash)
 
 	s.mtx.RLock()
 	defer s.mtx.RUnlock()
@@ -631,6 +635,9 @@ func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) {
 			return 0, err
 		}
 	}
+	if err := bw.Flush(); err != nil {
+		return 0, errors.WithStack(err)
+	}
 	return hash.Sum32(), nil
 }
 
@@ -788,6 +795,9 @@ func readMVCCSnapshotEntry(r io.Reader) ([]byte, []VersionedValue, bool, error)
 		}
 		return nil, nil, false, errors.WithStack(err)
 	}
+	if keyLen > maxSnapshotKeySize {
+		return nil, nil, false, errors.Newf("mvcc snapshot key too large: %d > %d", keyLen, maxSnapshotKeySize)
+	}
 
 	key := make([]byte, keyLen)
 	if _, err := io.ReadFull(r, key); err != nil {
@@ -798,6 +808,9 @@ func readMVCCSnapshotEntry(r io.Reader) ([]byte, []VersionedValue, bool, error)
 	if err := binary.Read(r, binary.LittleEndian, &versionCount); err != nil {
 		return nil, nil, false, errors.WithStack(err)
 	}
+	if versionCount > maxSnapshotVersionCount {
+		return nil, nil, false, errors.Newf("mvcc snapshot version count too large: %d > %d", versionCount, maxSnapshotVersionCount)
+	}
 	versions := make([]VersionedValue, 0, versionCount)
 	for i := uint64(0); i < versionCount; i++ {
 		version, err := readMVCCSnapshotVersion(r)
@@ -829,6 +842,9 @@ func readMVCCSnapshotVersion(r io.Reader) (VersionedValue, error) {
 	if err := binary.Read(r, binary.LittleEndian, &valueLen); err != nil {
 		return VersionedValue{}, errors.WithStack(err)
 	}
+	if valueLen > maxSnapshotValueSize {
+		return VersionedValue{}, errors.Newf("mvcc snapshot value too large: %d > %d", valueLen, maxSnapshotValueSize)
+	}
 	value := make([]byte, valueLen)
 	if _, err := io.ReadFull(r, value); err != nil {
 		return VersionedValue{}, errors.WithStack(err)

store/mvcc_store_snapshot_test.go

@@ -3,6 +3,9 @@ package store
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
+	"encoding/gob"
+	"hash/crc32"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -81,6 +84,43 @@ func TestMVCCStore_ApplyMutations_UnknownOp(t *testing.T) {
 	require.ErrorIs(t, err, ErrUnknownOp)
 }
 
+func TestMVCCStore_RestoreLegacySnapshot(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	// Build a legacy gob+crc32 snapshot payload.
+	legacy := mvccSnapshot{
+		LastCommitTS:  42,
+		MinRetainedTS: 10,
+		Entries: []mvccSnapshotEntry{
+			{
+				Key: []byte("legacy-key"),
+				Versions: []VersionedValue{
+					{TS: 42, Value: []byte("legacy-value"), Tombstone: false, ExpireAt: 0},
+				},
+			},
+		},
+	}
+
+	var payload bytes.Buffer
+	require.NoError(t, gob.NewEncoder(&payload).Encode(legacy))
+
+	checksum := crc32.ChecksumIEEE(payload.Bytes())
+	var cs [4]byte
+	binary.LittleEndian.PutUint32(cs[:], checksum)
+	full := append(payload.Bytes(), cs[:]...)
+
+	dst := newTestMVCCStore(t)
+	require.NoError(t, dst.Restore(bytes.NewReader(full)))
+
+	require.Equal(t, uint64(42), dst.LastCommitTS())
+
+	v, err := dst.GetAt(ctx, []byte("legacy-key"), 100)
+	require.NoError(t, err)
+	require.Equal(t, []byte("legacy-value"), v)
+}
+
 func snapshotBytes(t *testing.T, snap Snapshot) []byte {
 	t.Helper()
 

store/snapshot_pebble.go

@@ -70,26 +70,33 @@ func writePebbleSnapshotEntries(snap *pebble.Snapshot, w io.Writer) error {
 	if err != nil {
 		return errors.WithStack(err)
 	}
-	defer iter.Close()
 
 	for iter.First(); iter.Valid(); iter.Next() {
 		k := iter.Key()
 		v := iter.Value()
 
 		if err := binary.Write(w, binary.LittleEndian, uint64(len(k))); err != nil {
+			_ = iter.Close()
 			return errors.WithStack(err)
 		}
 		if _, err := w.Write(k); err != nil {
+			_ = iter.Close()
 			return errors.WithStack(err)
 		}
 		if err := binary.Write(w, binary.LittleEndian, uint64(len(v))); err != nil {
+			_ = iter.Close()
 			return errors.WithStack(err)
 		}
 		if _, err := w.Write(v); err != nil {
+			_ = iter.Close()
 			return errors.WithStack(err)
 		}
 	}
-	return nil
+	if err := iter.Error(); err != nil {
+		_ = iter.Close()
+		return errors.WithStack(err)
+	}
+	return errors.WithStack(iter.Close())
 }
 
 type countingWriter struct {