Merge pull request #190 from pushchain/feat/uea-migration-upgrade

feat: uea migration changes

Author: 0xNilesh

app/upgrades.go

@@ -8,6 +8,7 @@ import (
 	"github.com/pushchain/push-chain-node/app/upgrades"
 	aiauditfixes "github.com/pushchain/push-chain-node/app/upgrades/ai-audit-fixes"
 	aiauditfixes2 "github.com/pushchain/push-chain-node/app/upgrades/ai-audit-fixes-2"
+	ueamigration "github.com/pushchain/push-chain-node/app/upgrades/uea-migration"
 	ceagasandpayload "github.com/pushchain/push-chain-node/app/upgrades/cea-gas-and-payload"
 	ceapayloadverificationfix "github.com/pushchain/push-chain-node/app/upgrades/cea-payload-verification-fix"
 	chainmeta "github.com/pushchain/push-chain-node/app/upgrades/chain-meta"
@@ -55,6 +56,7 @@ var Upgrades = []upgrades.Upgrade{
 	ceapayloadverificationfix.NewUpgrade(),
 	aiauditfixes.NewUpgrade(),
 	aiauditfixes2.NewUpgrade(),
+	ueamigration.NewUpgrade(),
 }
 
 // RegisterUpgradeHandlers registers the chain upgrade handlers

app/upgrades/uea-migration/upgrade.go

@@ -0,0 +1,35 @@
+package ueamigration
+
+import (
+	"context"
+
+	storetypes "cosmossdk.io/store/types"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	"github.com/pushchain/push-chain-node/app/upgrades"
+)
+
+const UpgradeName = "uea-migration"
+
+func NewUpgrade() upgrades.Upgrade {
+	return upgrades.Upgrade{
+		UpgradeName:          UpgradeName,
+		CreateUpgradeHandler: CreateUpgradeHandler,
+		StoreUpgrades: storetypes.StoreUpgrades{
+			Added:   []string{},
+			Deleted: []string{},
+		},
+	}
+}
+
+func CreateUpgradeHandler(
+	mm upgrades.ModuleManager,
+	configurator module.Configurator,
+	ak *upgrades.AppKeepers,
+) upgradetypes.UpgradeHandler {
+	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		return mm.RunMigrations(ctx, configurator, fromVM)
+	}
+}

test/integration/uexecutor/inbound_cea_smart_contract_test.go

@@ -260,7 +260,13 @@ func TestInboundCEASmartContractRecipient(t *testing.T) {
 	})
 
 	t.Run("executeUniversalTx PCTx is recorded for smart contract recipient", func(t *testing.T) {
-		chainApp, ctx, vals, inbound, coreVals, _ := setupInboundCEASmartContractTest(t, 4)
+		chainApp, ctx, vals, inbound, coreVals, contractAddr := setupInboundCEASmartContractTest(t, 4)
+
+		// Fund the smart contract with upc so gas fee deduction succeeds
+		contractAccAddr := sdk.AccAddress(contractAddr.Bytes())
+		fundCoins := sdk.NewCoins(sdk.NewInt64Coin("upc", 1_000_000_000))
+		require.NoError(t, chainApp.BankKeeper.MintCoins(ctx, "mint", fundCoins))
+		require.NoError(t, chainApp.BankKeeper.SendCoinsFromModuleToAccount(ctx, "mint", contractAccAddr, fundCoins))
 
 		for i := 0; i < 3; i++ {
 			valAddr, err := sdk.ValAddressFromBech32(coreVals[i].OperatorAddress)
@@ -283,13 +289,51 @@ func TestInboundCEASmartContractRecipient(t *testing.T) {
 		require.Empty(t, callPcTx.ErrorMsg)
 	})
 
-	t.Run("EOA recipient receives deposit and executeUniversalTx call", func(t *testing.T) {
+	t.Run("gas fees deducted from smart contract recipient after executeUniversalTx", func(t *testing.T) {
+		chainApp, ctx, vals, inbound, coreVals, contractAddr := setupInboundCEASmartContractTest(t, 4)
+
+		// Fund the smart contract with upc so fee deduction can succeed
+		contractAccAddr := sdk.AccAddress(contractAddr.Bytes())
+		fundCoins := sdk.NewCoins(sdk.NewInt64Coin("upc", 1_000_000_000))
+		require.NoError(t, chainApp.BankKeeper.MintCoins(ctx, "mint", fundCoins))
+		require.NoError(t, chainApp.BankKeeper.SendCoinsFromModuleToAccount(ctx, "mint", contractAccAddr, fundCoins))
+
+		balanceBefore := chainApp.BankKeeper.GetBalance(ctx, contractAccAddr, "upc")
+
+		// Reach quorum
+		for i := 0; i < 3; i++ {
+			valAddr, err := sdk.ValAddressFromBech32(coreVals[i].OperatorAddress)
+			require.NoError(t, err)
+			coreValAcc := sdk.AccAddress(valAddr).String()
+
+			err = utils.ExecVoteInbound(t, ctx, chainApp, vals[i], coreValAcc, inbound)
+			require.NoError(t, err)
+		}
+
+		// Verify executeUniversalTx PCTx has gas_used > 0
+		utxKey := uexecutortypes.GetInboundUniversalTxKey(*inbound)
+		utx, found, err := chainApp.UexecutorKeeper.GetUniversalTx(ctx, utxKey)
+		require.NoError(t, err)
+		require.True(t, found)
+		require.GreaterOrEqual(t, len(utx.PcTx), 2, "should have deposit + executeUniversalTx PCTxs")
+
+		callPcTx := utx.PcTx[1]
+		require.Equal(t, "SUCCESS", callPcTx.Status)
+		require.Greater(t, callPcTx.GasUsed, uint64(0), "executeUniversalTx should report gas used")
+
+		// Verify upc balance decreased (gas was deducted)
+		balanceAfter := chainApp.BankKeeper.GetBalance(ctx, contractAccAddr, "upc")
+		require.True(t, balanceAfter.Amount.LT(balanceBefore.Amount),
+			"smart contract upc balance should decrease after gas fee deduction (before=%s, after=%s)",
+			balanceBefore.Amount, balanceAfter.Amount)
+	})
+
+	t.Run("EOA recipient receives deposit only, no executeUniversalTx", func(t *testing.T) {
 		chainApp, ctx, vals, _, coreVals, _ := setupInboundCEASmartContractTest(t, 4)
 		usdcAddress := utils.GetDefaultAddresses().ExternalUSDCAddr
 		testAddress := utils.GetDefaultAddresses().DefaultTestAddr
 
-		// TargetAddr2 is a plain EOA — deposit lands there and executeUniversalTx is called
-		// (calling to an EOA in EVM succeeds with empty output)
+		// TargetAddr2 is a plain EOA (no contract code deployed)
 		eoaRecipient := utils.GetDefaultAddresses().TargetAddr2
 
 		validUP := &uexecutortypes.UniversalPayload{
@@ -335,8 +379,15 @@ func TestInboundCEASmartContractRecipient(t *testing.T) {
 		require.NoError(t, err)
 		require.True(t, found)
 
-		// Expect 2 PCTxs: deposit + executeUniversalTx
-		require.GreaterOrEqual(t, len(utx.PcTx), 2, "should have deposit and executeUniversalTx PCTxs")
+		// EOA recipient: deposit PCTx only, no executeUniversalTx PCTx
+		// (code size check skips executeUniversalTx for EOAs, but payload execution via UEA may still run)
+		require.GreaterOrEqual(t, len(utx.PcTx), 1, "should have at least deposit PCTx for EOA recipient")
+
+		// Verify no executeUniversalTx-specific PCTx (the smart contract call path is skipped)
+		// The deposit PCTx should succeed
+		require.Equal(t, "SUCCESS", utx.PcTx[0].Status, "deposit should succeed for EOA recipient")
+
+		// No outbound should be created from executeUniversalTx (which was skipped)
 		require.Equal(t, "SUCCESS", utx.PcTx[0].Status, "deposit should succeed for EOA recipient")
 
 		// isCEA path never creates a revert outbound

test/integration/uexecutor/inbound_initiated_outbound_test.go

@@ -116,8 +116,8 @@ func setupInboundInitiatedOutboundTest(t *testing.T, numVals int) (*app.ChainApp
 
 	ueModuleAccAddress, _ := app.UexecutorKeeper.GetUeModuleAddress(ctx)
 	receipt, err := app.UexecutorKeeper.DeployUEAV2(ctx, ueModuleAccAddress, validUA)
-	ueaAddrHex := common.BytesToAddress(receipt.Ret)
 	require.NoError(t, err)
+	ueaAddrHex := common.BytesToAddress(receipt.Ret)
 
 	// signature
 	validVerificationData := "0xa7531ada733322bd6708c94cba5a7dbd1ce25bccf010f774777b039713fc330643e23b7ef2a4609244900c6ab9a03d83d3ecf73edf6b451f21cc7dbda625a3211b"

test/integration/uexecutor/vote_inbound_validation_test.go

@@ -123,13 +123,16 @@ func TestVoteInboundValidation(t *testing.T) {
 	t.Run("inbound with invalid payload records failed PCTx", func(t *testing.T) {
 		chainApp, ctx, vals, coreVals, _ := setupInboundValidationTest(t, 4)
 
-		// Construct a FUNDS_AND_PAYLOAD inbound with a malformed UniversalPayload:
-		// empty "to" encodes as zero address in RawPayload, passes ValidateForExecution,
-		// but fails at EVM execution level. The failure is recorded as a failed PCTx.
+		// Construct a FUNDS_AND_PAYLOAD inbound whose payload will revert at EVM execution.
+		// The UEA's executeUniversalTx skips verification for module-sender calls, so we
+		// must trigger a revert in the execution step itself.
+		// Strategy: call the Handler contract with an invalid function selector — the
+		// Handler has no fallback, so the low-level .call() returns success=false and
+		// the UEA reverts with ExecutionFailed().
 		malformedPayload := &uexecutortypes.UniversalPayload{
-			To:                   "", // empty "to" address -- fails at EVM execution
+			To:                   utils.GetDefaultAddresses().HandlerAddr.Hex(), // real contract, no fallback
 			Value:                "0",
-			Data:                 "0x",
+			Data:                 "0xdeadbeef", // invalid selector — no matching function
 			GasLimit:             "21000000",
 			MaxFeePerGas:         "1000000000",
 			MaxPriorityFeePerGas: "200000000",

test/utils/bytecode.go

@@ -184,7 +184,7 @@ func (k Keeper) CallUEAExecutePayload(
 		false, // gasless = false (@dev: we need gas to be emitted in the tx receipt)
 		false, // not a module sender
 		nil,
-		"executePayload",
+		"executeUniversalTx",
 		abiUniversalPayload,
 		verificationData,
 	)

x/uexecutor/keeper/evm.go

@@ -85,8 +85,13 @@ func (k Keeper) ExecuteInboundFundsAndPayload(ctx context.Context, utx types.Uni
 					}
 				}
 			} else {
-				// Non-UEA path (smart contract or EOA): deposit PRC20 and call executeUniversalTx
-				isSmartContract = true
+				// Non-UEA: check if recipient has code (smart contract) vs EOA
+				codeHash := k.evmKeeper.GetCodeHash(sdkCtx, ueaAddr)
+				if codeHash != types.EmptyCodeHash && codeHash != (common.Hash{}) {
+					// Smart contract: will call executeUniversalTx after deposit
+					isSmartContract = true
+				}
+				// EOA: just deposit, skip executeUniversalTx (no contract to call)
 				if inboundAmount.Sign() > 0 {
 					receipt, execErr = k.depositPRC20(
 						sdkCtx,
@@ -262,12 +267,20 @@ func (k Keeper) ExecuteInboundFundsAndPayload(ctx context.Context, utx types.Uni
 			BlockHeight: uint64(sdkCtx.BlockHeight()),
 			Status:      "FAILED",
 		}
+		if contractReceipt != nil {
+			callPcTx.TxHash = contractReceipt.Hash
+			callPcTx.GasUsed = contractReceipt.GasUsed
+		}
 		if contractErr != nil {
 			callPcTx.ErrorMsg = contractErr.Error()
 		} else {
-			callPcTx.TxHash = contractReceipt.Hash
-			callPcTx.GasUsed = contractReceipt.GasUsed
-			callPcTx.Status = "SUCCESS"
+			// Deduct gas fees from the recipient contract address
+			if feeErr := k.DeductGasFeesFromReceipt(ctx, sdkCtx, ueaAddr, contractReceipt, utx.InboundTx.UniversalPayload); feeErr != nil {
+				callPcTx.Status = "FAILED"
+				callPcTx.ErrorMsg = fmt.Sprintf("gas fee deduction failed: %s", feeErr.Error())
+			} else {
+				callPcTx.Status = "SUCCESS"
+			}
 		}
 		if updateErr := k.UpdateUniversalTx(ctx, universalTxKey, func(utx *types.UniversalTx) error {
 			utx.PcTx = append(utx.PcTx, &callPcTx)

x/uexecutor/keeper/execute_inbound_funds_and_payload.go

@@ -87,8 +87,12 @@ func (k Keeper) ExecuteInboundGasAndPayload(ctx context.Context, utx types.Unive
 							}
 						}
 					} else {
-						// Non-UEA path (smart contract or EOA): deposit + autoswap and call executeUniversalTx
-						isSmartContract = true
+						// Non-UEA: check if recipient has code (smart contract) vs EOA
+						codeHash := k.evmKeeper.GetCodeHash(sdkCtx, ueaAddr)
+						if codeHash != types.EmptyCodeHash && codeHash != (common.Hash{}) {
+							isSmartContract = true
+						}
+						// EOA: just deposit, skip executeUniversalTx
 						if amount.Sign() > 0 {
 							prc20AddrHex := common.HexToAddress(tokenConfig.NativeRepresentation.ContractAddress)
 							receipt, execErr = k.gasAndPayloadDepositAutoSwap(sdkCtx, prc20AddrHex, ueaAddr, amount)
@@ -261,12 +265,20 @@ func (k Keeper) ExecuteInboundGasAndPayload(ctx context.Context, utx types.Unive
 			BlockHeight: uint64(sdkCtx.BlockHeight()),
 			Status:      "FAILED",
 		}
+		if contractReceipt != nil {
+			callPcTx.TxHash = contractReceipt.Hash
+			callPcTx.GasUsed = contractReceipt.GasUsed
+		}
 		if contractErr != nil {
 			callPcTx.ErrorMsg = contractErr.Error()
 		} else if contractReceipt != nil {
-			callPcTx.TxHash = contractReceipt.Hash
-			callPcTx.GasUsed = contractReceipt.GasUsed
-			callPcTx.Status = "SUCCESS"
+			// Deduct gas fees from the recipient contract address
+			if feeErr := k.DeductGasFeesFromReceipt(ctx, sdkCtx, ueaAddr, contractReceipt, utx.InboundTx.UniversalPayload); feeErr != nil {
+				callPcTx.Status = "FAILED"
+				callPcTx.ErrorMsg = fmt.Sprintf("gas fee deduction failed: %s", feeErr.Error())
+			} else {
+				callPcTx.Status = "SUCCESS"
+			}
 		}
 		if updateErr := k.UpdateUniversalTx(ctx, universalTxKey, func(utx *types.UniversalTx) error {
 			utx.PcTx = append(utx.PcTx, &callPcTx)

x/uexecutor/keeper/execute_inbound_gas_and_payload.go

@@ -2,11 +2,10 @@ package keeper
 
 import (
 	"context"
-	"math/big"
+	"fmt"
 
 	"cosmossdk.io/errors"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	sdkErrors "github.com/cosmos/cosmos-sdk/types/errors"
 	vmtypes "github.com/cosmos/evm/x/vm/types"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/pushchain/push-chain-node/utils"
@@ -23,9 +22,8 @@ func (k Keeper) ExecutePayloadV2(ctx context.Context, evmFrom common.Address, ue
 		"from", evmFrom.Hex(),
 	)
 
-	// Step 1: Parse and validate payload and verificationData
-	payload, err := types.NewAbiUniversalPayload(universalPayload)
-	if err != nil {
+	// Step 1: Validate payload and verificationData early (fast-fail before EVM work)
+	if _, err := types.NewAbiUniversalPayload(universalPayload); err != nil {
 		return nil, errors.Wrapf(err, "invalid universal payload")
 	}
 
@@ -35,44 +33,23 @@ func (k Keeper) ExecutePayloadV2(ctx context.Context, evmFrom common.Address, ue
 	}
 
 	// Step 2: Execute payload through UEA
-	receipt, err := k.CallUEAExecutePayload(sdkCtx, evmFrom, ueaAddr, universalPayload, verificationDataVal)
-	if err != nil {
-		return nil, err
-	}
-
-	gasUnitsUsed := receipt.GasUsed
-	gasUnitsUsedBig := new(big.Int).SetUint64(gasUnitsUsed)
-
-	k.Logger().Debug("payload executed via UEA",
-		"uea", ueaAddr.Hex(),
-		"tx_hash", receipt.Hash,
-		"gas_used", gasUnitsUsed,
-	)
+	receipt, execErr := k.CallUEAExecutePayload(sdkCtx, evmFrom, ueaAddr, universalPayload, verificationDataVal)
 
-	// Step 3: Handle fee calculation and deduction
-	ueaAccAddr := sdk.AccAddress(ueaAddr.Bytes())
-
-	baseFee := k.feemarketKeeper.GetBaseFee(sdkCtx)
-	if baseFee.IsNil() {
-		return nil, errors.Wrapf(sdkErrors.ErrLogic, "base fee not found")
-	}
-
-	gasCost, err := k.CalculateGasCost(baseFee, payload.MaxFeePerGas, payload.MaxPriorityFeePerGas, gasUnitsUsed)
-	if err != nil {
-		return nil, errors.Wrapf(err, "failed to calculate gas cost")
+	// Step 3: Deduct gas fees regardless of success/failure.
+	// If deduction fails, return error so the caller records a FAILED PCTx.
+	// The receipt is still returned so callers can capture the tx hash.
+	if feeErr := k.DeductGasFeesFromReceipt(ctx, sdkCtx, ueaAddr, receipt, universalPayload); feeErr != nil {
+		return receipt, fmt.Errorf("gas fee deduction failed: %w", feeErr)
 	}
 
-	if gasUnitsUsedBig.Cmp(payload.GasLimit) > 0 {
-		return nil, errors.Wrapf(sdkErrors.ErrOutOfGas, "gas cost (%d) exceeds limit (%d)", gasCost, payload.GasLimit)
+	if execErr != nil {
+		return receipt, execErr
 	}
 
-	if err = k.DeductAndBurnFees(ctx, ueaAccAddr, gasCost); err != nil {
-		return nil, errors.Wrapf(err, "failed to deduct fees from %s", ueaAccAddr)
-	}
-
-	k.Logger().Debug("fees deducted for payload execution",
+	k.Logger().Debug("payload executed via UEA",
 		"uea", ueaAddr.Hex(),
-		"gas_cost", gasCost.String(),
+		"tx_hash", receipt.Hash,
+		"gas_used", receipt.GasUsed,
 	)
 
 	return receipt, nil