Prune vesting accounts balances (CosmWasm#1003)

* Prevent out of gas

* Prune vesting account denoms only

* Fix test state

* Move account exists error  up again

* Review feedback: better naming

x/wasm/keeper/keeper.go

@@ -22,7 +22,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/types/address"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
-	vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
+	vestingexported "github.com/cosmos/cosmos-sdk/x/auth/vesting/exported"
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
 	"github.com/tendermint/tendermint/libs/log"
 
@@ -57,11 +57,13 @@ type CoinTransferrer interface {
 	TransferCoins(ctx sdk.Context, fromAddr sdk.AccAddress, toAddr sdk.AccAddress, amt sdk.Coins) error
 }
 
-// CoinPruner handles the balances for accounts that are pruned on contract instantiate.
+// AccountPruner handles the balances and data cleanup for accounts that are pruned on contract instantiate.
 // This is an extension point to attach custom logic
-type CoinPruner interface {
-	// PruneBalances handle balances for given address
-	PruneBalances(ctx sdk.Context, contractAddress sdk.AccAddress) error
+type AccountPruner interface {
+	// CleanupExistingAccount handles the cleanup process for balances and data of the given account. The persisted account
+	// type is already reset to base account at this stage.
+	// The method returns true when the account address can be reused. Unsupported account types are rejected by returning false
+	CleanupExistingAccount(ctx sdk.Context, existingAccount authtypes.AccountI) (handled bool, err error)
 }
 
 // WasmVMResponseHandler is an extension point to handles the response data returned by a contract call.
@@ -82,19 +84,6 @@ var defaultAcceptedAccountTypes = map[reflect.Type]struct{}{
 	reflect.TypeOf(&authtypes.BaseAccount{}): {},
 }
 
-// list of account types that are replaced with base accounts. Chains importing wasmd
-// can overwrite this list with the WithPruneAccountTypesOnContractInstantiation option.
-//
-// contains vesting account types that can be created post genesis
-var defaultPruneAccountTypes = map[reflect.Type]struct{}{
-	reflect.TypeOf(&vestingtypes.DelayedVestingAccount{}):    {},
-	reflect.TypeOf(&vestingtypes.ContinuousVestingAccount{}): {},
-	// intentionally not added: genesis account types
-	// reflect.TypeOf(&vestingtypes.BaseVestingAccount{}):       {},
-	// reflect.TypeOf(&vestingtypes.PeriodicVestingAccount{}):   {},
-	// reflect.TypeOf(&vestingtypes.PermanentLockedAccount{}):   {},
-}
-
 // Keeper will have a reference to Wasmer with it's own data directory.
 type Keeper struct {
 	storeKey              sdk.StoreKey
@@ -113,8 +102,7 @@ type Keeper struct {
 	gasRegister          GasRegister
 	maxQueryStackSize    uint32
 	acceptedAccountTypes map[reflect.Type]struct{}
-	pruneAccountTypes    map[reflect.Type]struct{}
-	coinPruner           CoinPruner
+	accountPruner        AccountPruner
 }
 
 // NewKeeper creates a new contract Keeper instance
@@ -153,7 +141,7 @@ func NewKeeper(
 		wasmVM:               wasmer,
 		accountKeeper:        accountKeeper,
 		bank:                 NewBankCoinTransferrer(bankKeeper),
-		coinPruner:           NewCoinBurner(bankKeeper),
+		accountPruner:        NewVestingCoinBurner(bankKeeper),
 		portKeeper:           portKeeper,
 		capabilityKeeper:     capabilityKeeper,
 		messenger:            NewDefaultMessageHandler(router, channelKeeper, capabilityKeeper, bankKeeper, cdc, portSource),
@@ -162,7 +150,6 @@ func NewKeeper(
 		gasRegister:          NewDefaultWasmGasRegister(),
 		maxQueryStackSize:    types.DefaultMaxQueryStackSize,
 		acceptedAccountTypes: defaultAcceptedAccountTypes,
-		pruneAccountTypes:    defaultPruneAccountTypes,
 	}
 	keeper.wasmVMQueryHandler = DefaultQueryPlugins(bankKeeper, stakingKeeper, distKeeper, channelKeeper, queryRouter, keeper)
 	for _, o := range opts {
@@ -321,17 +308,18 @@ func (k Keeper) instantiate(ctx sdk.Context, codeID uint64, creator, admin sdk.A
 		if _, accept := k.acceptedAccountTypes[reflect.TypeOf(existingAcct)]; accept {
 			// keep account and balance as it is
 			k.Logger(ctx).Info("instantiate contract with existing account", "address", contractAddress.String())
-		} else if _, clear := k.pruneAccountTypes[reflect.TypeOf(existingAcct)]; clear {
-			k.Logger(ctx).Info("pruning existing account for contract instantiation", "address", contractAddress.String())
+		} else {
 			// consider an account in the wasmd namespace spam and overwrite it.
+			k.Logger(ctx).Info("pruning existing account for contract instantiation", "address", contractAddress.String())
 			contractAccount := k.accountKeeper.NewAccountWithAddress(ctx, contractAddress)
 			k.accountKeeper.SetAccount(ctx, contractAccount)
 			// also handle balance to not open cases where these accounts are abused and become liquid
-			if err := k.coinPruner.PruneBalances(ctx, contractAddress); err != nil {
-				return nil, nil, err
+			switch handled, err := k.accountPruner.CleanupExistingAccount(ctx, existingAcct); {
+			case err != nil:
+				return nil, nil, sdkerrors.Wrap(err, "prune balance")
+			case !handled:
+				return nil, nil, types.ErrAccountExists.Wrap("address is claimed by external account")
 			}
-		} else { // unknown account type
-			return nil, nil, types.ErrAccountExists.Wrapf("refusing to overwrite special account type:: %T", existingAcct)
 		}
 	} else {
 		// create an empty account (so we don't have issues later)
@@ -1161,32 +1149,41 @@ func (c BankCoinTransferrer) TransferCoins(parentCtx sdk.Context, fromAddr sdk.A
 	return nil
 }
 
-var _ CoinPruner = CoinBurner{}
+var _ AccountPruner = VestingCoinBurner{}
 
-// CoinBurner default implementation for CoinPruner to burn the coins
-type CoinBurner struct {
+// VestingCoinBurner default implementation for AccountPruner to burn the coins
+type VestingCoinBurner struct {
 	bank types.BankKeeper
 }
 
-// NewCoinBurner constructor
-func NewCoinBurner(bank types.BankKeeper) CoinBurner {
+// NewVestingCoinBurner constructor
+func NewVestingCoinBurner(bank types.BankKeeper) VestingCoinBurner {
 	if bank == nil {
 		panic("bank keeper must not be nil")
 	}
-	return CoinBurner{bank: bank}
+	return VestingCoinBurner{bank: bank}
 }
 
-// PruneBalances burns all coins owned by the account.
-func (b CoinBurner) PruneBalances(ctx sdk.Context, address sdk.AccAddress) error {
-	if amt := b.bank.GetAllBalances(ctx, address); !amt.IsZero() {
-		if err := b.bank.SendCoinsFromAccountToModule(ctx, address, types.ModuleName, amt); err != nil {
-			return sdkerrors.Wrap(err, "prune account balance")
-		}
-		if err := b.bank.BurnCoins(ctx, types.ModuleName, amt); err != nil {
-			return sdkerrors.Wrap(err, "burn account balance")
-		}
+// CleanupExistingAccount accepts only vesting account types to burns all their original vesting coin balances.
+// Other account types will be rejected and returned as unhandled.
+func (b VestingCoinBurner) CleanupExistingAccount(ctx sdk.Context, existingAcc authtypes.AccountI) (handled bool, err error) {
+	v, ok := existingAcc.(vestingexported.VestingAccount)
+	if !ok {
+		return false, nil
 	}
-	return nil
+
+	ctx = ctx.WithGasMeter(sdk.NewInfiniteGasMeter())
+	coinsToBurn := sdk.NewCoins()
+	for _, orig := range v.GetOriginalVesting() { // focus on the coin denoms that were setup originally; getAllBalances has some issues
+		coinsToBurn = append(coinsToBurn, b.bank.GetBalance(ctx, existingAcc.GetAddress(), orig.Denom))
+	}
+	if err := b.bank.SendCoinsFromAccountToModule(ctx, existingAcc.GetAddress(), types.ModuleName, coinsToBurn); err != nil {
+		return false, sdkerrors.Wrap(err, "prune account balance")
+	}
+	if err := b.bank.BurnCoins(ctx, types.ModuleName, coinsToBurn); err != nil {
+		return false, sdkerrors.Wrap(err, "burn account balance")
+	}
+	return true, nil
 }
 
 type msgDispatcher interface {

x/wasm/keeper/keeper_test.go

@@ -12,16 +12,16 @@ import (
 	"testing"
 	"time"
 
-	"github.com/cosmos/cosmos-sdk/baseapp"
-
 	wasmvm "github.com/CosmWasm/wasmvm"
 	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
+	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/crypto/keys/ed25519"
 	stypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	"github.com/cosmos/cosmos-sdk/x/auth/vesting"
 	vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 	distributiontypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
@@ -645,7 +645,7 @@ func TestInstantiateWithAccounts(t *testing.T) {
 			deposit:    sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(1_000))),
 			expBalance: sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(1_000))),
 		},
-		"prune listed DelayedVestingAccount gets overwritten": {
+		"prunable DelayedVestingAccount gets overwritten": {
 			account: vestingtypes.NewDelayedVestingAccount(
 				authtypes.NewBaseAccount(contractAddr, nil, 0, 0),
 				sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(1_000))), time.Now().Add(30*time.Hour).Unix()),
@@ -654,7 +654,7 @@ func TestInstantiateWithAccounts(t *testing.T) {
 			expAccount:  authtypes.NewBaseAccount(contractAddr, nil, lastAccountNumber+2, 0), // +1 for next seq, +1 for spec.account created
 			expBalance:  sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(1))),
 		},
-		"prune listed ContinuousVestingAccount gets overwritten": {
+		"prunable ContinuousVestingAccount gets overwritten": {
 			account: vestingtypes.NewContinuousVestingAccount(
 				authtypes.NewBaseAccount(contractAddr, nil, 0, 0),
 				sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(1_000))), time.Now().Add(time.Hour).Unix(), time.Now().Add(2*time.Hour).Unix()),
@@ -663,14 +663,14 @@ func TestInstantiateWithAccounts(t *testing.T) {
 			expAccount:  authtypes.NewBaseAccount(contractAddr, nil, lastAccountNumber+2, 0), // +1 for next seq, +1 for spec.account created
 			expBalance:  sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(1))),
 		},
-		"prune listed account without balance gets overwritten": {
+		"prunable account without balance gets overwritten": {
 			account: vestingtypes.NewContinuousVestingAccount(
 				authtypes.NewBaseAccount(contractAddr, nil, 0, 0),
 				sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(0))), time.Now().Add(time.Hour).Unix(), time.Now().Add(2*time.Hour).Unix()),
 			expAccount: authtypes.NewBaseAccount(contractAddr, nil, lastAccountNumber+2, 0), // +1 for next seq, +1 for spec.account created
 			expBalance: sdk.NewCoins(),
 		},
-		"unknown account type creates error": {
+		"unknown account type is rejected with error": {
 			account: authtypes.NewModuleAccount(
 				authtypes.NewBaseAccount(contractAddr, nil, 0, 0),
 				"testing",
@@ -2274,3 +2274,66 @@ func TestAppendToContractHistory(t *testing.T) {
 	gotHistory := keepers.WasmKeeper.GetContractHistory(ctx, contractAddr)
 	assert.Equal(t, orderedEntries, gotHistory)
 }
+
+func TestCoinBurnerPruneBalances(t *testing.T) {
+	parentCtx, keepers := CreateTestInput(t, false, AvailableCapabilities)
+	amts := sdk.NewCoins(sdk.NewInt64Coin("denom", 100))
+	senderAddr := keepers.Faucet.NewFundedRandomAccount(parentCtx, amts...)
+
+	// create vesting account
+	var vestingAddr sdk.AccAddress = rand.Bytes(types.ContractAddrLen)
+	msgCreateVestingAccount := vestingtypes.NewMsgCreateVestingAccount(senderAddr, vestingAddr, amts, time.Now().Add(time.Minute).Unix(), false)
+	_, err := vesting.NewMsgServerImpl(keepers.AccountKeeper, keepers.BankKeeper).CreateVestingAccount(sdk.WrapSDKContext(parentCtx), msgCreateVestingAccount)
+	require.NoError(t, err)
+	myVestingAccount := keepers.AccountKeeper.GetAccount(parentCtx, vestingAddr)
+	require.NotNil(t, myVestingAccount)
+
+	specs := map[string]struct {
+		setupAcc    func(t *testing.T, ctx sdk.Context) authtypes.AccountI
+		expBalances sdk.Coins
+		expHandled  bool
+		expErr      *sdkerrors.Error
+	}{
+		"vesting account - all removed": {
+			setupAcc:    func(t *testing.T, ctx sdk.Context) authtypes.AccountI { return myVestingAccount },
+			expBalances: sdk.NewCoins(),
+			expHandled:  true,
+		},
+		"vesting account with other tokens - only original denoms removed": {
+			setupAcc: func(t *testing.T, ctx sdk.Context) authtypes.AccountI {
+				keepers.Faucet.Fund(ctx, vestingAddr, sdk.NewCoin("other", sdk.NewInt(2)))
+				return myVestingAccount
+			},
+			expBalances: sdk.NewCoins(sdk.NewCoin("other", sdk.NewInt(2))),
+			expHandled:  true,
+		},
+		"non vesting account - not handled": {
+			setupAcc: func(t *testing.T, ctx sdk.Context) authtypes.AccountI {
+				return &authtypes.BaseAccount{Address: myVestingAccount.GetAddress().String()}
+			},
+			expBalances: sdk.NewCoins(sdk.NewCoin("denom", sdk.NewInt(100))),
+			expHandled:  false,
+		},
+	}
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			ctx, _ := parentCtx.CacheContext()
+			existingAccount := spec.setupAcc(t, ctx)
+			// overwrite account in store as in keeper before calling prune
+			keepers.AccountKeeper.SetAccount(ctx, keepers.AccountKeeper.NewAccountWithAddress(ctx, vestingAddr))
+
+			// when
+			noGasCtx := ctx.WithGasMeter(sdk.NewGasMeter(0)) // should not use callers gas
+			gotHandled, gotErr := NewVestingCoinBurner(keepers.BankKeeper).CleanupExistingAccount(noGasCtx, existingAccount)
+			// then
+			if spec.expErr != nil {
+				require.ErrorIs(t, gotErr, spec.expErr)
+				return
+			}
+			require.NoError(t, gotErr)
+			assert.Equal(t, spec.expBalances, keepers.BankKeeper.GetAllBalances(ctx, vestingAddr))
+			assert.Equal(t, spec.expHandled, gotHandled)
+			// and no out of gas panic
+		})
+	}
+}

x/wasm/keeper/options.go

@@ -99,14 +99,14 @@ func WithCoinTransferrer(x CoinTransferrer) Option {
 	})
 }
 
-// WithCoinPruner is an optional constructor parameter to set a custom type that handles balances
+// WithAccountPruner is an optional constructor parameter to set a custom type that handles balances and data cleanup
 // for accounts pruned on contract instantiate
-func WithCoinPruner(x CoinPruner) Option {
+func WithAccountPruner(x AccountPruner) Option {
 	if x == nil {
 		panic("must not be nil")
 	}
 	return optsFn(func(k *Keeper) {
-		k.coinPruner = x
+		k.accountPruner = x
 	})
 }
 
@@ -154,19 +154,6 @@ func WithAcceptedAccountTypesOnContractInstantiation(accts ...authtypes.AccountI
 	})
 }
 
-// WithPruneAccountTypesOnContractInstantiation sets the account types that should be cleared. Account types of this list
-// will be overwritten with the BaseAccount type and their balance burned, when they exist for an address on contract
-// instantiation.
-//
-// Values should be references and contain the `*vestingtypes.DelayedVestingAccount`, `*vestingtypes.ContinuousVestingAccount`
-// as post genesis account types with an open address range.
-func WithPruneAccountTypesOnContractInstantiation(accts ...authtypes.AccountI) Option {
-	m := asTypeMap(accts)
-	return optsFn(func(k *Keeper) {
-		k.pruneAccountTypes = m
-	})
-}
-
 func asTypeMap(accts []authtypes.AccountI) map[reflect.Type]struct{} {
 	m := make(map[reflect.Type]struct{}, len(accts))
 	for _, a := range accts {

x/wasm/keeper/options_test.go

@@ -95,20 +95,10 @@ func TestConstructorOptions(t *testing.T) {
 				assert.Equal(t, exp, k.acceptedAccountTypes)
 			},
 		},
-		"prune account types": {
-			srcOpt: WithPruneAccountTypesOnContractInstantiation(&authtypes.BaseAccount{}, &vestingtypes.ContinuousVestingAccount{}),
+		"account pruner": {
+			srcOpt: WithAccountPruner(VestingCoinBurner{}),
 			verify: func(t *testing.T, k Keeper) {
-				exp := map[reflect.Type]struct{}{
-					reflect.TypeOf(&authtypes.BaseAccount{}):                 {},
-					reflect.TypeOf(&vestingtypes.ContinuousVestingAccount{}): {},
-				}
-				assert.Equal(t, exp, k.pruneAccountTypes)
-			},
-		},
-		"coin pruner": {
-			srcOpt: WithCoinPruner(CoinBurner{}),
-			verify: func(t *testing.T, k Keeper) {
-				assert.Equal(t, CoinBurner{}, k.coinPruner)
+				assert.Equal(t, VestingCoinBurner{}, k.accountPruner)
 			},
 		},
 	}