[EVM] Merge balances on associate tx (sei-protocol#1630)

* add associate tx test

* fix migration of balances

* fix getAdmin function

* add log so we can debug key import on github actions

* fix test (multi-interaction cli step)

contracts/test/AssociateTest.js

@@ -0,0 +1,87 @@
+const { fundAddress, fundSeiAddress, getSeiBalance, associateKey, importKey, waitForReceipt, bankSend, evmSend, getNativeAccount} = require("./lib");
+const { expect } = require("chai");
+
+describe("Associate Balances", function () {
+
+    const keys = {
+        "test1": {
+            seiAddress: 'sei1nzdg7e6rvkrmvp5zzmp5tupuj0088nqsa4mze4',
+            evmAddress: '0x90684e7F229f2d8E2336661f79caB693E4228Ff7'
+        },
+        "test2": {
+            seiAddress: 'sei1jqgph9jpdtvv64e3rzegxtssvgmh7lxnn8vmdq',
+            evmAddress: '0x28b2B0621f76A2D08A9e04acb7F445E61ba5b7E7'
+        },
+        "test3": {
+            seiAddress: 'sei1qkawqt7dw09rkvn53lm2deamtfcpuq9v0h6zur',
+            evmAddress: '0xCb2FB25A6a34Ca874171Ac0406d05A49BC45a1cF'
+        }
+    }
+
+    const addresses = {
+        seiAddress: 'sei1nzdg7e6rvkrmvp5zzmp5tupuj0088nqsa4mze4',
+        evmAddress: '0x90684e7F229f2d8E2336661f79caB693E4228Ff7'
+    }
+
+    function truncate(num, byThisManyDecimals) {
+        return parseFloat(`${num}`.slice(0, 12))
+    }
+
+    async function verifyAssociation(seiAddr, evmAddr, associateFunc) {
+        const beforeSei = BigInt(await getSeiBalance(seiAddr))
+        const beforeEvm = await ethers.provider.getBalance(evmAddr)
+        const gas = await associateFunc(seiAddr)
+        const afterSei = BigInt(await getSeiBalance(seiAddr))
+        const afterEvm = await ethers.provider.getBalance(evmAddr)
+
+        // console.log(`SEI Balance (before): ${beforeSei}`)
+        // console.log(`EVM Balance (before): ${beforeEvm}`)
+        // console.log(`SEI Balance (after): ${afterSei}`)
+        // console.log(`EVM Balance (after): ${afterEvm}`)
+
+        const multiplier = BigInt(1000000000000)
+        expect(afterEvm).to.equal((beforeSei * multiplier) + beforeEvm - (gas * multiplier))
+        expect(afterSei).to.equal(truncate(beforeSei - gas))
+    }
+
+    before(async function(){
+        await importKey("test1", "../contracts/test/test1.key")
+        await importKey("test2", "../contracts/test/test2.key")
+        await importKey("test3", "../contracts/test/test3.key")
+    })
+
+    it("should associate with sei transaction", async function(){
+        const addr = keys.test1
+        await fundSeiAddress(addr.seiAddress, "10000000000")
+        await fundAddress(addr.evmAddress, "200");
+
+        await verifyAssociation(addr.seiAddress, addr.evmAddress, async function(){
+            await bankSend(addr.seiAddress, "test1")
+            return BigInt(20000)
+        })
+    });
+
+    it("should associate with evm transaction", async function(){
+        const addr = keys.test2
+        await fundSeiAddress(addr.seiAddress, "10000000000")
+        await fundAddress(addr.evmAddress, "200");
+
+        await verifyAssociation(addr.seiAddress, addr.evmAddress, async function(){
+            const txHash = await evmSend(addr.evmAddress, "test2", "0")
+            const receipt = await waitForReceipt(txHash)
+            return BigInt(receipt.gasUsed * (receipt.gasPrice / BigInt(1000000000000)))
+        })
+    });
+
+    it("should associate with associate transaction", async function(){
+        const addr = keys.test3
+        await fundSeiAddress(addr.seiAddress, "10000000000")
+        await fundAddress(addr.evmAddress, "200");
+
+        await verifyAssociation(addr.seiAddress, addr.evmAddress, async function(){
+            await associateKey("test3")
+            return BigInt(0)
+        })
+    });
+
+})

contracts/test/lib.js

@@ -10,27 +10,72 @@ async function delay() {
     await sleep(1000)
 }
 
-async function fundAddress(addr) {
-    return await execute(`seid tx evm send ${addr} 10000000000000000000 --from ${adminKeyName}`);
+async function fundAddress(addr, amount="10000000000000000000") {
+    const result = await evmSend(addr, adminKeyName, amount)
+    await delay()
+    return result
 }
 
-async function getAdmin() {
-    await associateAdmin()
-    const seiAddress = await getAdminSeiAddress()
+async function evmSend(addr, fromKey, amount="100000000000000000000000") {
+    const output = await execute(`seid tx evm send ${addr} ${amount} --from ${fromKey} -b block -y`);
+    return output.replace(/.*0x/, "0x").trim()
+}
+
+async function bankSend(toAddr, fromKey, amount="100000000000", denom="usei") {
+    const result = await execute(`seid tx bank send ${fromKey} ${toAddr} ${amount}${denom} -b block --fees 20000usei -y`);
+    await delay()
+    return result
+}
+
+async function fundSeiAddress(seiAddr, amount="100000000000", denom="usei") {
+    return await execute(`seid tx bank send ${adminKeyName} ${seiAddr} ${amount}${denom} -b block --fees 20000usei -y`);
+}
+
+async function getSeiBalance(seiAddr, denom="usei") {
+    const result = await execute(`seid query bank balances ${seiAddr} -o json`);
+    const balances = JSON.parse(result)
+    for(let b of balances.balances) {
+        if(b.denom === denom) {
+            return parseInt(b.amount, 10)
+        }
+    }
+    return 0
+}
+
+async function importKey(name, keyfile) {
+    try {
+        return await execute(`seid keys import ${name} ${keyfile}`, `printf "12345678\\n12345678\\n"`)
+    } catch(e) {
+        console.log("not importing key (skipping)")
+        console.log(e)
+    }
+}
+
+async function getNativeAccount(keyName) {
+    await associateKey(adminKeyName)
+    const seiAddress = await getKeySeiAddress(keyName)
+    await fundSeiAddress(seiAddress)
+    await delay()
     const evmAddress = await getEvmAddress(seiAddress)
     return {
         seiAddress,
         evmAddress
     }
 }
 
-async function getAdminSeiAddress() {
-    return (await execute(`seid keys show ${adminKeyName} -a`)).trim()
+async function getAdmin() {
+    await associateKey(adminKeyName)
+    return await getNativeAccount(adminKeyName)
+}
+
+async function getKeySeiAddress(name) {
+    return (await execute(`seid keys show ${name} -a`)).trim()
 }
 
-async function associateAdmin() {
+async function associateKey(keyName) {
     try {
-        return await execute(`seid tx evm associate-address --from ${adminKeyName}`)
+        await execute(`seid tx evm associate-address --from ${keyName} -b block`)
+        await delay()
     }catch(e){
         console.log("skipping associate")
     }
@@ -174,14 +219,14 @@ async function executeWasm(contractAddress, msg, coins = "0usei") {
     return JSON.parse(output);
 }
 
-async function execute(command) {
+async function execute(command, interaction=`printf "12345678\\n"`){
     return new Promise((resolve, reject) => {
         // Check if the Docker container 'sei-node-0' is running
         exec("docker ps --filter 'name=sei-node-0' --format '{{.Names}}'", (error, stdout, stderr) => {
             if (stdout.includes('sei-node-0')) {
                 // The container is running, modify the command to execute inside Docker
                 command = command.replace(/\.\.\//g, "/sei-protocol/sei-chain/");
-                const dockerCommand = `docker exec sei-node-0 /bin/bash -c 'export PATH=$PATH:/root/go/bin:/root/.foundry/bin && printf "12345678\\n" | ${command}'`;
+                const dockerCommand = `docker exec sei-node-0 /bin/bash -c 'export PATH=$PATH:/root/go/bin:/root/.foundry/bin && ${interaction} | ${command}'`;
                 execCommand(dockerCommand, resolve, reject);
             } else {
                 // The container is not running, execute command normally
@@ -205,8 +250,19 @@ function execCommand(command, resolve, reject) {
     });
 }
 
+async function waitForReceipt(txHash) {
+    let receipt = await ethers.provider.getTransactionReceipt(txHash)
+    while(!receipt) {
+        await delay()
+        receipt = await ethers.provider.getTransactionReceipt(txHash)
+    }
+    return receipt
+}
+
 module.exports = {
     fundAddress,
+    fundSeiAddress,
+    getSeiBalance,
     storeWasm,
     deployWasm,
     instantiateWasm,
@@ -220,4 +276,11 @@ module.exports = {
     deployEvmContract,
     deployErc20PointerForCw20,
     deployErc721PointerForCw721,
+    importKey,
+    getNativeAccount,
+    associateKey,
+    delay,
+    bankSend,
+    evmSend,
+    waitForReceipt,
 };

contracts/test/test1.key

@@ -0,0 +1,9 @@
+-----BEGIN TENDERMINT PRIVATE KEY-----
+kdf: bcrypt
+salt: 377151BDAA4DCDD8761F9489519D3E90
+type: secp256k1
+
+aMtUz1ow5FBmictYj67Wm5qJksJsiF0O86PrjOHnoINVBn0aCzYNxdkvNg0PP0Qb
+jHfcKaDPgk3PaeW16YqNJwveVX71A7lwpUmfT44=
+=46/d
+-----END TENDERMINT PRIVATE KEY-----

contracts/test/test2.key

@@ -0,0 +1,9 @@
+-----BEGIN TENDERMINT PRIVATE KEY-----
+kdf: bcrypt
+salt: 72E32FD4A4089E256F1EFCCBBA576592
+type: secp256k1
+
+MZ77pmO6izHdpkKVWYhGLGjZUAhf//WGURYVqUqhjJbCGtAXEkFKrjEAEpO3pyn1
+tYe8Hlry3BiqLoLY+j7CwE7W6g3WshJouXCrUiY=
+=VWqu
+-----END TENDERMINT PRIVATE KEY-----

contracts/test/test3.key

@@ -0,0 +1,9 @@
+-----BEGIN TENDERMINT PRIVATE KEY-----
+kdf: bcrypt
+salt: CDAEDE2709FFE1049C3E217BCD179BB2
+type: secp256k1
+
+l8hOkMcbc9ugxiUmS3k1wdziq6N+qhas25H0fjWD9tDK/BPn0QkY3pIz39TZJNuf
+/l9ImHhYxccWPoyTybKmXEjMaKLpFRZLdu+t844=
+=lvXk
+-----END TENDERMINT PRIVATE KEY-----

integration_test/evm_module/scripts/evm_tests.sh

@@ -2,3 +2,4 @@ cd contracts
 npm ci
 npx hardhat test --network seilocal test/EVMCompatabilityTest.js
 npx hardhat test --network seilocal test/EVMPrecompileTest.js
+npx hardhat test --network seilocal test/AssociateTest.js

x/evm/ante/preprocess.go

@@ -115,21 +115,7 @@ func (p *EVMPreprocessDecorator) AssociateAddresses(ctx sdk.Context, seiAddr sdk
 		}
 		p.accountKeeper.SetAccount(ctx, acc)
 	}
-	castAddr := sdk.AccAddress(evmAddr[:])
-	castAddrBalances := p.evmKeeper.BankKeeper().GetAllBalances(ctx, castAddr)
-	if !castAddrBalances.IsZero() {
-		if err := p.evmKeeper.BankKeeper().SendCoins(ctx, castAddr, seiAddr, castAddrBalances); err != nil {
-			return err
-		}
-	}
-	castAddrWei := p.evmKeeper.BankKeeper().GetWeiBalance(ctx, castAddr)
-	if !castAddrWei.IsZero() {
-		if err := p.evmKeeper.BankKeeper().SendCoinsAndWei(ctx, castAddr, seiAddr, sdk.ZeroInt(), castAddrWei); err != nil {
-			return err
-		}
-	}
-	p.evmKeeper.AccountKeeper().RemoveAccount(ctx, authtypes.NewBaseAccountWithAddress(castAddr))
-	return nil
+	return migrateBalance(ctx, p.evmKeeper, evmAddr, seiAddr)
 }
 
 // stateless
@@ -340,6 +326,24 @@ func NewEVMAddressDecorator(evmKeeper *evmkeeper.Keeper, accountKeeper *accountk
 	return &EVMAddressDecorator{evmKeeper: evmKeeper, accountKeeper: accountKeeper}
 }
 
+func migrateBalance(ctx sdk.Context, evmKeeper *evmkeeper.Keeper, evmAddr common.Address, seiAddr sdk.AccAddress) error {
+	castAddr := sdk.AccAddress(evmAddr[:])
+	castAddrBalances := evmKeeper.BankKeeper().GetAllBalances(ctx, castAddr)
+	if !castAddrBalances.IsZero() {
+		if err := evmKeeper.BankKeeper().SendCoins(ctx, castAddr, seiAddr, castAddrBalances); err != nil {
+			return err
+		}
+	}
+	castAddrWei := evmKeeper.BankKeeper().GetWeiBalance(ctx, castAddr)
+	if !castAddrWei.IsZero() {
+		if err := evmKeeper.BankKeeper().SendCoinsAndWei(ctx, castAddr, seiAddr, sdk.ZeroInt(), castAddrWei); err != nil {
+			return err
+		}
+	}
+	evmKeeper.AccountKeeper().RemoveAccount(ctx, authtypes.NewBaseAccountWithAddress(castAddr))
+	return nil
+}
+
 //nolint:revive
 func (p *EVMAddressDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) {
 	sigTx, ok := tx.(authsigning.SigVerifiableTx)
@@ -367,6 +371,10 @@ func (p *EVMAddressDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bo
 			continue
 		}
 		p.evmKeeper.SetAddressMapping(ctx, signer, evmAddr)
+		if err := migrateBalance(ctx, p.evmKeeper, evmAddr, signer); err != nil {
+			ctx.Logger().Error(fmt.Sprintf("failed to migrate EVM address balance (%s) %s", evmAddr.Hex(), err))
+			return ctx, err
+		}
 	}
 	return next(ctx, tx, simulate)
 }