Add checksum validation to store snapshots

store/memory_store.go

@@ -3,7 +3,9 @@ package store
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"encoding/gob"
+	"hash/crc32"
 	"io"
 	"log/slog"
 	"os"
@@ -209,22 +211,38 @@ func (s *memoryStore) Snapshot() (io.ReadWriter, error) {
 	s.mtx.RUnlock()
 
 	buf := &bytes.Buffer{}
-	err := gob.NewEncoder(buf).Encode(cl)
-	if err != nil {
+	if err := gob.NewEncoder(buf).Encode(cl); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	sum := crc32.ChecksumIEEE(buf.Bytes())
+	if err := binary.Write(buf, binary.LittleEndian, sum); err != nil {
 		return nil, errors.WithStack(err)
 	}
 
 	return buf, nil
 }
-func (s *memoryStore) Restore(buf io.Reader) error {
+func (s *memoryStore) Restore(r io.Reader) error {
 	s.mtx.Lock()
 	defer s.mtx.Unlock()
 
-	s.m = make(map[uint64][]byte)
-	err := gob.NewDecoder(buf).Decode(&s.m)
+	data, err := io.ReadAll(r)
 	if err != nil {
 		return errors.WithStack(err)
 	}
+	if len(data) < 4 {
+		return errors.WithStack(ErrInvalidChecksum)
+	}
+	payload := data[:len(data)-4]
+	expected := binary.LittleEndian.Uint32(data[len(data)-4:])
+	if crc32.ChecksumIEEE(payload) != expected {
+		return errors.WithStack(ErrInvalidChecksum)
+	}
+
+	s.m = make(map[uint64][]byte)
+	if err := gob.NewDecoder(bytes.NewReader(payload)).Decode(&s.m); err != nil {
+		return errors.WithStack(err)
+	}
 
 	return nil
 }

store/memory_store_test.go

@@ -1,6 +1,7 @@
 package store
 
 import (
+	"bytes"
 	"context"
 	"strconv"
 	"sync"
@@ -204,6 +205,46 @@ func TestMemoryStore_TTL(t *testing.T) {
 	wg.Wait()
 }
 
+func TestMemoryStore_SnapshotChecksum(t *testing.T) {
+	t.Parallel()
+
+	t.Run("success", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		st := NewMemoryStore()
+		assert.NoError(t, st.Put(ctx, []byte("foo"), []byte("bar")))
+
+		buf, err := st.Snapshot()
+		assert.NoError(t, err)
+
+		st2 := NewMemoryStore()
+		err = st2.Restore(bytes.NewReader(buf.(*bytes.Buffer).Bytes()))
+		assert.NoError(t, err)
+
+		v, err := st2.Get(ctx, []byte("foo"))
+		assert.NoError(t, err)
+		assert.Equal(t, []byte("bar"), v)
+	})
+
+	t.Run("corrupted", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		st := NewMemoryStore()
+		assert.NoError(t, st.Put(ctx, []byte("foo"), []byte("bar")))
+
+		buf, err := st.Snapshot()
+		assert.NoError(t, err)
+		data := buf.(*bytes.Buffer).Bytes()
+		corrupted := make([]byte, len(data))
+		copy(corrupted, data)
+		corrupted[0] ^= 0xff
+
+		st2 := NewMemoryStore()
+		err = st2.Restore(bytes.NewReader(corrupted))
+		assert.ErrorIs(t, err, ErrInvalidChecksum)
+	})
+}
+
 func TestMemoryStore_TTL_Txn(t *testing.T) {
 	ctx := context.Background()
 	t.Parallel()

store/rb_memory_store.go

@@ -3,7 +3,9 @@ package store
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"encoding/gob"
+	"hash/crc32"
 	"io"
 	"log/slog"
 	"os"
@@ -229,22 +231,46 @@ func (s *rbMemoryStore) Snapshot() (io.ReadWriter, error) {
 	s.mtx.RUnlock()
 
 	buf := &bytes.Buffer{}
-	err := gob.NewEncoder(buf).Encode(cl)
-	if err != nil {
+	if err := gob.NewEncoder(buf).Encode(cl); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	sum := crc32.ChecksumIEEE(buf.Bytes())
+	if err := binary.Write(buf, binary.LittleEndian, sum); err != nil {
 		return nil, errors.WithStack(err)
 	}
 
 	return buf, nil
 }
-func (s *rbMemoryStore) Restore(buf io.Reader) error {
+func (s *rbMemoryStore) Restore(r io.Reader) error {
 	s.mtx.Lock()
 	defer s.mtx.Unlock()
 
-	s.tree.Clear()
-	err := gob.NewDecoder(buf).Decode(&s.tree)
+	data, err := io.ReadAll(r)
 	if err != nil {
 		return errors.WithStack(err)
 	}
+	if len(data) < 4 {
+		return errors.WithStack(ErrInvalidChecksum)
+	}
+	payload := data[:len(data)-4]
+	expected := binary.LittleEndian.Uint32(data[len(data)-4:])
+	if crc32.ChecksumIEEE(payload) != expected {
+		return errors.WithStack(ErrInvalidChecksum)
+	}
+
+	var cl map[*[]byte][]byte
+	if err := gob.NewDecoder(bytes.NewReader(payload)).Decode(&cl); err != nil {
+		return errors.WithStack(err)
+	}
+
+	s.tree.Clear()
+	for k, v := range cl {
+		if k == nil {
+			continue
+		}
+		s.tree.Put(*k, v)
+	}
 
 	return nil
 }

store/rb_memory_store_test.go

@@ -1,6 +1,7 @@
 package store
 
 import (
+	"bytes"
 	"context"
 	"encoding/binary"
 	"strconv"
@@ -205,6 +206,46 @@ func TestRbMemoryStore_Txn(t *testing.T) {
 	})
 }
 
+func TestRbMemoryStore_SnapshotChecksum(t *testing.T) {
+	t.Parallel()
+
+	t.Run("success", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		st := NewRbMemoryStore()
+		assert.NoError(t, st.Put(ctx, []byte("foo"), []byte("bar")))
+
+		buf, err := st.Snapshot()
+		assert.NoError(t, err)
+
+		st2 := NewRbMemoryStore()
+		err = st2.Restore(bytes.NewReader(buf.(*bytes.Buffer).Bytes()))
+		assert.NoError(t, err)
+
+		v, err := st2.Get(ctx, []byte("foo"))
+		assert.NoError(t, err)
+		assert.Equal(t, []byte("bar"), v)
+	})
+
+	t.Run("corrupted", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		st := NewRbMemoryStore()
+		assert.NoError(t, st.Put(ctx, []byte("foo"), []byte("bar")))
+
+		buf, err := st.Snapshot()
+		assert.NoError(t, err)
+		data := buf.(*bytes.Buffer).Bytes()
+		corrupted := make([]byte, len(data))
+		copy(corrupted, data)
+		corrupted[0] ^= 0xff
+
+		st2 := NewRbMemoryStore()
+		err = st2.Restore(bytes.NewReader(corrupted))
+		assert.ErrorIs(t, err, ErrInvalidChecksum)
+	})
+}
+
 func TestRbMemoryStore_TTL(t *testing.T) {
 	ctx := context.Background()
 	t.Parallel()

store/store.go

@@ -10,6 +10,7 @@ import (
 var ErrKeyNotFound = errors.New("not found")
 var ErrUnknownOp = errors.New("unknown op")
 var ErrNotSupported = errors.New("not supported")
+var ErrInvalidChecksum = errors.New("invalid checksum")
 
 type KVPair struct {
 	Key   []byte