fix: calling precompile through sc with sc state update (#2853)

* fix calling precompile through sc with sc state update

* changelog

* extend comment a bit

changelog.md

@@ -40,6 +40,7 @@
 * [2672](https://github.com/zeta-chain/node/pull/2672) - check observer set for duplicates when adding a new observer or updating an existing one
 * [2787](https://github.com/zeta-chain/node/pull/2787) - ask for 3 accounts (signer, pda, system_program) on solana gateway deposit
 * [2842](https://github.com/zeta-chain/node/pull/2842) - fix: move interval assignment out of cctx loop in EVM outbound tx scheduler
+* [2853](https://github.com/zeta-chain/node/pull/2853) - calling precompile through sc with sc state update
 
 ## v19.0.0
 

e2e/contracts/teststaking/TestStaking.abi

@@ -14,6 +14,19 @@
     "stateMutability": "payable",
     "type": "fallback"
   },
+  {
+    "inputs": [],
+    "name": "counter",
+    "outputs": [
+      {
+        "internalType": "uint256",
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "stateMutability": "view",
+    "type": "function"
+  },
   {
     "inputs": [],
     "name": "depositWZETA",
@@ -143,6 +156,64 @@
     "stateMutability": "nonpayable",
     "type": "function"
   },
+  {
+    "inputs": [
+      {
+        "internalType": "address",
+        "name": "staker",
+        "type": "address"
+      },
+      {
+        "internalType": "string",
+        "name": "validator",
+        "type": "string"
+      },
+      {
+        "internalType": "uint256",
+        "name": "amount",
+        "type": "uint256"
+      }
+    ],
+    "name": "stakeAndRevert",
+    "outputs": [
+      {
+        "internalType": "bool",
+        "name": "",
+        "type": "bool"
+      }
+    ],
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
+  {
+    "inputs": [
+      {
+        "internalType": "address",
+        "name": "staker",
+        "type": "address"
+      },
+      {
+        "internalType": "string",
+        "name": "validator",
+        "type": "string"
+      },
+      {
+        "internalType": "uint256",
+        "name": "amount",
+        "type": "uint256"
+      }
+    ],
+    "name": "stakeWithStateUpdate",
+    "outputs": [
+      {
+        "internalType": "bool",
+        "name": "",
+        "type": "bool"
+      }
+    ],
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
   {
     "inputs": [
       {

e2e/contracts/teststaking/TestStaking.sol

@@ -52,6 +52,9 @@ contract TestStaking {
     WZETA wzeta;
     address owner;
 
+    // @dev used to test state change in smart contract
+    uint256 public counter = 0;
+
     constructor(address _wzeta) {
         wzeta = WZETA(_wzeta);
         owner = msg.sender;
@@ -76,6 +79,20 @@ contract TestStaking {
         return staking.stake(staker, validator, amount);
     }
 
+    function stakeWithStateUpdate(address staker, string memory validator, uint256 amount) external onlyOwner returns (bool)  {
+        counter = counter + 1;
+        bool success = staking.stake(staker, validator, amount);
+        counter = counter + 1;
+        return success;
+    }
+
+    function stakeAndRevert(address staker, string memory validator, uint256 amount) external onlyOwner returns (bool)  {
+        counter = counter + 1;
+        staking.stake(staker, validator, amount);
+        counter = counter + 1;
+        revert("testrevert");
+    }
+
     function unstake(
         address staker,
         string memory validator,

e2e/e2etests/test_precompiles_staking_through_contract.go

@@ -4,10 +4,12 @@ import (
 	"math/big"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/stretchr/testify/require"
 
+	"github.com/zeta-chain/node/cmd/zetacored/config"
 	"github.com/zeta-chain/node/e2e/contracts/teststaking"
 	"github.com/zeta-chain/node/e2e/runner"
 	"github.com/zeta-chain/node/e2e/utils"
@@ -58,15 +60,78 @@ func TestPrecompilesStakingThroughContract(r *runner.E2ERunner, args []string) {
 	utils.MustWaitForTxReceipt(r.Ctx, r.ZEVMClient, tx, r.Logger, r.ReceiptTimeout)
 	r.ZEVMAuth.Value = big.NewInt(0)
 
-	tx, err = testStaking.WithdrawWZETA(r.ZEVMAuth, big.NewInt(100000000000))
+	stakeAmount := 100000000000
+	tx, err = testStaking.WithdrawWZETA(r.ZEVMAuth, big.NewInt(int64(stakeAmount)))
 	require.NoError(r, err)
 	utils.MustWaitForTxReceipt(r.Ctx, r.ZEVMClient, tx, r.Logger, r.ReceiptTimeout)
 
-	// stake 3 to validator1 using testStaking smart contract
-	tx, err = testStaking.Stake(r.ZEVMAuth, testStakingAddr, validators[0].OperatorAddress, big.NewInt(3))
+	// bank balance at the start
+	balanceBefore, err := r.BankClient.Balance(r.Ctx, &banktypes.QueryBalanceRequest{
+		Address: sdk.AccAddress(testStakingAddr.Bytes()).String(),
+		Denom:   config.BaseDenom,
+	})
+	require.NoError(r, err)
+	require.Equal(r, int64(stakeAmount), balanceBefore.Balance.Amount.Int64())
+
+	// stake 3 to validator1 and revert in same function
+	tx, err = testStaking.StakeAndRevert(r.ZEVMAuth, testStakingAddr, validators[0].OperatorAddress, big.NewInt(3))
+	require.NoError(r, err)
+	utils.MustWaitForTxReceipt(r.Ctx, r.ZEVMClient, tx, r.Logger, r.ReceiptTimeout)
+
+	// check that bank balance was not changed because of revert in testStaking contract
+	balanceAfterRevert, err := r.BankClient.Balance(r.Ctx, &banktypes.QueryBalanceRequest{
+		Address: sdk.AccAddress(testStakingAddr.Bytes()).String(),
+		Denom:   config.BaseDenom,
+	})
+	require.NoError(r, err)
+	require.Equal(r, balanceBefore.Balance.Amount.Int64(), balanceAfterRevert.Balance.Amount.Int64())
+
+	// check that counter was not updated
+	counter, err := testStaking.Counter(&bind.CallOpts{})
+	require.NoError(r, err)
+	require.Equal(r, int64(0), counter.Int64())
+
+	// check that shares are still 0
+	sharesAfterRevert, err := testStaking.GetShares(&bind.CallOpts{}, testStakingAddr, validators[0].OperatorAddress)
+	require.NoError(r, err)
+	require.Equal(r, int64(0), sharesAfterRevert.Int64())
+
+	// stake 1 to validator1 using testStaking smart contract without smart contract state update
+	tx, err = testStaking.Stake(r.ZEVMAuth, testStakingAddr, validators[0].OperatorAddress, big.NewInt(1))
+	require.NoError(r, err)
+	utils.MustWaitForTxReceipt(r.Ctx, r.ZEVMClient, tx, r.Logger, r.ReceiptTimeout)
+
+	// check that bank balance is reduced by 1
+	balanceAfterStake, err := r.BankClient.Balance(r.Ctx, &banktypes.QueryBalanceRequest{
+		Address: sdk.AccAddress(testStakingAddr.Bytes()).String(),
+		Denom:   config.BaseDenom,
+	})
+	require.NoError(r, err)
+	require.Equal(r, balanceBefore.Balance.Amount.Int64()-1, balanceAfterStake.Balance.Amount.Int64())
+
+	// stake 2 more to validator1 using testStaking smart contract with smart contract state update
+	tx, err = testStaking.StakeWithStateUpdate(
+		r.ZEVMAuth,
+		testStakingAddr,
+		validators[0].OperatorAddress,
+		big.NewInt(2),
+	)
 	require.NoError(r, err)
 	utils.MustWaitForTxReceipt(r.Ctx, r.ZEVMClient, tx, r.Logger, r.ReceiptTimeout)
 
+	// check that bank balance is reduced by 2 more, 3 in total
+	balanceAfterStake, err = r.BankClient.Balance(r.Ctx, &banktypes.QueryBalanceRequest{
+		Address: sdk.AccAddress(testStakingAddr.Bytes()).String(),
+		Denom:   config.BaseDenom,
+	})
+	require.NoError(r, err)
+	require.Equal(r, balanceBefore.Balance.Amount.Int64()-3, balanceAfterStake.Balance.Amount.Int64())
+
+	// check that counter is updated
+	counter, err = testStaking.Counter(&bind.CallOpts{})
+	require.NoError(r, err)
+	require.Equal(r, int64(2), counter.Int64())
+
 	// check shares are set to 3
 	sharesAfterVal1, err := testStaking.GetShares(&bind.CallOpts{}, testStakingAddr, validators[0].OperatorAddress)
 	require.NoError(r, err)

precompiles/staking/staking.go

@@ -182,6 +182,7 @@ func (c *Contract) GetShares(
 
 func (c *Contract) Stake(
 	ctx sdk.Context,
+	evm *vm.EVM,
 	contract *vm.Contract,
 	method *abi.Method,
 	args []interface{},
@@ -231,6 +232,14 @@ func (c *Contract) Stake(
 		return nil, err
 	}
 
+	// if caller is not the same as origin it means call is coming through smart contract,
+	// and because state of smart contract calling precompile might be updated as well
+	// manually reduce amount in stateDB, so it is properly reflected in bank module
+	stateDB := evm.StateDB.(ptypes.ExtStateDB)
+	if contract.CallerAddress != evm.Origin {
+		stateDB.SubBalance(stakerAddress, amount)
+	}
+
 	return method.Outputs.Pack(true)
 }
 
@@ -389,7 +398,7 @@ func (c *Contract) Run(evm *vm.EVM, contract *vm.Contract, _ bool) ([]byte, erro
 	case StakeMethodName:
 		var res []byte
 		execErr := stateDB.ExecuteNativeAction(contract.Address(), nil, func(ctx sdk.Context) error {
-			res, err = c.Stake(ctx, contract, method, args)
+			res, err = c.Stake(ctx, evm, contract, method, args)
 			return err
 		})
 		if execErr != nil {