From 8eea651425b58119040504d5676290def1fef6fc Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Tue, 17 Sep 2024 12:02:04 +0000
Subject: [PATCH] feat(runtime(v2)): add sanity checks on store upgrades
 (backport #21748) (#21778)

Co-authored-by: Julien Robert <julien@rbrt.fr>
---
 runtime/store.go      | 42 ++++++++++++++++++++++++++++
 runtime/store_test.go | 65 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+)
 create mode 100644 runtime/store_test.go

diff --git a/runtime/store.go b/runtime/store.go
index 1585bc54af..a38b24ad2f 100644
--- a/runtime/store.go
+++ b/runtime/store.go
@@ -2,6 +2,8 @@ package runtime
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"io"
 
 	dbm "github.com/cosmos/cosmos-db"
@@ -186,6 +188,11 @@ func KVStoreAdapter(store store.KVStore) storetypes.KVStore {
 // UpgradeStoreLoader is used to prepare baseapp with a fixed StoreLoader
 // pattern. This is useful for custom upgrade loading logic.
 func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades) baseapp.StoreLoader {
+	// sanity checks on store upgrades
+	if err := checkStoreUpgrade(storeUpgrades); err != nil {
+		panic(err)
+	}
+
 	return func(ms storetypes.CommitMultiStore) error {
 		if upgradeHeight == ms.LastCommitID().Version+1 {
 			// Check if the current commit version and upgrade height matches
@@ -202,3 +209,38 @@ func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades)
 		return baseapp.DefaultStoreLoader(ms)
 	}
 }
+
+// checkStoreUpgrade performs sanity checks on the store upgrades
+func checkStoreUpgrade(storeUpgrades *store.StoreUpgrades) error {
+	if storeUpgrades == nil {
+		return errors.New("store upgrades cannot be nil")
+	}
+
+	// check for duplicates
+	exists := make(map[string]bool)
+	for _, key := range storeUpgrades.Added {
+		if exists[key] {
+			return fmt.Errorf("store upgrade has duplicate key %s in added", key)
+		}
+
+		if storeUpgrades.IsDeleted(key) {
+			return fmt.Errorf("store upgrade has key %s in both added and deleted", key)
+		}
+
+		exists[key] = true
+	}
+	exists = make(map[string]bool)
+	for _, key := range storeUpgrades.Deleted {
+		if exists[key] {
+			return fmt.Errorf("store upgrade has duplicate key %s in deleted", key)
+		}
+
+		if storeUpgrades.IsAdded(key) {
+			return fmt.Errorf("store upgrade has key %s in both added and deleted", key)
+		}
+
+		exists[key] = true
+	}
+
+	return nil
+}
diff --git a/runtime/store_test.go b/runtime/store_test.go
new file mode 100644
index 0000000000..a583a08d03
--- /dev/null
+++ b/runtime/store_test.go
@@ -0,0 +1,65 @@
+package runtime
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	corestore "cosmossdk.io/core/store"
+)
+
+func TestCheckStoreUpgrade(t *testing.T) {
+	tests := []struct {
+		name          string
+		storeUpgrades *corestore.StoreUpgrades
+		errMsg        string
+	}{
+		{
+			name:          "Nil StoreUpgrades",
+			storeUpgrades: nil,
+			errMsg:        "store upgrades cannot be nil",
+		},
+		{
+			name: "Valid StoreUpgrades",
+			storeUpgrades: &corestore.StoreUpgrades{
+				Added:   []string{"store1", "store2"},
+				Deleted: []string{"store3", "store4"},
+			},
+		},
+		{
+			name: "Duplicate key in Added",
+			storeUpgrades: &corestore.StoreUpgrades{
+				Added:   []string{"store1", "store2", "store1"},
+				Deleted: []string{"store3"},
+			},
+			errMsg: "store upgrade has duplicate key store1 in added",
+		},
+		{
+			name: "Duplicate key in Deleted",
+			storeUpgrades: &corestore.StoreUpgrades{
+				Added:   []string{"store1"},
+				Deleted: []string{"store2", "store3", "store2"},
+			},
+			errMsg: "store upgrade has duplicate key store2 in deleted",
+		},
+		{
+			name: "Key in both Added and Deleted",
+			storeUpgrades: &corestore.StoreUpgrades{
+				Added:   []string{"store1", "store2"},
+				Deleted: []string{"store2", "store3"},
+			},
+			errMsg: "store upgrade has key store2 in both added and deleted",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := checkStoreUpgrade(tt.storeUpgrades)
+			if tt.errMsg == "" {
+				require.NoError(t, err)
+			} else {
+				require.ErrorContains(t, err, tt.errMsg)
+			}
+		})
+	}
+}