have PDA reimburse the ATA creation fee;

assuming the withdrawal already paid at least this much in ZRC20 withdraw fee.
zeta-chain · Dec 3, 2024 · 9118d4c · 9118d4c
1 parent 7fb0a66
commit 9118d4c
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 41 deletions.
diff --git a/programs/protocol-contracts-solana/src/lib.rs b/programs/protocol-contracts-solana/src/lib.rs
@@ -338,6 +338,7 @@ pub mod gateway {
         message_hash: [u8; 32],
         nonce: u64,
     ) -> Result<()> {
+
         let pda = &mut ctx.accounts.pda;
         // let program_id = &mut ctx.accounts
         if nonce != pda.nonce {
@@ -386,6 +387,9 @@ pub mod gateway {
             Errors::SPLAtaAndMintAddressMismatch,
         );
 
+        let cost_gas = 5000; // default gas cost in lamports
+        let cost_ata_create = &mut 2039280; // default SPL ATA account creation rent fee in lamports
+
         // test whether the recipient_ata is created or not; if not, create it
         let recipient_ata_account = ctx.accounts.recipient_ata.to_account_info();
         if recipient_ata_account.lamports() == 0
@@ -397,8 +401,7 @@ pub mod gateway {
                 recipient_ata_account.key(),
                 ctx.accounts.recipient.key(),
             );
-            let signer_info = &ctx.accounts.signer.to_account_info();
-            let bal_before = signer_info.lamports();
+            let bal_before = ctx.accounts.signer.lamports();
             invoke(
                 &create_associated_token_account(
                     ctx.accounts.signer.to_account_info().key,
@@ -419,22 +422,15 @@ pub mod gateway {
                         .clone(),
                 ],
             )?;
-            let bal_after = signer_info.lamports();
+            let bal_after = ctx.accounts.signer.lamports();
+            *cost_ata_create = bal_before - bal_after;
 
             msg!("Associated token account for recipient created!");
             msg!(
-                "Refunding the rent paid by the signer {:?}",
+                "Refunding the rent ({:?} lamports) paid by the signer {:?}",
+                cost_ata_create,
                 ctx.accounts.signer.to_account_info().key
             );
-
-            let rent_payer_info = ctx.accounts.rent_payer_pda.to_account_info();
-            let cost = bal_before - bal_after;
-            rent_payer_info.sub_lamports(cost)?;
-            signer_info.add_lamports(cost)?;
-            msg!(
-                "Signer refunded the ATA account creation rent amount {:?} lamports",
-                cost,
-            );
         }
 
         let xfer_ctx = CpiContext::new_with_signer(
@@ -452,6 +448,12 @@ pub mod gateway {
 
         transfer_checked(xfer_ctx, amount, decimals)?;
         msg!("withdraw spl token successfully");
+        // Note: this pda.sub_lamports() must be done here due to this issue https://github.com/solana-labs/solana/issues/9711
+        // otherwise the previous CPI calls might fail with error:
+        // "sum of account balances before and after instruction do not match"
+        let reimbursement = cost_gas + *cost_ata_create;
+        pda.sub_lamports(reimbursement)?;
+        ctx.accounts.signer.add_lamports(reimbursement)?;
 
         Ok(())
     }

diff --git a/tests/protocol-contracts-solana.ts b/tests/protocol-contracts-solana.ts
@@ -343,33 +343,6 @@ describe("Gateway", () => {
 
     });
 
-    it("withdraw SPL token to a non-existent account should succeed by creating it", async () => {
-        let seeds = [Buffer.from("rent-payer", "utf-8")];
-        const [rentPayerPda] = anchor.web3.PublicKey.findProgramAddressSync(
-            seeds,
-            gatewayProgram.programId,
-        );
-        let rentPayerPdaBal0 = await conn.getBalance(rentPayerPda);
-        let pda_ata = await spl.getAssociatedTokenAddress(mint.publicKey, pdaAccount, true);
-        const pdaAccountData = await gatewayProgram.account.pda.fetch(pdaAccount);
-        const hexAddr = bufferToHex(Buffer.from(pdaAccountData.tssAddress));
-        const amount = new anchor.BN(500_000);
-        const nonce = pdaAccountData.nonce;
-        const wallet2 = anchor.web3.Keypair.generate();
-
-        const to = await spl.getAssociatedTokenAddress(mint.publicKey, wallet2.publicKey);
-
-        let to_ata_bal = await conn.getBalance(to);
-        expect(to_ata_bal).to.be.eq(0); // the new ata account (owned by wallet2) should be non-existent;
-        const txsig = await withdrawSplToken(mint, usdcDecimals, amount, nonce, pda_ata, to, wallet2.publicKey,  keyPair, gatewayProgram);
-        to_ata_bal = await conn.getBalance(to);
-        expect(to_ata_bal).to.be.gt(2_000_000); // the new ata account (owned by wallet2) should be created
-
-        // rent_payer_pda should have reduced balance
-
-        let rentPayerPdaBal1 = await conn.getBalance(rentPayerPda);
-        expect(rentPayerPdaBal0-rentPayerPdaBal1).to.be.eq(to_ata_bal); // rentPayer pays rent
-    });
 
     it("fails to deposit if receiver is empty address", async() => {
         try {
@@ -415,7 +388,38 @@ describe("Gateway", () => {
         let bal3 = await conn.getBalance(to);
         expect(bal3).to.be.gte(500_000_000);
     })
-
+
+    it("withdraw SPL token to a non-existent account should succeed by creating it", async () => {
+        let seeds = [Buffer.from("rent-payer", "utf-8")];
+        const [rentPayerPda] = anchor.web3.PublicKey.findProgramAddressSync(
+            seeds,
+            gatewayProgram.programId,
+        );
+        let rentPayerPdaBal0 = await conn.getBalance(pdaAccount);
+        let pda_ata = await spl.getAssociatedTokenAddress(mint.publicKey, pdaAccount, true);
+        const pdaAccountData = await gatewayProgram.account.pda.fetch(pdaAccount);
+        const hexAddr = bufferToHex(Buffer.from(pdaAccountData.tssAddress));
+        const amount = new anchor.BN(500_000);
+        const nonce = pdaAccountData.nonce;
+        const wallet2 = anchor.web3.Keypair.generate();
+
+
+        const to = await spl.getAssociatedTokenAddress(mint.publicKey, wallet2.publicKey);
+
+        let to_ata_bal = await conn.getBalance(to);
+        expect(to_ata_bal).to.be.eq(0); // the new ata account (owned by wallet2) should be non-existent;
+        const txsig = await withdrawSplToken(mint, usdcDecimals, amount, nonce, pda_ata, to, wallet2.publicKey,  keyPair, gatewayProgram);
+        to_ata_bal = await conn.getBalance(to);
+        expect(to_ata_bal).to.be.gt(2_000_000); // the new ata account (owned by wallet2) should be created
+
+        // rent_payer_pda should have reduced balance
+
+        let rentPayerPdaBal1 = await conn.getBalance(pdaAccount);
+        // expected reimbursement to be gas fee (5000 lamports) + ATA creation cost 2039280 lamports
+        expect(rentPayerPdaBal0-rentPayerPdaBal1).to.be.eq(to_ata_bal + 5000); // rentPayer pays rent
+    });
+
+
     it("fails to deposit and call if receiver is empty address", async() => {
         try {
             await gatewayProgram.methods.depositAndCall(new anchor.BN(1_000_000_000), Array(20).fill(0), Buffer.from("hello", "utf-8")).accounts({}).rpc();