fix to overwrite empty account when we create delegator module account

initia-labs · Jun 19, 2024 · 28076ca · 28076ca
1 parent d5dd63d
commit 28076ca
Show file tree

Hide file tree

Showing 6 changed files with 50 additions and 11 deletions.
diff --git a/x/move/keeper/handler.go b/x/move/keeper/handler.go
@@ -386,10 +386,6 @@ func (k Keeper) handleExecuteResponse(
 		// increase global account number if the given account is not exists
 		if !k.authKeeper.HasAccount(ctx, addr) {
 			k.authKeeper.NextAccountNumber(ctx)
-		} else {
-			// remove account if it already exists
-			// to avoid collection's primary key conflict
-			k.authKeeper.RemoveAccount(ctx, accI)
 		}
 
 		// write or overwrite account

diff --git a/x/move/keeper/staking.go b/x/move/keeper/staking.go
@@ -128,17 +128,27 @@ func (k Keeper) hasZeroRewards(ctx context.Context, validatorAddr sdk.ValAddress
 // consequentially delegate the deposited coins to a validator.
 func (k Keeper) DelegateToValidator(ctx context.Context, valAddr sdk.ValAddress, delCoins sdk.Coins) (sdk.DecCoins, error) {
 	delegatorModuleName := types.GetDelegatorModuleName(valAddr)
-	macc := k.authKeeper.GetModuleAccount(ctx, delegatorModuleName)
+	delModuleAddr := authtypes.NewModuleAddress(delegatorModuleName)
 
-	// register module account if not registered
-	if macc == nil {
-		macc = authtypes.NewEmptyModuleAccount(delegatorModuleName)
+	if macc := k.authKeeper.GetAccount(ctx, delModuleAddr); macc != nil {
+		if _, ok := macc.(sdk.ModuleAccountI); !ok {
+			if !types.IsEmptyAccount(macc) {
+				return sdk.NewDecCoins(), types.ErrAddressAlreadyTaken.Wrapf("module account %s is already taken", delModuleAddr.String())
+			}
+
+			// overwrite empty account with module account
+			newAcc := authtypes.NewEmptyModuleAccount(delegatorModuleName)
+			newAcc.SetAccountNumber(macc.GetAccountNumber())
+
+			k.authKeeper.SetModuleAccount(ctx, newAcc)
+		}
+	} else {
+		// register module account if not registered
+		macc := authtypes.NewEmptyModuleAccount(delegatorModuleName)
 		maccI := (k.authKeeper.NewAccount(ctx, macc)).(sdk.ModuleAccountI) // set the account number
 		k.authKeeper.SetModuleAccount(ctx, maccI)
 	}
 
-	delModuleAddr := macc.GetAddress()
-
 	// send staking coin move module to validator module account
 	// delegated coins are burned, so we should mint coins to module account
 	err := k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.MoveStakingModuleName, delModuleAddr, delCoins)

diff --git a/x/move/keeper/staking_test.go b/x/move/keeper/staking_test.go
@@ -100,6 +100,28 @@ func TestDelegateToValidator(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestDelegateToValidator_OverwriteEmptyAccount(t *testing.T) {
+	ctx, input := createDefaultTestInput(t)
+	valAddr := createValidatorWithBalance(t, ctx, input, 100_000_000, 100_000)
+
+	input.Faucet.Fund(ctx, types.MoveStakingModuleAddress, sdk.NewCoin(bondDenom, math.NewInt(100_000_000)))
+	moveBalance := input.BankKeeper.GetBalance(ctx, types.MoveStakingModuleAddress, bondDenom).Amount.Uint64()
+	require.Equal(t, uint64(100_000_000), moveBalance)
+
+	// fund to delegator module account to create empty account
+	delegatorModuleAddr := types.GetDelegatorModuleAddress(valAddr)
+	input.Faucet.Fund(ctx, delegatorModuleAddr, sdk.NewCoin(bondDenom, math.NewInt(100_000_000)))
+
+	_, err := input.MoveKeeper.DelegateToValidator(ctx, valAddr, sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100))))
+	require.NoError(t, err)
+
+	moveBalance = input.BankKeeper.GetBalance(ctx, types.MoveStakingModuleAddress, bondDenom).Amount.Uint64()
+	require.Equal(t, uint64(99_999_900), moveBalance)
+
+	_, err = input.MoveKeeper.DelegateToValidator(ctx, valAddr, sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100_000_000_000))))
+	require.Error(t, err)
+}
+
 func TestAmountToShare(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 	valAddr := createValidatorWithBalance(t, ctx, input, 100_000_000, 100_000)

diff --git a/x/move/types/auth.go b/x/move/types/auth.go
@@ -38,3 +38,12 @@ func NewTableAccountWithAddress(addr sdk.AccAddress) *TableAccount {
 func (ma TableAccount) SetPubKey(pubKey cryptotypes.PubKey) error {
 	return fmt.Errorf("not supported for table accounts")
 }
+
+// IsEmptyAccount checks if the account is empty.
+func IsEmptyAccount(account sdk.AccountI) bool {
+	_, isModuleAccount := account.(sdk.ModuleAccountI)
+	_, isObjectAccount := account.(*ObjectAccount)
+	_, isTableAccount := account.(*TableAccount)
+
+	return !isModuleAccount && !isObjectAccount && !isTableAccount && account.GetPubKey() == nil
+}
diff --git a/x/move/types/errors.go b/x/move/types/errors.go
@@ -44,4 +44,7 @@ var (
 
 	// ErrNotSupportedStargateQuery error raised when the stargate request is not supported or accepted
 	ErrNotSupportedStargateQuery = errorsmod.Register(ModuleName, 14, "not supported stargate query")
+
+	// ErrAddressAlreadyTaken error raised when the address is already taken
+	ErrAddressAlreadyTaken = errorsmod.Register(ModuleName, 15, "address already taken")
 )
diff --git a/x/move/types/expected_keeper.go b/x/move/types/expected_keeper.go
@@ -20,7 +20,6 @@ type AccountKeeper interface {
 	GetAccount(ctx context.Context, addr sdk.AccAddress) sdk.AccountI
 	SetAccount(ctx context.Context, acc sdk.AccountI)
 	HasAccount(ctx context.Context, addr sdk.AccAddress) bool
-	RemoveAccount(ctx context.Context, acc sdk.AccountI)
 
 	GetModuleAddress(name string) sdk.AccAddress
 	GetModuleAccount(ctx context.Context, moduleName string) sdk.ModuleAccountI