Merge pull request #1954 from lavanet/fix-vault-stake-bug

fix: vault double stake bug

app/app.go

@@ -171,6 +171,7 @@ var Upgrades = []upgrades.Upgrade{
 	upgrades.Upgrade_4_2_0,
 	upgrades.Upgrade_5_0_0,
 	upgrades.Upgrade_5_1_0,
+	upgrades.Upgrade_5_2_0,
 }
 
 // this line is used by starport scaffolding # stargate/wasm/app/enabledProposals
@@ -1086,7 +1087,7 @@ func initParamsKeeper(appCodec codec.BinaryCodec, legacyAmino *codec.LegacyAmino
 	paramsKeeper.Subspace(stakingtypes.ModuleName)
 	paramsKeeper.Subspace(distrtypes.ModuleName)
 	paramsKeeper.Subspace(slashingtypes.ModuleName)
-	paramsKeeper.Subspace(govtypes.ModuleName).WithKeyTable(v1.ParamKeyTable()) //nolint:staticcheck
+	paramsKeeper.Subspace(govtypes.ModuleName).WithKeyTable(v1.ParamKeyTable()) //nolint
 	paramsKeeper.Subspace(crisistypes.ModuleName)
 	paramsKeeper.Subspace(ibctransfertypes.ModuleName)
 	paramsKeeper.Subspace(ibcexported.ModuleName)

app/upgrades/empty_upgrades.go

@@ -75,3 +75,9 @@ var Upgrade_5_1_0 = Upgrade{
 	CreateUpgradeHandler: defaultUpgradeHandler,
 	StoreUpgrades:        store.StoreUpgrades{},
 }
+
+var Upgrade_5_2_0 = Upgrade{
+	UpgradeName:          "v5.2.0",
+	CreateUpgradeHandler: defaultUpgradeHandler,
+	StoreUpgrades:        store.StoreUpgrades{},
+}

x/epochstorage/keeper/keeper.go

@@ -34,7 +34,7 @@ type (
 		stakeEntries        *collections.IndexedMap[collections.Triple[uint64, string, collections.Pair[uint64, string]], types.StakeEntry, types.EpochChainIdProviderIndexes]
 		stakeEntriesCurrent *collections.IndexedMap[collections.Pair[string, string], types.StakeEntry, types.ChainIdVaultIndexes]
 		epochHashes         collections.Map[uint64, []byte]
-		providersMetaData   collections.Map[string, types.ProviderMetadata]
+		providersMetaData   *collections.IndexedMap[string, types.ProviderMetadata, types.MetadataVaultIndexes]
 	}
 )
 
@@ -79,7 +79,10 @@ func NewKeeper(
 
 		epochHashes: collections.NewMap(sb, types.EpochHashesPrefix, "epoch_hashes", collections.Uint64Key, collections.BytesValue),
 
-		providersMetaData: collections.NewMap(sb, types.ProviderMetaDataPrefix, "provider_metadata", collections.StringKey, collcompat.ProtoValue[types.ProviderMetadata](cdc)),
+		providersMetaData: collections.NewIndexedMap(sb, types.ProviderMetaDataPrefix, "provider_metadata",
+			collections.StringKey,
+			collcompat.ProtoValue[types.ProviderMetadata](cdc),
+			types.NewMetadataVaultIndexes(sb)),
 	}
 
 	keeper.AddFixationRegistry(string(types.KeyEpochBlocks), func(ctx sdk.Context) any { return keeper.EpochBlocksRaw(ctx) })

x/epochstorage/keeper/migrations.go

@@ -124,3 +124,39 @@ func (m Migrator) SetEpochHashForMigrator(ctx sdk.Context, epoch uint64, hash []
 		panic(err)
 	}
 }
+
+// fix provider metadata to be consistent with the stake entries
+func (m Migrator) MigrateVersion8To9(ctx sdk.Context) error {
+	allMetadata, err := m.keeper.GetAllMetadata(ctx)
+	if err != nil {
+		return err
+	}
+
+	for _, metadata := range allMetadata {
+		if metadata.Provider != metadata.Vault {
+			_, err := m.keeper.GetMetadata(ctx, metadata.Vault)
+			if err == nil {
+				fmt.Println("exists for vault with different provider: ", metadata.Vault, " and provider: ", metadata.Provider, "removing metadata")
+				err := m.keeper.RemoveMetadata(ctx, metadata.Vault)
+				if err != nil {
+					return err
+				}
+				continue
+			}
+		}
+
+		chains := []string{}
+		for _, chainID := range metadata.Chains {
+			stakeEntry, found := m.keeper.GetStakeEntryCurrent(ctx, chainID, metadata.Provider)
+			if found {
+				chains = append(chains, stakeEntry.Chain)
+			} else {
+				fmt.Println("stake entry not found for provider: ", metadata.Provider, " and chain: ", chainID)
+			}
+		}
+		metadata.Chains = chains
+		m.keeper.SetMetadata(ctx, metadata)
+	}
+
+	return nil
+}

x/epochstorage/keeper/provider_metadata.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/lavanet/lava/v5/utils"
 	"github.com/lavanet/lava/v5/x/epochstorage/types"
 )
 
@@ -33,3 +34,21 @@ func (k Keeper) GetAllMetadata(ctx sdk.Context) ([]types.ProviderMetadata, error
 	}
 	return iter.Values()
 }
+
+// GetProviderMetadataByVault gets the provider metadata for a specific vault address
+func (k Keeper) GetProviderMetadataByVault(ctx sdk.Context, vault string) (val types.ProviderMetadata, found bool) {
+	pk, err := k.providersMetaData.Indexes.Index.MatchExact(ctx, vault)
+	if err != nil {
+		return types.ProviderMetadata{}, false
+	}
+
+	entry, err := k.providersMetaData.Get(ctx, pk)
+	if err != nil {
+		utils.LavaFormatError("GetProviderMetadataByVault: Get with primary key failed", err,
+			utils.LogAttr("vault", vault),
+		)
+		return types.ProviderMetadata{}, false
+	}
+
+	return entry, true
+}

x/epochstorage/module.go

@@ -135,6 +135,11 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 		// panic:ok: at start up, migration cannot proceed anyhow
 		panic(fmt.Errorf("%s: failed to register migration to v8: %w", types.ModuleName, err))
 	}
+
+	// register v8 -> v9 migration
+	if err := cfg.RegisterMigration(types.ModuleName, 8, migrator.MigrateVersion8To9); err != nil {
+		panic(fmt.Errorf("%s: failed to register migration to v9: %w", types.ModuleName, err))
+	}
 }
 
 // RegisterInvariants registers the capability module's invariants.
@@ -159,7 +164,7 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 }
 
 // ConsensusVersion implements ConsensusVersion.
-func (AppModule) ConsensusVersion() uint64 { return 8 }
+func (AppModule) ConsensusVersion() uint64 { return 9 }
 
 // BeginBlock executes all ABCI BeginBlock logic respective to the capability module.
 func (am AppModule) BeginBlock(ctx sdk.Context, _ abci.RequestBeginBlock) {

x/epochstorage/types/provider_metdata.go

@@ -0,0 +1,41 @@
+package types
+
+import (
+	fmt "fmt"
+
+	"cosmossdk.io/collections"
+	"cosmossdk.io/collections/indexes"
+	"github.com/lavanet/lava/v5/utils"
+)
+
+var MetadataVaultIndexesPrefix = collections.NewPrefix([]byte("MetadataVaultIndexes/"))
+
+// ChainIdVaultIndexes defines a secondary unique index for the keeper's stakeEntriesCurrent indexed map
+// Normally, a current stake entry can be accessed with the primary key: [chainID, address]
+// The new set of indexes, ChainIdVaultIndexes, allows accessing the stake entries with [chainID, vault]
+type MetadataVaultIndexes struct {
+	Index *indexes.Unique[string, string, ProviderMetadata]
+}
+
+func (c MetadataVaultIndexes) IndexesList() []collections.Index[string, ProviderMetadata] {
+	return []collections.Index[string, ProviderMetadata]{c.Index}
+}
+
+func NewMetadataVaultIndexes(sb *collections.SchemaBuilder) MetadataVaultIndexes {
+	return MetadataVaultIndexes{
+		Index: indexes.NewUnique(sb, MetadataVaultIndexesPrefix, "provider_metadata_by_vault",
+			collections.StringKey,
+			collections.StringKey,
+			func(pk string, entry ProviderMetadata) (string, error) {
+				if entry.Vault == "" {
+					return "",
+						utils.LavaFormatError("NewMetadataVaultIndexes: cannot create new MetadataVault index",
+							fmt.Errorf("empty vault address"),
+							utils.LogAttr("provider", entry.Provider),
+						)
+				}
+				return entry.Vault, nil
+			},
+		),
+	}
+}

x/pairing/keeper/msg_server_stake_provider_test.go

@@ -1245,3 +1245,20 @@ func TestDelegatorAfterProviderUnstakeAndStake(t *testing.T) {
 		})
 	}
 }
+
+// TestStakeProviderWithSameVault tests that a vault address cannot be used to stake for multiple providers.
+// It verifies that:
+// 1. A vault can stake for itself as a provider
+// 2. The same vault cannot stake for a different provider address
+func TestStakeProviderWithSameVault(t *testing.T) {
+	ts := newTester(t)
+	SetupForSingleProviderTests(ts, 0, 3, 0)
+
+	provider, _ := ts.AddAccount(common.PROVIDER, 1, 100*testBalance)
+
+	err := ts.StakeProvider(provider.GetVaultAddr(), provider.GetVaultAddr(), ts.Spec(SpecName(0)), 100)
+	require.NoError(t, err)
+
+	err = ts.StakeProvider(provider.GetVaultAddr(), provider.Addr.String(), ts.Spec(SpecName(0)), 100)
+	require.Error(t, err)
+}

x/pairing/keeper/staking.go

@@ -37,6 +37,16 @@ func (k Keeper) StakeNewEntry(ctx sdk.Context, validator, creator, chainID strin
 		metadata.Chains = lavaslices.AddUnique(metadata.Chains, chainID)
 	}
 
+	vaultMetadata, found := k.epochStorageKeeper.GetProviderMetadataByVault(ctx, creator)
+	if found {
+		if vaultMetadata.Provider != provider || vaultMetadata.Vault != creator {
+			return utils.LavaFormatWarning("provider metadata mismatch", fmt.Errorf("provider metadata mismatch"),
+				utils.LogAttr("provider", provider),
+				utils.LogAttr("vault", creator),
+			)
+		}
+	}
+
 	spec, err := k.specKeeper.GetExpandedSpec(ctx, specChainID)
 	if err != nil || !spec.Enabled {
 		return utils.LavaFormatWarning("spec not found or not active", err,

x/pairing/types/expected_keepers.go

@@ -55,6 +55,7 @@ type EpochstorageKeeper interface {
 	GetEpochHash(ctx sdk.Context, epoch uint64) []byte
 	GetMetadata(ctx sdk.Context, provider string) (epochstoragetypes.ProviderMetadata, error)
 	SetMetadata(ctx sdk.Context, metadata epochstoragetypes.ProviderMetadata)
+	GetProviderMetadataByVault(ctx sdk.Context, vault string) (epochstoragetypes.ProviderMetadata, bool)
 }
 
 type AccountKeeper interface {