Merge pull request #206 from pushchain/feat/tss-fund-migration-fixes

feat: added dynamic gas limit changes in fund migration

Author: 0xNilesh

api/utss/v1/types.pulsar.go

@@ -10,6 +10,7 @@ import (
 	aiauditfixes2 "github.com/pushchain/push-chain-node/app/upgrades/ai-audit-fixes-2"
 	purgeexpiredoutbounds "github.com/pushchain/push-chain-node/app/upgrades/purge-expired-outbounds"
 	removeutxverifier "github.com/pushchain/push-chain-node/app/upgrades/remove-utxverifier"
+	tssfundmigrationfixes "github.com/pushchain/push-chain-node/app/upgrades/tss-fund-migration-fixes"
 	tssmigration "github.com/pushchain/push-chain-node/app/upgrades/tss-migration"
 	ueamigration "github.com/pushchain/push-chain-node/app/upgrades/uea-migration"
 	ceagasandpayload "github.com/pushchain/push-chain-node/app/upgrades/cea-gas-and-payload"
@@ -63,6 +64,7 @@ var Upgrades = []upgrades.Upgrade{
 	tssmigration.NewUpgrade(),
 	purgeexpiredoutbounds.NewUpgrade(),
 	removeutxverifier.NewUpgrade(),
+	tssfundmigrationfixes.NewUpgrade(),
 }
 
 // RegisterUpgradeHandlers registers the chain upgrade handlers

app/upgrades.go

@@ -0,0 +1,53 @@
+package tssfundmigrationfixes
+
+import (
+	"context"
+
+	storetypes "cosmossdk.io/store/types"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	"github.com/pushchain/push-chain-node/app/upgrades"
+)
+
+const UpgradeName = "tss-fund-migration-fixes"
+
+func NewUpgrade() upgrades.Upgrade {
+	return upgrades.Upgrade{
+		UpgradeName:          UpgradeName,
+		CreateUpgradeHandler: CreateUpgradeHandler,
+		StoreUpgrades: storetypes.StoreUpgrades{
+			Added:   []string{},
+			Deleted: []string{},
+		},
+	}
+}
+
+// CreateUpgradeHandler runs the utss v3 → v4 migration which backfills
+// FundMigration.l1_gas_fee on records stored before the field existed.
+// The new gas_limit and l1_gas_fee values used by InitiateFundMigration are
+// sourced from UniversalCore's tssFundMigrationGasLimitByChainNamespace and
+// l1GasFeeByChainNamespace mappings at call time — no state seeding required.
+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) {
+		sdkCtx := sdk.UnwrapSDKContext(ctx)
+		logger := sdkCtx.Logger().With("upgrade", UpgradeName)
+		logger.Info("Starting upgrade handler")
+		logger.Info("Feature: FundMigration.gas_limit and l1_gas_fee now sourced from UniversalCore per-chain mappings")
+
+		versionMap, err := mm.RunMigrations(ctx, configurator, fromVM)
+		if err != nil {
+			logger.Error("RunMigrations failed", "error", err)
+			return nil, err
+		}
+
+		logger.Info("Upgrade complete", "upgrade", UpgradeName)
+		return versionMap, nil
+	}
+}

app/upgrades/tss-fund-migration-fixes/upgrade.go

@@ -0,0 +1,22 @@
+package tssfundmigrationfixes_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	tssfundmigrationfixes "github.com/pushchain/push-chain-node/app/upgrades/tss-fund-migration-fixes"
+)
+
+// TestNewUpgrade_Identity verifies the upgrade descriptor carries the expected
+// name, wires up a non-nil handler factory, and declares no store additions or
+// deletions (the migration is in-place on existing kv keys).
+func TestNewUpgrade_Identity(t *testing.T) {
+	u := tssfundmigrationfixes.NewUpgrade()
+
+	require.Equal(t, "tss-fund-migration-fixes", u.UpgradeName)
+	require.NotNil(t, u.CreateUpgradeHandler, "upgrade must expose a handler factory")
+	require.Empty(t, u.StoreUpgrades.Added, "no new KV stores expected")
+	require.Empty(t, u.StoreUpgrades.Deleted, "no KV stores deleted")
+	require.Empty(t, u.StoreUpgrades.Renamed, "no KV stores renamed")
+}

app/upgrades/tss-fund-migration-fixes/upgrade_test.go

@@ -103,5 +103,6 @@ message FundMigration {
   int64  completed_block    = 9;
   string tx_hash            = 10;
   string gas_price          = 11; // gas price from oracle (wei)
-  uint64 gas_limit          = 12; // gas limit for native transfer (21000)
+  uint64 gas_limit          = 12; // gas limit sourced from UniversalCore per chain namespace
+  string l1_gas_fee         = 13; // L1 data-availability fee (wei) from UniversalCore; 0 for non-L2 chains
 }

proto/utss/v1/types.proto

@@ -2,10 +2,15 @@ package integrationtest
 
 import (
 	"fmt"
+	"math/big"
 	"strconv"
+	"strings"
 	"testing"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/ethereum/go-ethereum/accounts/abi"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/stretchr/testify/require"
 
 	"github.com/pushchain/push-chain-node/app"
@@ -18,12 +23,84 @@ import (
 
 const testChain = "eip155:11155111"
 
+// universalCoreSetupABI exposes the admin methods needed to configure
+// per-chain mappings during test setup. These are intentionally kept out of
+// the production ABI (x/uexecutor/types/abi.go) — Go-side keeper code never
+// calls them; only tests do.
+const universalCoreSetupABI = `[
+    {
+      "type": "function",
+      "name": "grantRole",
+      "inputs": [
+        { "name": "role",    "type": "bytes32", "internalType": "bytes32" },
+        { "name": "account", "type": "address", "internalType": "address" }
+      ],
+      "outputs": [],
+      "stateMutability": "nonpayable"
+    },
+    {
+      "type": "function",
+      "name": "setL1GasFeeByChain",
+      "inputs": [
+        { "name": "chainNamespace", "type": "string",  "internalType": "string" },
+        { "name": "l1GasFee",       "type": "uint256", "internalType": "uint256" }
+      ],
+      "outputs": [],
+      "stateMutability": "nonpayable"
+    },
+    {
+      "type": "function",
+      "name": "setTssFundMigrationGasLimitByChain",
+      "inputs": [
+        { "name": "chainNamespace", "type": "string",  "internalType": "string" },
+        { "name": "gasLimit",       "type": "uint256", "internalType": "uint256" }
+      ],
+      "outputs": [],
+      "stateMutability": "nonpayable"
+    }
+]`
+
+// seedFundMigrationChainValues grants MANAGER_ROLE to the admin and seeds the
+// per-chain tss-fund-migration gas limit and L1 gas fee on UniversalCore.
+// InitiateFundMigration rejects a zero gas limit, so without this seeding the
+// keeper read returns 0 and the migration fails validation.
+func seedFundMigrationChainValues(
+	t *testing.T,
+	chainApp *app.ChainApp,
+	ctx sdk.Context,
+	admin common.Address,
+	chain string,
+	gasLimit, l1GasFee *big.Int,
+) {
+	t.Helper()
+
+	handlerAddr := utils.GetDefaultAddresses().HandlerAddr
+	setupABI, err := abi.JSON(strings.NewReader(universalCoreSetupABI))
+	require.NoError(t, err)
+
+	managerRole := crypto.Keccak256Hash([]byte("MANAGER_ROLE"))
+	var roleArg [32]byte
+	copy(roleArg[:], managerRole.Bytes())
+
+	_, err = chainApp.EVMKeeper.CallEVM(ctx, setupABI, admin, handlerAddr, true, "grantRole", roleArg, admin)
+	require.NoError(t, err, "grant MANAGER_ROLE")
+
+	_, err = chainApp.EVMKeeper.CallEVM(ctx, setupABI, admin, handlerAddr, true, "setTssFundMigrationGasLimitByChain", chain, gasLimit)
+	require.NoError(t, err, "seed tss fund migration gas limit")
+
+	_, err = chainApp.EVMKeeper.CallEVM(ctx, setupABI, admin, handlerAddr, true, "setL1GasFeeByChain", chain, l1GasFee)
+	require.NoError(t, err, "seed l1 gas fee")
+}
+
 // setupFundMigrationTest initializes app with validators, a finalized keygen key, and a chain config.
 // Returns app, ctx, validator addresses, and the finalized key ID.
 func setupFundMigrationTest(t *testing.T, numVals int, outboundEnabled bool) (*app.ChainApp, sdk.Context, []string, string) {
 	t.Helper()
 
-	app, ctx, _, validators := utils.SetAppWithMultipleValidators(t, numVals)
+	app, ctx, baseAccounts, validators := utils.SetAppWithMultipleValidators(t, numVals)
+
+	admin := common.BytesToAddress(baseAccounts[0].GetAddress().Bytes())
+	seedFundMigrationChainValues(t, app, ctx, admin, testChain, big.NewInt(21000), big.NewInt(150))
 
 	// Register universal validators
 	universalVals := make([]string, len(validators))
@@ -129,7 +206,10 @@ func TestInitiateFundMigration(t *testing.T) {
 		require.Equal(t, utsstypes.FundMigrationStatus_FUND_MIGRATION_STATUS_PENDING, migration.Status)
 		require.Equal(t, oldKeyId, migration.OldKeyId)
 		require.Equal(t, testChain, migration.Chain)
+		// GasLimit and L1GasFee come from UniversalCore's per-chain mappings,
+		// seeded by seedFundMigrationChainValues.
 		require.Equal(t, uint64(21000), migration.GasLimit)
+		require.Equal(t, "150", migration.L1GasFee)
 		require.NotEmpty(t, migration.GasPrice)
 
 		// Verify pending index

test/integration/utss/fund_migration_test.go

@@ -30,6 +30,7 @@ type FundMigrationData struct {
 	To       string   // New TSS address (derived from current pubkey)
 	GasPrice *big.Int // Gas price from the migration event
 	GasLimit uint64   // Gas limit from the migration event
+	L1GasFee *big.Int // Extra L1 data-availability fee (wei); 0 for non-L2 chains
 }
 
 // UnsignedSigningReq contains the request for signing an outbound transaction

test/utils/bytecode.go

@@ -471,26 +471,29 @@ func (tb *TxBuilder) GetGasFeeUsed(ctx context.Context, txHash string) (string,
 }
 
 // GetFundMigrationSigningRequest builds a native token transfer for fund migration,
-// transferring the maximum possible balance (balance minus gas cost).
+// transferring the maximum possible balance (balance minus gas cost minus L1 fee).
 // Fund migration only triggers when outbound is disabled and no pending outbounds remain,
 // so the balance at signing time will equal the balance at broadcast time.
+// L1GasFee covers OP-stack sequencer data-availability charges; 0 for non-L2 chains.
 func (tb *TxBuilder) GetFundMigrationSigningRequest(ctx context.Context, data *common.FundMigrationData, nonce uint64) (*common.UnsignedSigningReq, error) {
 	fromAddr := ethcommon.HexToAddress(data.From)
 	toAddr := ethcommon.HexToAddress(data.To)
 
 	if data.GasPrice == nil || data.GasPrice.Sign() == 0 {
 		return nil, fmt.Errorf("gas price must be provided for fund migration")
 	}
+	if data.GasLimit == 0 {
+		return nil, fmt.Errorf("gas limit must be provided for fund migration")
+	}
 
 	balance, err := tb.rpcClient.GetBalance(ctx, fromAddr)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get balance of %s: %w", data.From, err)
 	}
 
-	gasCost := new(big.Int).Mul(data.GasPrice, new(big.Int).SetUint64(data.GasLimit))
-	maxTransfer := new(big.Int).Sub(balance, gasCost)
-	if maxTransfer.Sign() <= 0 {
-		return nil, fmt.Errorf("insufficient balance for gas: balance=%s gasCost=%s", balance.String(), gasCost.String())
+	maxTransfer, err := computeFundMigrationTransfer(balance, data.GasPrice, data.GasLimit, data.L1GasFee)
+	if err != nil {
+		return nil, err
 	}
 
 	tb.logger.Info().
@@ -499,6 +502,7 @@ func (tb *TxBuilder) GetFundMigrationSigningRequest(ctx context.Context, data *c
 		Str("balance", balance.String()).
 		Str("gas_price", data.GasPrice.String()).
 		Uint64("gas_limit", data.GasLimit).
+		Str("l1_gas_fee", l1GasFeeString(data.L1GasFee)).
 		Str("transfer_amount", maxTransfer.String()).
 		Msg("building fund migration tx")
 
@@ -521,6 +525,9 @@ func (tb *TxBuilder) GetFundMigrationSigningRequest(ctx context.Context, data *c
 }
 
 // BroadcastFundMigrationTx assembles and broadcasts a signed fund migration transaction.
+// The sweep amount must be recomputed here using the same formula as signing
+// (balance - gasPrice*gasLimit - l1GasFee); otherwise the broadcast tx hash
+// diverges from the signed hash.
 func (tb *TxBuilder) BroadcastFundMigrationTx(ctx context.Context, req *common.UnsignedSigningReq, data *common.FundMigrationData, signature []byte) (string, error) {
 	if len(signature) != 65 {
 		return "", fmt.Errorf("signature must be 65 bytes [r(32)|s(32)|v(1)], got %d", len(signature))
@@ -529,6 +536,9 @@ func (tb *TxBuilder) BroadcastFundMigrationTx(ctx context.Context, req *common.U
 	if data.GasPrice == nil || data.GasPrice.Sign() == 0 {
 		return "", fmt.Errorf("gas price must be provided for fund migration")
 	}
+	if data.GasLimit == 0 {
+		return "", fmt.Errorf("gas limit must be provided for fund migration")
+	}
 
 	fromAddr := ethcommon.HexToAddress(data.From)
 	toAddr := ethcommon.HexToAddress(data.To)
@@ -538,10 +548,9 @@ func (tb *TxBuilder) BroadcastFundMigrationTx(ctx context.Context, req *common.U
 		return "", fmt.Errorf("failed to get balance of %s: %w", data.From, err)
 	}
 
-	gasCost := new(big.Int).Mul(data.GasPrice, new(big.Int).SetUint64(data.GasLimit))
-	maxTransfer := new(big.Int).Sub(balance, gasCost)
-	if maxTransfer.Sign() <= 0 {
-		return "", fmt.Errorf("insufficient balance for gas during broadcast")
+	maxTransfer, err := computeFundMigrationTransfer(balance, data.GasPrice, data.GasLimit, data.L1GasFee)
+	if err != nil {
+		return "", err
 	}
 
 	tx := types.NewTransaction(
@@ -574,3 +583,31 @@ func (tb *TxBuilder) BroadcastFundMigrationTx(ctx context.Context, req *common.U
 
 	return txHashStr, nil
 }
+
+// computeFundMigrationTransfer returns the native amount to sweep from the old
+// TSS address to the new one: balance - (gasPrice * gasLimit) - l1GasFee.
+// The l1GasFee covers OP-stack sequencer data-availability charges (0 for
+// non-L2 chains). All validators must compute the same value — any drift
+// here breaks the TSS signing hash.
+func computeFundMigrationTransfer(balance, gasPrice *big.Int, gasLimit uint64, l1GasFee *big.Int) (*big.Int, error) {
+	gasCost := new(big.Int).Mul(gasPrice, new(big.Int).SetUint64(gasLimit))
+	totalFee := new(big.Int).Set(gasCost)
+	if l1GasFee != nil && l1GasFee.Sign() > 0 {
+		totalFee.Add(totalFee, l1GasFee)
+	}
+	maxTransfer := new(big.Int).Sub(balance, totalFee)
+	if maxTransfer.Sign() <= 0 {
+		return nil, fmt.Errorf("insufficient balance for gas: balance=%s gasCost=%s l1GasFee=%s",
+			balance.String(), gasCost.String(), l1GasFeeString(l1GasFee))
+	}
+	return maxTransfer, nil
+}
+
+// l1GasFeeString returns a stable decimal representation of the L1 gas fee
+// for logging / error messages, treating nil as "0".
+func l1GasFeeString(v *big.Int) string {
+	if v == nil {
+		return "0"
+	}
+	return v.String()
+}