[NONEVM-3009] [Testing] Contracts don't lose balance on fees

contracts/contracts/ccip/ccipsend_executor/contract.tolk

@@ -54,7 +54,7 @@ fun onBouncedMessage(in: InMessageBounced) {
 fun init(onrampSend: OnRamp_Send, config: CCIPSendExecutor_Config): CCIPSendExecutor<CCIPSendExecutor_State_Initialized> {
     val st = lazy CCIPSendExecutor_InitialData.fromCell(contract.getData());
     return CCIPSendExecutor<CCIPSendExecutor_State_Initialized> {
-        messageID: st.messageId,
+        messageID: st.messageID,
         onrampSend: onrampSend,
         addresses: CCIPSendExecutor_Addresses {
             onramp: st.onramp,
@@ -165,7 +165,7 @@ fun CCIPSendExecutor<T>.exitSuccessfully(self, fee: coins) {
         value: 0,
         dest: self.addresses.load().onramp,
         body: OnRamp_ExecutorFinishedSuccessfully {
-            msgId: self.messageID,
+            msgID: self.messageID,
             msg: self.onrampSend.msg,
             metadata: self.onrampSend.metadata,
             fee,

contracts/contracts/ccip/ccipsend_executor/types.tolk

@@ -2,7 +2,7 @@ import "../onramp/messages.tolk";
 
 struct CCIPSendExecutor_InitialData {
     onramp: address,
-    messageId: uint224,
+    messageID: uint224,
 }
 
 struct CCIPSendExecutor_Data {

contracts/contracts/ccip/onramp/contract.tolk

@@ -92,16 +92,7 @@ fun send(payload: OnRamp_Send, sender: address, jettonWallet: address? = null) {
     val executeMsg = createMessage({
         bounce: true,
         value: 0,
-        dest: AutoDeployAddress {
-            stateInit: ContractState {
-                code: st.executor_code,
-                data: CCIPSendExecutor_InitialData {
-                    onramp: contract.getAddress(),
-                    messageId: st.currentMessageId,
-                }.toCell(),
-            }
-            // TODO use toShard so these contracts live in the same shard as the onramp
-        },
+        dest: executorAddress(st.executorCode, st.currentMessageId),
         body: CCIPSendExecutor_Execute {
             onrampSend: payload,
             config: CCIPSendExecutor_Config {
@@ -118,9 +109,22 @@ fun send(payload: OnRamp_Send, sender: address, jettonWallet: address? = null) {
     st.store();
 }
 
+@inline
+fun executorAddress(executorCode: cell, messageID: int): AutoDeployAddress {
+    return AutoDeployAddress {
+        stateInit: ContractState {
+            code: executorCode,
+            data: CCIPSendExecutor_InitialData {
+                onramp: contract.getAddress(),
+                messageID: messageID
+            }.toCell(),
+        }
+    }
+}
+
 fun commit(payload: OnRamp_ExecutorFinishedSuccessfully, sender: address) {
     var st = lazy OnRamp_Storage.load();
-    // TODO validate sender is executor msg.msgId
+    assert(executorAddress(st.executorCode, payload.msgID).addressMatches(sender)) throw Error.Unauthorized;
 
     val ccipsend: Router_CCIPSend = payload.msg.load();
     val metadata = payload.metadata;
@@ -187,7 +191,7 @@ fun commit(payload: OnRamp_ExecutorFinishedSuccessfully, sender: address) {
 
 fun replyWithError(payload: OnRamp_ExecutorFinishedWithError, sender: address) {
     var st = lazy OnRamp_Storage.load();
-    // TODO validate sender is executor msg.msgId
+    assert(executorAddress(st.executorCode, payload.msgID).addressMatches(sender)) throw Error.Unauthorized;
 
     val ccipsend: Router_CCIPSend = payload.msg.load();
     val metadata = payload.metadata;

contracts/contracts/ccip/onramp/messages.tolk

@@ -42,7 +42,7 @@ struct (0x10000003) OnRamp_SetDynamicConfig {
 
 // crc32('OnRamp_ExecutorFinishedSuccessfully')
 struct (0xCFA6B336) OnRamp_ExecutorFinishedSuccessfully {
-    msgId: uint224,
+    msgID: uint224,
     msg: Cell<Router_CCIPSend>
     metadata: Metadata
     fee: coins

contracts/contracts/ccip/onramp/storage.tolk

@@ -19,7 +19,7 @@ struct OnRamp_Storage {
     config: Cell<OnRamp_DynamicConfig>;
 
     destChainConfigs: map<uint64, OnRamp_DestChainConfig>; // chainSelector -> DestChainConfig
-    executor_code: cell; // code for CCIPSendExecutor
+    executorCode: cell; // code for CCIPSendExecutor
     currentMessageId: uint224;
 }
 

contracts/tests/ccip/CCIPRouter.spec.ts

@@ -1,4 +1,4 @@
-import { Blockchain, BlockchainTransaction, SandboxContract, TreasuryContract } from '@ton/sandbox'
+import { Blockchain, printTransactionFees, SandboxContract, TreasuryContract } from '@ton/sandbox'
 import {
   toNano,
   Address,
@@ -28,7 +28,8 @@ import { newWithdrawableSpec } from '../lib/funding/WithdrawableSpec'
 import * as ownable2step from '../../wrappers/libraries/access/Ownable2Step'
 import * as UpgradeableSpec from '../lib/versioning/UpgradeableSpec'
 import * as TypeAndVersionSpec from '../lib/versioning/TypeAndVersionSpec'
-import { dump } from '../utils/prettyPrint'
+import { dump, prettifyAddressesMap } from '../utils/prettyPrint'
+import { mapOpcode } from '../utils/opcodes'
 
 const CHAINSEL_EVM_TEST_90000001 = 909606746561742123n
 const CHAINSEL_EVM_TEST_90000002 = 5548718428018410741n
@@ -531,87 +532,9 @@ describe('Router', () => {
     }
   })
 
-  it('doesnt lose balance on messageSent fees', async () => {
-    const initialOnRampBalance = (await blockchain.getContract(onRamp.address)).balance
-
-    const ccipSend: rt.CCIPSend = {
-      queryID: 1,
-      destChainSelector: CHAINSEL_EVM_TEST_90000001,
-      receiver: EVM_ADDRESS,
-      data: Cell.EMPTY,
-      tokenAmounts: [],
-      feeToken: TEST_TOKEN_ADDR,
-      extraArgs: rt.builder.data.extraArgs
-        .encode({
-          kind: 'generic-v2',
-          gasLimit: 100n,
-          allowOutOfOrderExecution: true,
-        })
-        .asCell(),
-    }
-
-    const originalSentValue = toNano('0.5')
-    const valueFromExecutor = toNano('0.4')
-    const ccipFee = toNano('0.01')
-    const result = await onRamp.sendExecutorFinishedSuccessfully(deployer.getSender(), {
-      value: valueFromExecutor,
-      body: {
-        messageID: 42n,
-        msg: rt.builder.message.in.ccipSend.encode(ccipSend).asCell(),
-        metadata: {
-          sender: deployer.address,
-          value: originalSentValue,
-        },
-        fee: ccipFee,
-      },
-    })
-
-    expect(result.transactions).toHaveTransaction({
-      from: deployer.address,
-      to: onRamp.address,
-      success: true,
-    })
-
-    expect(result.transactions).toHaveTransaction({
-      from: onRamp.address,
-      to: router.address,
-      success: true,
-      op: rt.Opcodes.messageSent,
-    })
-
-    expect(result.transactions).toHaveTransaction({
-      from: router.address,
-      to: deployer.address,
-      success: true,
-      op: rt.OutgoingOpcodes.ccipSendACK,
-    })
-
-    const finalOnRampBalance = (await blockchain.getContract(onRamp.address)).balance
-
-    const relayTX = result.transactions.find((tx) => {
-      return (
-        tx.inMessage != null &&
-        tx.inMessage != undefined &&
-        tx.inMessage.info.src != null &&
-        tx.inMessage.info.src != undefined &&
-        tx.inMessage.info.src instanceof Address &&
-        tx.inMessage.info.src.equals(deployer.address) &&
-        tx.inMessage.info.dest != null &&
-        tx.inMessage.info.dest != undefined &&
-        tx.inMessage.info.dest instanceof Address &&
-        tx.inMessage.info.dest.equals(onRamp.address) &&
-        tx.description.type === 'generic'
-      )
-    }) as BlockchainTransaction & {
-      inMessage: Message & { info: CommonMessageInfoInternal }
-      description: TransactionDescriptionGeneric
-    }
-    const rentFee = relayTX.description.storagePhase?.storageFeesCollected ?? 0n
-
-    expect(finalOnRampBalance).toBe(initialOnRampBalance - rentFee + ccipFee)
-  })
-
   it('onramp arbitrary message passing', async () => {
+    // Track initial balance to verify fees are handled correctly
+    const initialOnRampBalance = (await blockchain.getContract(onRamp.address)).balance
     const ccipSend: rt.CCIPSend = {
       queryID: 1,
       destChainSelector: CHAINSEL_EVM_TEST_90000001,
@@ -753,6 +676,85 @@ describe('Router', () => {
           })
         },
       })
+
+      printTransactionFees(result.transactions, mapOpcode)
+      const addresses = prettifyAddressesMap(result.transactions)
+
+      result.transactions.forEach((tx) => {
+        if (
+          tx.inMessage &&
+          tx.inMessage.info.type === 'internal' &&
+          tx.description.type === 'generic'
+        ) {
+          const inValue = tx.inMessage.info.value.coins
+          const outValue = tx.outMessages
+            .values()
+            .reduce(
+              (acc, msg) => acc + (msg.info.type === 'internal' ? msg.info.value.coins : 0n),
+              0n,
+            )
+
+          const fees = {
+            inFwdFee: tx.inMessage.info.forwardFee,
+            gasFees:
+              tx.description.computePhase.type === 'vm' ? tx.description.computePhase.gasFees : 0n,
+            actionFees: tx.description.actionPhase?.totalActionFees ?? 0n,
+            fwdFees: tx.description.actionPhase?.totalFwdFees ?? 0n,
+            storageFees: tx.description.storagePhase?.storageFeesCollected ?? 0n,
+          }
+          const totalFees = [fees.actionFees, fees.gasFees, fees.storageFees].reduce(
+            (a, b) => a + b,
+            0n,
+          )
+
+          console.log(
+            `Balance check for tx from ${addresses.get(tx.inMessage.info.src.toRawString())} to ${addresses.get(tx.inMessage.info.dest.toRawString())}:\n`,
+          )
+          // table format
+          console.table({
+            'In Value': inValue,
+            'Out Value': outValue,
+            'In Fwd Fee': fees.inFwdFee,
+            'Gas Fees': fees.gasFees,
+            'Action Fees': fees.actionFees,
+            'Fwd Fees': fees.fwdFees,
+            'Storage Fees': fees.storageFees,
+            'Total Fees': totalFees,
+            'In Value - Out Value - Fees': inValue - outValue - totalFees,
+          })
+        }
+      })
+
+      // Verify balance handling: OnRamp doesn't lose balance on messageSent fees
+      const finalOnRampBalance = (await blockchain.getContract(onRamp.address)).balance
+      const rentFees = result.transactions
+        .filter((tx) => {
+          return (
+            tx.inMessage != null &&
+            tx.inMessage != undefined &&
+            tx.inMessage.info.dest != null &&
+            tx.inMessage.info.dest != undefined &&
+            tx.inMessage.info.dest instanceof Address &&
+            tx.inMessage.info.dest.equals(router.address)
+          )
+        })
+        .reduce((acc, tx) => {
+          switch (tx.description.type) {
+            case 'generic': {
+              const rentFee = tx.description.storagePhase?.storageFeesCollected ?? 0n
+              return acc + rentFee
+            }
+            case 'storage': {
+              const rentFee = tx.description.storagePhase.storageFeesCollected
+              return acc + rentFee
+            }
+          }
+          return acc
+        }, 0n)
+
+      // The final balance should be initial balance minus rent fees plus the fee that was paid
+      // (the fee comes from the validated fee calculation above)
+      expect(finalOnRampBalance).toBe(initialOnRampBalance - rentFees + amount.fee)
     }
   })
 

contracts/tests/utils/opcodes.ts

@@ -0,0 +1,34 @@
+import * as rt from '../../wrappers/ccip/Router'
+import * as onr from '../../wrappers/ccip/OnRamp'
+import * as fq from '../../wrappers/ccip/FeeQuoter'
+import * as sx from '../../wrappers/ccip/CCIPSendExecutor'
+// import * as rx from '../../wrappers/ccip/CCIPReceiveExecutor'
+import * as offr from '../../wrappers/ccip/OffRamp'
+
+// Create a comprehensive opcode mapping
+const createOpcodeMapping = () => {
+  const mapping: Record<number, string> = {}
+
+  for (const [ops, name] of [
+    [rt.Opcodes, 'Router.in'],
+    [rt.OutgoingOpcodes, 'Router.out'],
+    [onr.Opcodes, 'OnRamp.in'],
+    [fq.Opcodes, 'FeeQuoter.in'],
+    [sx.Opcodes, 'SendExecutor.in'],
+    // [rx.Opcodes, 'ReceiveExecutor.in'],
+    [offr.Opcodes, 'OffRamp.in'],
+  ]) {
+    for (const [key, value] of Object.entries(ops)) {
+      mapping[value as number] = `${name}.${key}`
+    }
+  }
+
+  return mapping
+}
+
+export const OPCODE_MAPPING = createOpcodeMapping()
+
+// Useful to use with printTransactionFees from @ton/sandbox
+export const mapOpcode = (op: number): string | undefined => {
+  return OPCODE_MAPPING[op]
+}