fix!: invalid return value from keeper GetLastCompleteUpgrade - breaking change version

CHANGELOG.md

@@ -272,6 +272,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### State Machine Breaking
 
+* (x/upgrade) [\#11800](https://github.com/cosmos/cosmos-sdk/pull/11800) Fix `GetLastCompleteUpgrade` to properly return the latest upgrade.
 * [\#10564](https://github.com/cosmos/cosmos-sdk/pull/10564) Fix bug when updating allowance inside AllowedMsgAllowance
 * (x/auth)[\#9596](https://github.com/cosmos/cosmos-sdk/pull/9596) Enable creating periodic vesting accounts with a transactions instead of requiring them to be created in genesis.
 * (x/bank) [\#9611](https://github.com/cosmos/cosmos-sdk/pull/9611) Introduce a new index to act as a reverse index between a denomination and address allowing to query for

x/upgrade/keeper/keeper.go

@@ -14,7 +14,6 @@ import (
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 
 	"github.com/cosmos/cosmos-sdk/codec"
-	"github.com/cosmos/cosmos-sdk/internal/conv"
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -239,28 +238,43 @@ func (k Keeper) GetUpgradedConsensusState(ctx sdk.Context, lastHeight int64) ([]
 func (k Keeper) GetLastCompletedUpgrade(ctx sdk.Context) (string, int64) {
 	iter := sdk.KVStoreReversePrefixIterator(ctx.KVStore(k.storeKey), []byte{types.DoneByte})
 	defer iter.Close()
+
 	if iter.Valid() {
-		return parseDoneKey(iter.Key()), int64(binary.BigEndian.Uint64(iter.Value()))
+		return parseDoneKey(iter.Key())
 	}
 
 	return "", 0
 }
 
-// parseDoneKey - split upgrade name from the done key
-func parseDoneKey(key []byte) string {
-	kv.AssertKeyAtLeastLength(key, 2)
-	return string(key[1:])
+// parseDoneKey - split upgrade name and height from the done key
+func parseDoneKey(key []byte) (string, int64) {
+	// 1 byte for the DoneByte + 8 bytes height + at least 1 byte for the name
+	kv.AssertKeyAtLeastLength(key, 10)
+	height := binary.BigEndian.Uint64(key[1:9])
+	return string(key[9:]), int64(height)
+}
+
+// encodeDoneKey - concatenate DoneByte, height and upgrade name to form the done key
+func encodeDoneKey(name string, height int64) []byte {
+	key := make([]byte, 9+len(name)) // 9 = donebyte + uint64 len
+	key[0] = types.DoneByte
+	binary.BigEndian.PutUint64(key[1:9], uint64(height))
+	copy(key[9:], name)
+	return key
 }
 
 // GetDoneHeight returns the height at which the given upgrade was executed
 func (k Keeper) GetDoneHeight(ctx sdk.Context, name string) int64 {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte{types.DoneByte})
-	bz := store.Get(conv.UnsafeStrToBytes(name))
-	if len(bz) == 0 {
-		return 0
-	}
+	iter := sdk.KVStorePrefixIterator(ctx.KVStore(k.storeKey), []byte{types.DoneByte})
+	defer iter.Close()
 
-	return int64(binary.BigEndian.Uint64(bz))
+	for ; iter.Valid(); iter.Next() {
+		upgradeName, height := parseDoneKey(iter.Key())
+		if upgradeName == name {
+			return height
+		}
+	}
+	return 0
 }
 
 // ClearIBCState clears any planned IBC state
@@ -303,10 +317,8 @@ func (k Keeper) GetUpgradePlan(ctx sdk.Context) (plan types.Plan, havePlan bool)
 
 // setDone marks this upgrade name as being done so the name can't be reused accidentally
 func (k Keeper) setDone(ctx sdk.Context, name string) {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte{types.DoneByte})
-	bz := make([]byte, 8)
-	binary.BigEndian.PutUint64(bz, uint64(ctx.BlockHeight()))
-	store.Set([]byte(name), bz)
+	store := ctx.KVStore(k.storeKey)
+	store.Set(encodeDoneKey(name, ctx.BlockHeight()), []byte{1})
 }
 
 // HasHandler returns true iff there is a handler registered for this name

x/upgrade/keeper/keeper_test.go

@@ -278,6 +278,42 @@ func (s *KeeperTestSuite) TestLastCompletedUpgrade() {
 	require.Equal(int64(15), height)
 }
 
+// This test ensures that `GetLastDoneUpgrade` always returns the last upgrade according to the block height
+// it was executed at, rather than using an ordering based on upgrade names.
+func (s *KeeperTestSuite) TestLastCompletedUpgradeOrdering() {
+	keeper := s.app.UpgradeKeeper
+	require := s.Require()
+
+	// apply first upgrade
+	keeper.SetUpgradeHandler("test-v0.9", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
+
+	keeper.ApplyUpgrade(s.ctx, types.Plan{
+		Name:   "test-v0.9",
+		Height: 10,
+	})
+
+	name, height := keeper.GetLastCompletedUpgrade(s.ctx)
+	require.Equal("test-v0.9", name)
+	require.Equal(int64(10), height)
+
+	// apply second upgrade
+	keeper.SetUpgradeHandler("test-v0.10", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
+
+	newCtx := s.ctx.WithBlockHeight(15)
+	keeper.ApplyUpgrade(newCtx, types.Plan{
+		Name:   "test-v0.10",
+		Height: 15,
+	})
+
+	name, height = keeper.GetLastCompletedUpgrade(newCtx)
+	require.Equal("test-v0.10", name)
+	require.Equal(int64(15), height)
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/upgrade/keeper/migrations.go

@@ -0,0 +1,43 @@
+package keeper
+
+import (
+	"encoding/binary"
+
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
+)
+
+// Migrator is a struct for handling in-place store migrations.
+type Migrator struct {
+	keeper Keeper
+}
+
+// NewMigrator returns a new Migrator.
+func NewMigrator(keeper Keeper) Migrator {
+	return Migrator{keeper: keeper}
+}
+
+// Migrate1to2 migrates from version 1 to 2.
+func (m Migrator) Migrate1to2(ctx sdk.Context) error {
+	return migrateDoneUpgradeKeys(ctx, m.keeper.storeKey)
+}
+
+func migrateDoneUpgradeKeys(ctx sdk.Context, storeKey storetypes.StoreKey) error {
+	store := ctx.KVStore(storeKey)
+	oldDoneStore := prefix.NewStore(store, []byte{types.DoneByte})
+	oldDoneStoreIter := oldDoneStore.Iterator(nil, nil)
+	defer oldDoneStoreIter.Close()
+
+	for ; oldDoneStoreIter.Valid(); oldDoneStoreIter.Next() {
+		oldKey := oldDoneStoreIter.Key()
+		upgradeName := string(oldKey)
+		upgradeHeight := int64(binary.BigEndian.Uint64(oldDoneStoreIter.Value()))
+		newKey := encodeDoneKey(upgradeName, upgradeHeight)
+
+		store.Set(newKey, []byte{1})
+		oldDoneStore.Delete(oldKey)
+	}
+	return nil
+}

x/upgrade/keeper/migrations_test.go

@@ -0,0 +1,59 @@
+package keeper
+
+import (
+	"encoding/binary"
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/testutil"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
+	"github.com/stretchr/testify/require"
+)
+
+type storedUpgrade struct {
+	name   string
+	height int64
+}
+
+func encodeOldDoneKey(upgrade storedUpgrade) []byte {
+	return append([]byte{types.DoneByte}, []byte(upgrade.name)...)
+}
+
+func TestMigrateDoneUpgradeKeys(t *testing.T) {
+	upgradeKey := sdk.NewKVStoreKey("upgrade")
+	ctx := testutil.DefaultContext(upgradeKey, sdk.NewTransientStoreKey("transient_test"))
+	store := ctx.KVStore(upgradeKey)
+
+	testCases := []struct {
+		name     string
+		upgrades []storedUpgrade
+	}{
+		{
+			name: "valid upgrades",
+			upgrades: []storedUpgrade{
+				{name: "some-other-upgrade", height: 1},
+				{name: "test02", height: 2},
+				{name: "test01", height: 3},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		for _, upgrade := range tc.upgrades {
+			bz := make([]byte, 8)
+			binary.BigEndian.PutUint64(bz, uint64(upgrade.height))
+			oldKey := encodeOldDoneKey(upgrade)
+			store.Set(oldKey, bz)
+		}
+
+		err := migrateDoneUpgradeKeys(ctx, upgradeKey)
+		require.NoError(t, err)
+
+		for _, upgrade := range tc.upgrades {
+			newKey := encodeDoneKey(upgrade.name, upgrade.height)
+			oldKey := encodeOldDoneKey(upgrade)
+			require.Nil(t, store.Get(oldKey))
+			require.Equal(t, []byte{1}, store.Get(newKey))
+		}
+	}
+}

x/upgrade/module.go

@@ -22,6 +22,10 @@ func init() {
 	types.RegisterLegacyAminoCodec(codec.NewLegacyAmino())
 }
 
+const (
+	consensusVersion uint64 = 2
+)
+
 var (
 	_ module.AppModule      = AppModule{}
 	_ module.AppModuleBasic = AppModuleBasic{}
@@ -95,6 +99,9 @@ func (am AppModule) LegacyQuerierHandler(legacyQuerierCdc *codec.LegacyAmino) sd
 func (am AppModule) RegisterServices(cfg module.Configurator) {
 	types.RegisterMsgServer(cfg.MsgServer(), keeper.NewMsgServerImpl(am.keeper))
 	types.RegisterQueryServer(cfg.QueryServer(), am.keeper)
+
+	m := keeper.NewMigrator(am.keeper)
+	cfg.RegisterMigration(types.ModuleName, 1, m.Migrate1to2)
 }
 
 // InitGenesis is ignored, no sense in serializing future upgrades
@@ -118,7 +125,7 @@ func (am AppModule) ExportGenesis(_ sdk.Context, cdc codec.JSONCodec) json.RawMe
 }
 
 // ConsensusVersion implements AppModule/ConsensusVersion.
-func (AppModule) ConsensusVersion() uint64 { return 1 }
+func (AppModule) ConsensusVersion() uint64 { return consensusVersion }
 
 // BeginBlock calls the upgrade module hooks
 //