modify simapp

simapp/app.go

@@ -280,30 +280,6 @@ func NewSimApp(
 	bApp.SetParamStore(app.ConsensusParamsKeeper.ParamsStore)
 
 	// add keepers
-	accountsKeeper, err := accounts.NewKeeper(
-		appCodec,
-		runtime.NewEnvironment(runtime.NewKVStoreService(keys[accounts.StoreKey]), logger),
-		signingCtx.AddressCodec(),
-		appCodec,
-		app.MsgServiceRouter(),
-		app.GRPCQueryRouter(),
-		appCodec.InterfaceRegistry(),
-		// TESTING: do not add
-		accountstd.AddAccount("counter", counter.NewAccount),
-		accountstd.AddAccount("aa_minimal", account_abstraction.NewMinimalAbstractedAccount),
-		// Lockup account
-		accountstd.AddAccount(lockup.CONTINUOUS_LOCKING_ACCOUNT, lockup.NewContinuousLockingAccount),
-		accountstd.AddAccount(lockup.PERIODIC_LOCKING_ACCOUNT, lockup.NewPeriodicLockingAccount),
-		accountstd.AddAccount(lockup.DELAYED_LOCKING_ACCOUNT, lockup.NewDelayedLockingAccount),
-		accountstd.AddAccount(lockup.PERMANENT_LOCKING_ACCOUNT, lockup.NewPermanentLockingAccount),
-		// PRODUCTION: add
-		baseaccount.NewAccount("base", txConfig.SignModeHandler()),
-	)
-	if err != nil {
-		panic(err)
-	}
-	app.AccountsKeeper = accountsKeeper
-
 	app.AuthKeeper = authkeeper.NewAccountKeeper(runtime.NewEnvironment(runtime.NewKVStoreService(keys[authtypes.StoreKey]), logger), appCodec, authtypes.ProtoBaseAccount, maccPerms, signingCtx.AddressCodec(), sdk.Bech32MainPrefix, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
 	app.BankKeeper = bankkeeper.NewBaseKeeper(
@@ -354,6 +330,30 @@ func NewSimApp(
 		stakingtypes.NewMultiStakingHooks(app.DistrKeeper.Hooks(), app.SlashingKeeper.Hooks()),
 	)
 
+	accountsKeeper, err := accounts.NewKeeper(
+		appCodec,
+		runtime.NewEnvironment(runtime.NewKVStoreService(keys[accounts.StoreKey]), logger),
+		signingCtx.AddressCodec(),
+		appCodec,
+		app.MsgServiceRouter(),
+		app.GRPCQueryRouter(),
+		appCodec.InterfaceRegistry(),
+		// TESTING: do not add
+		accountstd.AddAccount("counter", counter.NewAccount),
+		accountstd.AddAccount("aa_minimal", account_abstraction.NewMinimalAbstractedAccount),
+		// Lockup account
+		lockup.NewContinuousLockingAccount(app.StakingKeeper, app.BankKeeper),
+		lockup.NewDelayedLockingAccount(app.StakingKeeper, app.BankKeeper),
+		lockup.NewPeriodicLockingAccount(app.StakingKeeper, app.BankKeeper),
+		lockup.NewPermanentLockingAccount(app.StakingKeeper, app.BankKeeper),
+		// PRODUCTION: add
+		baseaccount.NewAccount("base", txConfig.SignModeHandler()),
+	)
+	if err != nil {
+		panic(err)
+	}
+	app.AccountsKeeper = accountsKeeper
+
 	app.CircuitKeeper = circuitkeeper.NewKeeper(runtime.NewEnvironment(runtime.NewKVStoreService(keys[circuittypes.StoreKey]), logger), appCodec, authtypes.NewModuleAddress(govtypes.ModuleName).String(), app.AuthKeeper.AddressCodec())
 	app.BaseApp.SetCircuitBreaker(&app.CircuitKeeper)
 

x/accounts/defaults/lockup/continuous_locking_account.go

@@ -86,6 +86,12 @@ func (cva *ContinuousLockingAccount) WithdrawUnlockedCoins(ctx context.Context,
 	return cva.BaseLockup.WithdrawUnlockedCoins(ctx, msg, cva.GetLockedCoinsWithDenoms)
 }
 
+func (cva *ContinuousLockingAccount) ClawbackFunds(ctx context.Context, msg *types.MsgClawback) (
+	*types.MsgClawbackResponse, error,
+) {
+	return cva.BaseLockup.ClawbackFunds(ctx, msg, cva.GetLockedCoinsWithDenoms)
+}
+
 // GetLockCoinsInfo returns the total number of unlocked and locked coins.
 func (cva ContinuousLockingAccount) GetLockCoinsInfo(ctx context.Context, blockTime time.Time) (unlockedCoins, lockedCoins sdk.Coins, err error) {
 	unlockedCoins = sdk.Coins{}
@@ -220,6 +226,7 @@ func (cva ContinuousLockingAccount) RegisterExecuteHandlers(builder *accountstd.
 	accountstd.RegisterExecuteHandler(builder, cva.Undelegate)
 	accountstd.RegisterExecuteHandler(builder, cva.SendCoins)
 	accountstd.RegisterExecuteHandler(builder, cva.WithdrawUnlockedCoins)
+	accountstd.RegisterExecuteHandler(builder, cva.ClawbackFunds)
 }
 
 func (cva ContinuousLockingAccount) RegisterQueryHandlers(builder *accountstd.QueryBuilder) {

x/accounts/defaults/lockup/delayed_locking_account.go

@@ -63,6 +63,12 @@ func (dva *DelayedLockingAccount) WithdrawUnlockedCoins(ctx context.Context, msg
 	return dva.BaseLockup.WithdrawUnlockedCoins(ctx, msg, dva.GetLockedCoinsWithDenoms)
 }
 
+func (dva *DelayedLockingAccount) ClawbackFunds(ctx context.Context, msg *types.MsgClawback) (
+	*types.MsgClawbackResponse, error,
+) {
+	return dva.BaseLockup.ClawbackFunds(ctx, msg, dva.GetLockedCoinsWithDenoms)
+}
+
 // GetLockCoinsInfo returns the total number of unlocked and locked coins.
 func (dva DelayedLockingAccount) GetLockCoinsInfo(ctx context.Context, blockTime time.Time) (sdk.Coins, sdk.Coins, error) {
 	endTime, err := dva.EndTime.Get(ctx)
@@ -153,6 +159,7 @@ func (dva DelayedLockingAccount) RegisterExecuteHandlers(builder *accountstd.Exe
 	accountstd.RegisterExecuteHandler(builder, dva.Undelegate)
 	accountstd.RegisterExecuteHandler(builder, dva.SendCoins)
 	accountstd.RegisterExecuteHandler(builder, dva.WithdrawUnlockedCoins)
+	accountstd.RegisterExecuteHandler(builder, dva.ClawbackFunds)
 }
 
 func (dva DelayedLockingAccount) RegisterQueryHandlers(builder *accountstd.QueryBuilder) {

x/accounts/defaults/lockup/lockup.go

@@ -405,6 +405,9 @@ func (bva *BaseLockup) ClawbackFunds(
 
 		}
 		clawbackTokens = append(clawbackTokens, sdk.NewCoin(denom, lockedAmt))
+
+		// clear the lock token tracking
+		bva.OriginalLocking.Remove(ctx, denom)
 	}
 	if len(clawbackTokens) == 0 {
 		return nil, fmt.Errorf("no tokens available for clawback")

x/accounts/defaults/lockup/periodic_locking_account.go

@@ -135,6 +135,12 @@ func (pva *PeriodicLockingAccount) WithdrawUnlockedCoins(ctx context.Context, ms
 	return pva.BaseLockup.WithdrawUnlockedCoins(ctx, msg, pva.GetLockedCoinsWithDenoms)
 }
 
+func (pva *PeriodicLockingAccount) ClawbackFunds(ctx context.Context, msg *types.MsgClawback) (
+	*types.MsgClawbackResponse, error,
+) {
+	return pva.BaseLockup.ClawbackFunds(ctx, msg, pva.GetLockedCoinsWithDenoms)
+}
+
 // IterateSendEnabledEntries iterates over all the SendEnabled entries.
 func (pva PeriodicLockingAccount) IteratePeriods(
 	ctx context.Context,
@@ -339,6 +345,7 @@ func (pva PeriodicLockingAccount) RegisterExecuteHandlers(builder *accountstd.Ex
 	accountstd.RegisterExecuteHandler(builder, pva.Undelegate)
 	accountstd.RegisterExecuteHandler(builder, pva.SendCoins)
 	accountstd.RegisterExecuteHandler(builder, pva.WithdrawUnlockedCoins)
+	accountstd.RegisterExecuteHandler(builder, pva.ClawbackFunds)
 }
 
 func (pva PeriodicLockingAccount) RegisterQueryHandlers(builder *accountstd.QueryBuilder) {

x/accounts/defaults/lockup/permanent_locking_account.go

@@ -44,7 +44,7 @@ func (plva PermanentLockingAccount) Init(ctx context.Context, msg *types.MsgInit
 
 // GetlockedCoinsWithDenoms returns the total number of locked coins. If no coins are
 // locked, nil is returned.
-func (plva PermanentLockingAccount) GetlockedCoinsWithDenoms(ctx context.Context, blockTime time.Time, denoms ...string) (sdk.Coins, error) {
+func (plva PermanentLockingAccount) GetLockedCoinsWithDenoms(ctx context.Context, blockTime time.Time, denoms ...string) (sdk.Coins, error) {
 	vestingCoins := sdk.Coins{}
 	for _, denom := range denoms {
 		originalVestingAmt, err := plva.OriginalLocking.Get(ctx, denom)
@@ -59,7 +59,7 @@ func (plva PermanentLockingAccount) GetlockedCoinsWithDenoms(ctx context.Context
 func (plva *PermanentLockingAccount) Delegate(ctx context.Context, msg *types.MsgDelegate) (
 	*types.MsgExecuteMessagesResponse, error,
 ) {
-	return plva.BaseLockup.Delegate(ctx, msg, plva.GetlockedCoinsWithDenoms)
+	return plva.BaseLockup.Delegate(ctx, msg, plva.GetLockedCoinsWithDenoms)
 }
 
 func (plva *PermanentLockingAccount) Undelegate(ctx context.Context, msg *types.MsgUndelegate) (
@@ -71,7 +71,13 @@ func (plva *PermanentLockingAccount) Undelegate(ctx context.Context, msg *types.
 func (plva *PermanentLockingAccount) SendCoins(ctx context.Context, msg *types.MsgSend) (
 	*types.MsgExecuteMessagesResponse, error,
 ) {
-	return plva.BaseLockup.SendCoins(ctx, msg, plva.GetlockedCoinsWithDenoms)
+	return plva.BaseLockup.SendCoins(ctx, msg, plva.GetLockedCoinsWithDenoms)
+}
+
+func (plva *PermanentLockingAccount) ClawbackFunds(ctx context.Context, msg *types.MsgClawback) (
+	*types.MsgClawbackResponse, error,
+) {
+	return plva.BaseLockup.ClawbackFunds(ctx, msg, plva.GetLockedCoinsWithDenoms)
 }
 
 func (plva PermanentLockingAccount) QueryLockupAccountInfo(ctx context.Context, req *types.QueryLockupAccountInfoRequest) (
@@ -104,6 +110,7 @@ func (plva PermanentLockingAccount) RegisterExecuteHandlers(builder *accountstd.
 	accountstd.RegisterExecuteHandler(builder, plva.Delegate)
 	accountstd.RegisterExecuteHandler(builder, plva.Undelegate)
 	accountstd.RegisterExecuteHandler(builder, plva.SendCoins)
+	accountstd.RegisterExecuteHandler(builder, plva.ClawbackFunds)
 }
 
 func (plva PermanentLockingAccount) RegisterQueryHandlers(builder *accountstd.QueryBuilder) {