Merge pull request #419 from bootjp/copilot/sub-pr-414

store: fix native snapshot restore bounds, sentinel errors, and reduce test allocations

Author: bootjp

store/lsm_store.go

@@ -683,7 +683,7 @@ func readRestoreEntry(r io.Reader, keyBuf *[]byte) (kLen, vLen int, eof bool, er
 	if _, err = io.ReadFull(r, (*keyBuf)[:kLen]); err != nil {
 		return 0, 0, false, errors.WithStack(err)
 	}
-	vLen, err = readRestoreFieldLen(r, "snapshot value", maxSnapshotValueSize)
+	vLen, err = readRestoreFieldLen(r, "snapshot value", maxSnapshotValueSize+valueHeaderSize)
 	if err != nil {
 		return 0, 0, false, err
 	}
@@ -700,7 +700,14 @@ func readRestoreFieldLen(r io.Reader, field string, maxSize int) (int, error) {
 
 func restoreFieldLenInt(raw uint64, field string, maxSize int) (int, error) {
 	if raw > uint64(math.MaxInt) || int(raw) > maxSize {
-		return 0, errors.WithStack(errors.Newf("%s length %d > %d", field, raw, maxSize))
+		switch field {
+		case "snapshot key":
+			return 0, errors.WithStack(errors.Wrapf(ErrSnapshotKeyTooLarge, "length %d > %d", raw, maxSize))
+		case "snapshot value":
+			return 0, errors.WithStack(errors.Wrapf(ErrValueTooLarge, "length %d > %d", raw, maxSize))
+		default:
+			return 0, errors.WithStack(errors.Newf("%s length %d > %d", field, raw, maxSize))
+		}
 	}
 	return int(raw), nil
 }

store/lsm_store_test.go

@@ -514,9 +514,12 @@ func TestPebbleStore_Restore_NativePebbleAtomic(t *testing.T) {
 }
 
 // TestPebbleStore_PutAt_ValueTooLarge verifies that PutAt rejects values
-// exceeding maxSnapshotValueSize (256 MiB) to ensure write and restore are
-// consistent.
+// exceeding maxSnapshotValueSize to ensure write and restore are consistent.
 func TestPebbleStore_PutAt_ValueTooLarge(t *testing.T) {
+	orig := maxSnapshotValueSize
+	maxSnapshotValueSize = 100
+	t.Cleanup(func() { maxSnapshotValueSize = orig })
+
 	dir, err := os.MkdirTemp("", "pebble-value-limit-*")
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
@@ -535,6 +538,10 @@ func TestPebbleStore_PutAt_ValueTooLarge(t *testing.T) {
 // TestPebbleStore_ApplyMutations_ValueTooLarge verifies that ApplyMutations
 // rejects Put mutations with values exceeding maxSnapshotValueSize.
 func TestPebbleStore_ApplyMutations_ValueTooLarge(t *testing.T) {
+	orig := maxSnapshotValueSize
+	maxSnapshotValueSize = 100
+	t.Cleanup(func() { maxSnapshotValueSize = orig })
+
 	dir, err := os.MkdirTemp("", "pebble-apply-value-limit-*")
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
@@ -555,6 +562,10 @@ func TestPebbleStore_ApplyMutations_ValueTooLarge(t *testing.T) {
 // TestMVCCStore_PutAt_ValueTooLarge verifies that the in-memory mvccStore
 // also rejects oversized values.
 func TestMVCCStore_PutAt_ValueTooLarge(t *testing.T) {
+	orig := maxSnapshotValueSize
+	maxSnapshotValueSize = 100
+	t.Cleanup(func() { maxSnapshotValueSize = orig })
+
 	s := NewMVCCStore()
 	defer s.Close()
 
@@ -568,6 +579,10 @@ func TestMVCCStore_PutAt_ValueTooLarge(t *testing.T) {
 // TestMVCCStore_ApplyMutations_ValueTooLarge verifies that the in-memory
 // mvccStore rejects oversized Put mutations.
 func TestMVCCStore_ApplyMutations_ValueTooLarge(t *testing.T) {
+	orig := maxSnapshotValueSize
+	maxSnapshotValueSize = 100
+	t.Cleanup(func() { maxSnapshotValueSize = orig })
+
 	s := NewMVCCStore()
 	defer s.Close()
 

store/mvcc_store.go

@@ -28,11 +28,16 @@ type VersionedValue struct {
 const (
 	checksumSize            = 4
 	mvccSnapshotVersion     = uint32(1)
-	maxSnapshotKeySize      = 1 << 20   // 1 MiB per key
-	maxSnapshotVersionCount = 1 << 20   // 1M versions per key
-	maxSnapshotValueSize    = 256 << 20 // 256 MiB per value; prevents OOM from malformed snapshots
+	maxSnapshotKeySize      = 1 << 20 // 1 MiB per key
+	maxSnapshotVersionCount = 1 << 20 // 1M versions per key
 )
 
+// maxSnapshotValueSize caps the allowed size of a single value during snapshot
+// restore and write paths to prevent OOM from malformed or adversarial snapshots.
+// Declared as a var so tests can temporarily lower the limit without allocating
+// hundreds of MiB.
+var maxSnapshotValueSize = 256 << 20 // 256 MiB
+
 var mvccSnapshotMagic = [8]byte{'E', 'K', 'V', 'M', 'V', 'C', 'C', '2'}
 
 // mvccSnapshot is retained for backward compatibility with older gob-backed
@@ -848,7 +853,7 @@ 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 {
+	if valueLen > uint64(maxSnapshotValueSize) {
 		return VersionedValue{}, errors.Wrapf(ErrValueTooLarge, "%d > %d", valueLen, maxSnapshotValueSize)
 	}
 	value := make([]byte, valueLen)