From 40a6c46a2fdb97295b325297fc30b2fee6b412e4 Mon Sep 17 00:00:00 2001
From: Mathieu Hofman <mathieu@agoric.com>
Date: Tue, 12 Sep 2023 20:28:58 +0000
Subject: [PATCH] fix(snapshots): raise the per snapshot item limit

---
 snapshots/manager.go      |  4 +-
 snapshots/manager_test.go | 78 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/snapshots/manager.go b/snapshots/manager.go
index 1ec632718ecf..9e0a5dd79a82 100644
--- a/snapshots/manager.go
+++ b/snapshots/manager.go
@@ -21,7 +21,9 @@ const (
 
 	chunkBufferSize = 4
 
-	snapshotMaxItemSize = int(64e6) // SDK has no key/value size limit, so we set an arbitrary limit
+	// snapshotMaxItemSize limits the size of both KVStore entries and snapshot
+	// extension payloads during a state-sync restore.
+	snapshotMaxItemSize = int(512e6)
 )
 
 // operation represents a Manager operation. Only one operation can be in progress at a time.
diff --git a/snapshots/manager_test.go b/snapshots/manager_test.go
index efd3a7f78dfd..6542205b3ac7 100644
--- a/snapshots/manager_test.go
+++ b/snapshots/manager_test.go
@@ -2,6 +2,7 @@ package snapshots_test
 
 import (
 	"errors"
+	"io"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -218,3 +219,80 @@ func TestManager_Restore(t *testing.T) {
 	})
 	require.NoError(t, err)
 }
+
+const snapshotMaxItemSize = int(512e6) // Copied from github.com/cosmos/cosmos-sdk/snapshots
+
+func TestManager_RestoreLargeItem(t *testing.T) {
+	store := setupStore(t)
+	target := &mockSnapshotter{}
+	extSnapshotter := newExtSnapshotter(0)
+	manager := snapshots.NewManager(store, target)
+	err := manager.RegisterExtensions(extSnapshotter)
+	require.NoError(t, err)
+
+	// The protobuf wrapper introduces an extra byte
+	largeItem := make([]byte, snapshotMaxItemSize-1)
+	expectItems := [][]byte{largeItem}
+
+	chunks := snapshotItems(expectItems, newExtSnapshotter(1))
+
+	// Starting a restore works
+	err = manager.Restore(types.Snapshot{
+		Height:   3,
+		Format:   2,
+		Hash:     []byte{1, 2, 3},
+		Chunks:   1,
+		Metadata: types.Metadata{ChunkHashes: checksums(chunks)},
+	})
+	require.NoError(t, err)
+
+	// Feeding the chunks should work
+	for i, chunk := range chunks {
+		done, err := manager.RestoreChunk(chunk)
+		require.NoError(t, err)
+		if i == len(chunks)-1 {
+			assert.True(t, done)
+		} else {
+			assert.False(t, done)
+		}
+	}
+
+	assert.Equal(t, expectItems, target.items)
+	assert.Equal(t, 1, len(extSnapshotter.state))
+}
+
+func TestManager_CannotRestoreTooLargeItem(t *testing.T) {
+	store := setupStore(t)
+	target := &mockSnapshotter{}
+	extSnapshotter := newExtSnapshotter(0)
+	manager := snapshots.NewManager(store, target)
+	err := manager.RegisterExtensions(extSnapshotter)
+	require.NoError(t, err)
+
+	// The protobuf wrapper introduces an extra byte
+	largeItem := make([]byte, snapshotMaxItemSize)
+	expectItems := [][]byte{largeItem}
+
+	chunks := snapshotItems(expectItems, newExtSnapshotter(1))
+
+	// Starting a restore works
+	err = manager.Restore(types.Snapshot{
+		Height:   3,
+		Format:   2,
+		Hash:     []byte{1, 2, 3},
+		Chunks:   1,
+		Metadata: types.Metadata{ChunkHashes: checksums(chunks)},
+	})
+	require.NoError(t, err)
+
+	// Feeding the chunks fails
+	for _, chunk := range chunks {
+		_, err = manager.RestoreChunk(chunk)
+
+		if err != nil {
+			break
+		}
+	}
+	require.Error(t, err)
+	require.True(t, errors.Is(err, io.ErrShortBuffer))
+}