fix(statedb): snapshot locked balance on statedb account

tests/integration/precompiles/staking/test_integration.go

@@ -22,6 +22,7 @@ import (
 	"github.com/cosmos/evm/precompiles/testutil"
 	"github.com/cosmos/evm/precompiles/testutil/contracts"
 	cosmosevmutil "github.com/cosmos/evm/testutil/constants"
+	basefactory "github.com/cosmos/evm/testutil/integration/base/factory"
 	"github.com/cosmos/evm/testutil/integration/evm/network"
 	"github.com/cosmos/evm/testutil/integration/evm/utils"
 	testutiltx "github.com/cosmos/evm/testutil/tx"
@@ -35,6 +36,8 @@ import (
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/types/query"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 )
 
@@ -428,6 +431,90 @@ func TestPrecompileIntegrationTestSuite(t *testing.T, create network.CreateEvmAp
 				})
 			})
 
+			Context("from a vesting account", func() {
+				var (
+					vestAddr     sdk.AccAddress
+					vestPriv     *ethsecp256k1.PrivKey
+					amtLocked    math.Int
+					amtSpendable math.Int
+					preBal       math.Int
+					preSupply    math.Int
+				)
+
+				BeforeEach(func() {
+					// setup vesting account to delegate from
+					vestAddr, vestPriv = testutiltx.NewAccAddressAndKey()
+					amtLocked = math.NewInt(2e18)
+					amtSpendable = math.NewInt(2e18)
+
+					funder := s.keyring.GetKey(0)
+					startTime := s.network.GetContext().BlockTime().Unix()
+					createMsg := &vestingtypes.MsgCreateVestingAccount{
+						FromAddress: funder.AccAddr.String(),
+						ToAddress:   vestAddr.String(),
+						Amount:      sdk.NewCoins(sdk.NewCoin(s.bondDenom, amtLocked)),
+						EndTime:     startTime + 365*24*3600,
+						Delayed:     false,
+					}
+					sendMsg := banktypes.NewMsgSend(
+						funder.AccAddr, vestAddr,
+						sdk.NewCoins(sdk.NewCoin(s.bondDenom, amtSpendable)),
+					)
+					_, err := s.factory.CommitCosmosTx(funder.Priv, basefactory.CosmosTxArgs{
+						Msgs: []sdk.Msg{createMsg, sendMsg},
+					})
+					Expect(err).To(BeNil(), "error while submitting vesting setup tx")
+
+					ctx := s.network.GetContext()
+					_, ok := s.network.App.GetAccountKeeper().GetAccount(ctx, vestAddr).(*vestingtypes.ContinuousVestingAccount)
+					Expect(ok).To(BeTrue(), "expected vesting account to persist after tx commit")
+					spendable := s.network.App.GetBankKeeper().SpendableCoin(ctx, vestAddr, s.bondDenom).Amount
+					Expect(spendable).To(Equal(amtSpendable), "unexpected spendable balance after vesting setup")
+
+					preBalRes, err := s.grpcHandler.GetBalanceFromBank(vestAddr, s.bondDenom)
+					Expect(err).To(BeNil(), "error while getting pre balance")
+					preBal = preBalRes.Balance.Amount
+					Expect(preBal).To(Equal(amtLocked.Add(amtSpendable)), "expected vester pre bank balance to equal OV + extra")
+
+					preSupplyRes, err := s.grpcHandler.GetTotalSupply()
+					Expect(err).To(BeNil(), "error while getting pre supply")
+					preSupply = preSupplyRes.Supply.AmountOf(s.bondDenom)
+				})
+
+				It("should preserve bank balance and total supply when delegating within spendable", func() {
+					// delegating less than spendable balance and less than
+					// locked balance
+					delAmt := big.NewInt(1e18)
+					gasPrice := big.NewInt(1e9)
+
+					callArgs.Args = []interface{}{
+						common.BytesToAddress(vestAddr), valAddr.String(), delAmt,
+					}
+					delTxArgs := txArgs
+					delTxArgs.GasPrice = gasPrice
+
+					logCheckArgs := passCheck.WithExpEvents(staking.EventTypeDelegate)
+					res, _, err := s.factory.CallContractAndCheckLogs(
+						vestPriv, delTxArgs, callArgs, logCheckArgs,
+					)
+					Expect(err).To(BeNil(), "error while calling the smart contract: %v", err)
+					Expect(s.network.NextBlock()).To(BeNil())
+
+					postBalRes, err := s.grpcHandler.GetBalanceFromBank(vestAddr, s.bondDenom)
+					Expect(err).To(BeNil(), "error while getting post balance")
+					postSupplyRes, err := s.grpcHandler.GetTotalSupply()
+					Expect(err).To(BeNil(), "error while getting post supply")
+					postBal := postBalRes.Balance.Amount
+					postSupply := postSupplyRes.Supply.AmountOf(s.bondDenom)
+
+					gasCost := new(big.Int).Mul(gasPrice, big.NewInt(res.GasUsed))
+					expBalDrop := new(big.Int).Add(delAmt, gasCost)
+					actualBalDrop := preBal.Sub(postBal).BigInt()
+					Expect(actualBalDrop).To(Equal(expBalDrop), "vesting bank balance dropped by more than delegation amount + gas")
+					Expect(postSupply).To(Equal(preSupply), "unexpected total supply after delegating from vesting account")
+				})
+			})
+
 			Context("on behalf of another account", func() {
 				It("should not delegate if delegator address is not the msg.sender", func() {
 					delegator := s.keyring.GetKey(0)

tests/integration/precompiles/staking/test_staking.go

@@ -393,7 +393,7 @@ func (s *PrecompileTestSuite) TestRun() {
 				s.Require().NoError(err, "failed to pack input")
 				return input
 			},
-			21559, // use enough gas to avoid out of gas error
+			25000, // use enough gas to avoid out of gas error
 			true,
 			false,
 			true,
@@ -404,7 +404,7 @@ func (s *PrecompileTestSuite) TestRun() {
 			func(_ keyring.Key) []byte {
 				return []byte("invalid")
 			},
-			21559, // use enough gas to avoid out of gas error
+			25000, // use enough gas to avoid out of gas error
 			false,
 			false,
 			true,

tests/integration/x/vm/test_statedb.go

@@ -911,6 +911,68 @@ func (s *KeeperTestSuite) TestAddSlotToAccessList() {
 // 	}
 // }
 
+// TestGetAccountLocked verifies Keeper.GetAccount snapshots LockedCoins for
+// the EVM denom into the Account at load time. The snapshot powers the commit
+// path's bank-balance reconstruction without re-reading LockedCoins after a
+// precompile may have mutated DelegatedVesting on a vesting account.
+func (s *KeeperTestSuite) TestGetAccountLocked() {
+	addr := utiltx.GenerateAddress()
+	bondDenom := s.Network.GetBaseDenom()
+
+	testCases := []struct {
+		name      string
+		malleate  func()
+		expLocked *big.Int
+	}{
+		{
+			"non-existent account returns nil",
+			func() {},
+			nil, // GetAccount returns nil entirely; checked separately
+		},
+		{
+			"base account has zero Locked snapshot",
+			func() {
+				ctx := s.Network.GetContext()
+				err := s.Network.App.GetBankKeeper().SendCoins(ctx, s.Keyring.GetAccAddr(0), addr.Bytes(), sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100))))
+				s.Require().NoError(err)
+			},
+			big.NewInt(0),
+		},
+		{
+			"vesting account snapshots OriginalVesting at start time",
+			func() {
+				ctx := s.Network.GetContext()
+				accAddr := sdk.AccAddress(addr.Bytes())
+				err := s.Network.App.GetBankKeeper().SendCoins(ctx, s.Keyring.GetAccAddr(0), accAddr, sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100))))
+				s.Require().NoError(err)
+
+				baseAccount := s.Network.App.GetAccountKeeper().GetAccount(ctx, accAddr).(*authtypes.BaseAccount)
+				currTime := ctx.BlockTime().Unix()
+				acc, err := vestingtypes.NewContinuousVestingAccount(baseAccount, sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100))), currTime, currTime+100)
+				s.Require().NoError(err)
+				s.Network.App.GetAccountKeeper().SetAccount(ctx, acc)
+			},
+			big.NewInt(100),
+		},
+	}
+
+	for _, tc := range testCases {
+		s.Run(tc.name, func() {
+			s.SetupTest()
+			tc.malleate()
+			acc := s.Network.App.GetEVMKeeper().GetAccount(s.Network.GetContext(), addr)
+			if tc.expLocked == nil {
+				s.Require().Nil(acc, "expected nil account")
+				return
+			}
+			s.Require().NotNil(acc, "expected non-nil account")
+			locked := acc.LockedBalanceSnapshot()
+			s.Require().NotNil(locked, "expected Locked snapshot to be populated")
+			s.Require().Zero(tc.expLocked.Cmp(locked), "Locked snapshot mismatch: want %s got %s", tc.expLocked, locked)
+		})
+	}
+}
+
 func (s *KeeperTestSuite) TestSetBalance() {
 	amount := common.U2560
 	totalBalance := common.U2560
@@ -1048,6 +1110,96 @@ func (s *KeeperTestSuite) TestSetBalance() {
 	}
 }
 
+func (s *KeeperTestSuite) TestSetBalanceWithLocked() {
+	amount := common.U2560
+	var locked *big.Int
+	addr := utiltx.GenerateAddress()
+
+	testCases := []struct {
+		name           string
+		addr           common.Address
+		malleate       func()
+		expErr         bool
+		expTotalAmount func() *uint256.Int
+		expSpendable   func() *uint256.Int
+	}{
+		{
+			"non vesting account: locked overrides reread",
+			addr,
+			func() {
+				amount = uint256.NewInt(100)
+				locked = big.NewInt(50)
+			},
+			false,
+			func() *uint256.Int {
+				return uint256.NewInt(150)
+			},
+			func() *uint256.Int {
+				// All funds are spendable on a non base account, the snapshot
+				// only inflates the bank balance reconstruction.
+				return uint256.NewInt(150)
+			},
+		},
+		{
+			"vesting account: locked snapshot beats current LockedCoins re-read",
+			addr,
+			func() {
+				ctx := s.Network.GetContext()
+				accAddr := sdk.AccAddress(addr.Bytes())
+				err := s.Network.App.GetBankKeeper().SendCoins(ctx, s.Keyring.GetAccAddr(0), accAddr, sdk.NewCoins(sdk.NewCoin(s.Network.GetBaseDenom(), math.NewInt(100))))
+				s.Require().NoError(err)
+
+				baseAccount := s.Network.App.GetAccountKeeper().GetAccount(ctx, accAddr).(*authtypes.BaseAccount)
+				baseDenom := s.Network.GetBaseDenom()
+				currTime := s.Network.GetContext().BlockTime().Unix()
+
+				// setup vesting account with 40 locked to simulate a spend
+				acc, err := vestingtypes.NewContinuousVestingAccount(
+					baseAccount,
+					sdk.NewCoins(sdk.NewCoin(baseDenom, math.NewInt(40))),
+					currTime, currTime+100,
+				)
+				s.Require().NoError(err)
+				s.Network.App.GetAccountKeeper().SetAccount(ctx, acc)
+
+				amount = uint256.NewInt(100)
+
+				// override locked to 100
+				locked = big.NewInt(100)
+			},
+			false,
+			func() *uint256.Int {
+				// ensure we used the override of 100
+				return uint256.NewInt(200)
+			},
+			func() *uint256.Int {
+				// spendable = bank balance − current LockedCoins = 200 − 40 = 160.
+				return uint256.NewInt(160)
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		s.Run(tc.name, func() {
+			s.SetupTest()
+
+			tc.malleate()
+			err := s.Network.App.GetEVMKeeper().SetBalanceWithLocked(s.Network.GetContext(), tc.addr, amount, locked)
+			if tc.expErr {
+				s.Require().Error(err)
+				return
+			}
+
+			balance := s.Network.App.GetEVMKeeper().GetBalance(s.Network.GetContext(), tc.addr)
+			s.Require().NoError(err)
+			s.Require().Equal(tc.expTotalAmount(), balance)
+
+			spendable := s.Network.App.GetEVMKeeper().SpendableCoin(s.Network.GetContext(), tc.addr)
+			s.Require().Equal(tc.expSpendable(), spendable)
+		})
+	}
+}
+
 func (s *KeeperTestSuite) TestDeleteAccount() {
 	var (
 		ctx          sdk.Context

x/erc20/keeper/dynamic_precompiles.go

@@ -45,21 +45,14 @@ func (k Keeper) RegisterERC20CodeHash(ctx sdk.Context, erc20Addr common.Address)
 		k.evmKeeper.SetCode(ctx, codeHash, bytecode)
 	}
 
-	var (
-		nonce   uint64
-		balance = common.U2560
-	)
-	// keep balance and nonce if account exists
-	if acc := k.evmKeeper.GetAccount(ctx, erc20Addr); acc != nil {
-		nonce = acc.Nonce
-		balance = acc.Balance
+	// reuse account with modified code hash if it already exists
+	acc := k.evmKeeper.GetAccount(ctx, erc20Addr)
+	if acc == nil {
+		acc = statedb.NewEmptyAccount()
 	}
+	acc.CodeHash = codeHash
 
-	return k.evmKeeper.SetAccount(ctx, erc20Addr, statedb.Account{
-		CodeHash: codeHash,
-		Nonce:    nonce,
-		Balance:  balance,
-	})
+	return k.evmKeeper.SetAccount(ctx, erc20Addr, *acc)
 }
 
 // UnRegisterERC20CodeHash sets the codehash for the account to an empty one

x/vm/keeper/keeper.go

@@ -337,6 +337,18 @@ func (k *Keeper) SpendableCoin(ctx sdk.Context, addr common.Address) *uint256.In
 	return result
 }
 
+// lockedCoin loads account's locked balance of the gas token.
+func (k *Keeper) lockedCoin(ctx sdk.Context, addr common.Address) *big.Int {
+	ctx, span := ctx.StartSpan(tracer, "LockedCoin", trace.WithAttributes(attribute.String("address", addr.Hex())))
+	defer span.End()
+	cosmosAddr := sdk.AccAddress(addr.Bytes())
+
+	lockedCoins := k.bankWrapper.LockedCoins(ctx, cosmosAddr)
+	lockedGasCoin := lockedCoins.AmountOf(types.GetEVMCoinDenom())
+
+	return lockedGasCoin.BigInt()
+}
+
 // GetBalance load account's balance of gas token.
 func (k *Keeper) GetBalance(ctx sdk.Context, addr common.Address) *uint256.Int {
 	ctx, span := ctx.StartSpan(tracer, "GetBalance", trace.WithAttributes(attribute.String("address", addr.Hex())))