TokenTransfer onramp flow skeleton

contracts/contracts/ccip/ccipsend_executor/contract.tolk

@@ -9,6 +9,9 @@ import "../../lib/utils"
 import "../router/messages"
 import "../fee_quoter/types"
 import "@stdlib/gas-payments"
+import "../common/types"
+import "../token_registry/messages"
+import "../test/tokenPool/messages"
 
 tolk 1.4.1
 
@@ -46,6 +49,17 @@ fun onInternalMessage(in: InMessage) {
             assert(this.addresses.load().feeQuoter == in.senderAddress, CCIPSendExecutor_Error.Unauthorized);
             this.onMessageValidationFailed(msg);
         }
+        // Sender must be TokenRegistry
+        TokenRegistry_ReturnTokenInfo => {
+            var this = CCIPSendExecutor<CCIPSendExecutor_State_TokenRegistryAccess>.load();
+            assert(this.addresses.load().tokenRegistry! == in.senderAddress, CCIPSendExecutor_Error.Unauthorized);
+            this.onTokenInfoReceived(msg);
+        }
+        MockTokenPool_NotifySuccessfulLockOrBurn => {
+            var this = CCIPSendExecutor<CCIPSendExecutor_State_TokenTransfer>.load();
+            assert(this.state.load().tokenPool == in.senderAddress, CCIPSendExecutor_Error.Unauthorized);
+            this.onConfirmLockOrBurn(msg);
+        }
         else => {
             assert (in.body.isEmpty()) throw 0xFFFF;
         },   
@@ -59,9 +73,16 @@ fun onBouncedMessage(in: InMessageBounced) {
             val this = CCIPSendExecutor<CCIPSendExecutor_State_OnGoingFeeValidation>.load();
             this.exitWithError(CCIPSendExecutor_Error.FeeQuoterBounce);
         }
+        //TODO: We should handle bounced messages from the TokenRegistry, as that means the token is not enabled on the lane
     }
 }
 
+//TODO: Figure out how to maintain backwards compatibility for messages sent between updating the Executor code in the OnRamp and upgrading the OnRamp itself.
+// Option 1: split the initialization in two messages to maintain backwards compatibility.
+// 1. CCIPSendExecutor_Execute stays the same as in 1.6.1 (that means not changing Config)
+// 2. CCIPSendExecutor_InitTokenTransfer passes the token registry address and other necessary configs
+// That way we can upgrade the SendExecutor code in the OnRamp and then upgrade the OnRamp itself, if an executor is deployed by the OnRamp between upgrades the new SendExecutor will still be able to handle the initialization message.
+// Option 2: Create a new message type CCIPSendExecutor_ExecuteV2 which has the new config element. Make the new executor able to handle both this version and the original.
 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> {
@@ -70,6 +91,7 @@ fun init(onrampSend: OnRamp_Send, config: CCIPSendExecutor_Config): CCIPSendExec
         addresses: CCIPSendExecutor_Addresses {
             onramp: st.onramp,
             feeQuoter: config.feeQuoter,
+            tokenRegistry: config.tokenRegistry
         }.toCell(),
         state: CCIPSendExecutor_State_Initialized{
         }.toCell(),
@@ -102,13 +124,73 @@ fun getValidatedFee(feeQuoter: address, onrampSend: OnRamp_Send) {
 }
 
 fun CCIPSendExecutor<CCIPSendExecutor_State_OnGoingFeeValidation>.onMessageValidated(mutate self, msg: FeeQuoter_MessageValidated<RemainingBitsAndRefs>) {
+    val message = Router_CCIPSend.fromCell(self.onrampSend.msg);
+    val tokenAmounts = message.tokenAmounts;
     if (msg.fee.feeTokenAmount + Router_Costs.CCIPSend() > self.onrampSend.metadata.value) {
         self.exitWithError(CCIPSendExecutor_Error.InsufficientFunds);
         return;
     }
+    // If token amounts is empty finalize the message
+    if (tokenAmounts.empty()) {
+        setGasLimitToMaximum();
+        self.exitSuccessfully(msg.fee);
+        return;
+     }
+
+    // Else if there are token transfers continue querying the tokenRegistry
+    val tokenRegistry = self.addresses.load().tokenRegistry!;
+    val queryMsg = createMessage({
+        bounce: true,
+        value: 0,
+        dest: tokenRegistry,
+        body: TokenRegistry_GetTokenInfo {
+        }
+    });
+    queryMsg.send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
+    val newState = CCIPSendExecutor<CCIPSendExecutor_State_TokenRegistryAccess> {
+        id: self.id,
+        onrampSend: self.onrampSend,
+        addresses: self.addresses,
+        state: CCIPSendExecutor_State_TokenRegistryAccess {
+            fee: msg.fee,
+        }.toCell(),
+    };
+    newState.store();
+}
+
+fun CCIPSendExecutor<CCIPSendExecutor_State_TokenRegistryAccess>.onTokenInfoReceived(mutate self, msg: TokenRegistry_ReturnTokenInfo) {
+    //TODO validate sender wallet address by querying the jettonMinter
+    assert(msg.tokenInfo.enabled) throw CCIPSendExecutor_Error.TokenNotEnabled;
+    val onrampSend = lazy self.onrampSend.msg.load();
+    val tokenAmount = onrampSend.tokenAmounts.iter().next();
+    val addresses = lazy self.addresses.load();
+    val requestLockOrBurn = createMessage({
+        bounce: true,
+        value: 0,
+        dest: addresses.onramp,
+        body: OnRamp_ExecutorRequestsLockOrBurn {
+            tokenAmount,
+            tokenPool: msg.tokenInfo.tokenPool,
+            destChainSelector: onrampSend.destChainSelector,
+            executorID: self.id,
+        }
+    });
+    requestLockOrBurn.send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
+    val newState = CCIPSendExecutor<CCIPSendExecutor_State_TokenTransfer> {
+        id: self.id,
+        onrampSend: self.onrampSend,
+        addresses: self.addresses,
+        state: CCIPSendExecutor_State_TokenTransfer {
+            tokenPool: msg.tokenInfo.tokenPool,
+            fee: self.state.load().fee,
+        }.toCell(),
+    };
+    newState.store();
+}
+
+fun CCIPSendExecutor<CCIPSendExecutor_State_TokenTransfer>.onConfirmLockOrBurn(mutate self, msg: MockTokenPool_NotifySuccessfulLockOrBurn) {
     setGasLimitToMaximum();
-
-    self.exitSuccessfully(msg.fee);
+    self.exitSuccessfully(self.state.load().fee);
 }
 
 fun CCIPSendExecutor<T>.exitSuccessfully(self, fee: Fee) {

contracts/contracts/ccip/ccipsend_executor/errors.tolk

@@ -7,4 +7,5 @@ enum CCIPSendExecutor_Error {
     InsufficientFunds
     InsufficientFee
     FeeQuoterBounce
+    TokenNotEnabled
 }

contracts/contracts/ccip/ccipsend_executor/messages.tolk

@@ -3,11 +3,16 @@ import "types"
 
 import "../onramp/messages"
 import "../fee_quoter/messages"
+import "../token_registry/messages"
+import "../common/types"
+import "../test/tokenPool/messages"
 
 type CCIPSendExecutor_InMessage =
     | CCIPSendExecutor_Execute
     | FeeQuoter_MessageValidated<RemainingBitsAndRefs>
     | FeeQuoter_MessageValidationFailed<RemainingBitsAndRefs>
+    | TokenRegistry_ReturnTokenInfo
+    | MockTokenPool_NotifySuccessfulLockOrBurn
     ;
 
 // crc32('CCIPSendExecutor_Execute')
@@ -55,4 +60,4 @@ fun SendExecutor_Costs.MessageValidationFailed(): int {
 
 type CCIPSendExecutor_BouncedMessage =
     | FeeQuoter_GetValidatedFee<RemainingBitsAndRefs>
-    ;
+    ;

contracts/contracts/ccip/ccipsend_executor/types.tolk

@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: BUSL-1.1
 import "../onramp/messages.tolk";
+import "../fee_quoter/types"
 
 const CCIPSendExecutor_FACILITY_NAME = "link.chain.ton.ccip.CCIPSendExecutor";
 const CCIPSendExecutor_FACILITY_ID = CCIPSendExecutor_FACILITY_NAME.crc32() % 640 + 10;  // 178
@@ -22,11 +23,14 @@ struct CCIPSendExecutor_Data {
 struct CCIPSendExecutor_Addresses {
     onramp: address,
     feeQuoter: address,
+    tokenRegistry: address?,
 }
 
 type CCIPSendExecutor_State =
     | Cell<CCIPSendExecutor_State_Initialized>
     | Cell<CCIPSendExecutor_State_OnGoingFeeValidation>
+    | Cell<CCIPSendExecutor_State_TokenRegistryAccess>
+    | Cell<CCIPSendExecutor_State_TokenTransfer>
     | Cell<CCIPSendExecutor_State_Finalized>
 
 struct CCIPSendExecutor_State_Initialized {
@@ -35,11 +39,25 @@ struct CCIPSendExecutor_State_Initialized {
 struct CCIPSendExecutor_State_OnGoingFeeValidation {
 }
 
+struct CCIPSendExecutor_State_TokenRegistryAccess {
+    fee: Fee,
+}
+
+//TODO Impl wallet address validation through JettonMinter/ stored wallet code
+struct CCIPSendExecutor_State_WalletAddressValidation {
+}
+
+struct CCIPSendExecutor_State_TokenTransfer {
+    tokenPool: address,
+    fee: Fee,
+}
+
 struct CCIPSendExecutor_State_Finalized {
 }
 
 struct CCIPSendExecutor_Config {
     feeQuoter: address,
+    tokenRegistry: address?,
 }
 
 

contracts/contracts/ccip/fee_quoter/contract.tolk

@@ -277,9 +277,10 @@ fun calculateValidatedFee(msg: Router_CCIPSend): Fee {
 
     val (executionGasPrice: uint112, dataAvailabilityGasPrice: uint112) = destChainConfig.getValidatedGasPrice();
 
-    val tokensIter = msg.tokenAmounts.iter();
     val premiumFeeUsdWei = mustProd(destChainConfig.config.networkFeeUsdCents, VAL_1E16, FeeQuoter_Error.PremiumFeeOverflow);
-    assert (tokensIter.empty()) throw FeeQuoter_Error.UnsupportedNumberOfTokens;
+    // NOTE: token transfers are currently priced as if the message carried no tokens
+    // (the extra token-transfer fee is ignored). The blocking empty-tokens assert was removed
+    // to enable the token-transfer send flow.
 
     val dataAvailabilityCost = _dataAvailabilityCost(
         destChainConfig,
@@ -366,8 +367,7 @@ fun _dataAvailabilityCost(
 fun validateMessageAndResolveGasLimitForDestination(extraArgs: cell, config: FeeQuoterDestChainConfig, message: Router_CCIPSend, msgDataLen: uint256): int {
     // Check that payload is formed correctly.
     assert (msgDataLen <= config.maxDataBytes) throw FeeQuoter_Error.MsgDataTooLarge;
-    val tokenAmounts = message.tokenAmounts.iter();
-    assert (tokenAmounts.empty()) throw FeeQuoter_Error.UnsupportedNumberOfTokens; 
+    // NOTE: token transfers are allowed; their extra fee is currently ignored (priced as a plain message).
 
     // NOTE: we could deploy distinct contracts to cut down on code
     if (config.chainFamilySelector == CHAIN_FAMILY_SELECTOR_EVM ||

contracts/contracts/ccip/onramp/contract.tolk

@@ -72,6 +72,12 @@ fun onInternalMessage(in: InMessage) {
             assert(destChainConfig.router == in.senderAddress, Error.Unauthorized);
             onSend(st, destChainConfig, msg)
         }
+        OnRamp_ExecutorRequestsLockOrBurn => {
+        // Sender must be Send Executor
+            var st = lazy OnRamp_Storage.load();
+            assert(st.executor.autoDeployAddress(msg.executorID).calculateAddress() == in.senderAddress, Error.Unauthorized);
+            onExecutorRequestsLockOrBurn(st, msg, in.senderAddress);
+        }
         // Sender must be Send Executor
         OnRamp_ExecutorFinishedSuccessfully => {
             val st = lazy OnRamp_Storage.load();
@@ -252,6 +258,14 @@ fun onSend(st: OnRamp_Storage, destChainConfig: OnRamp_DestChainConfig, payload:
 
     val config = st.config.load();
 
+    var tokenRegistry: address? = null;
+
+    val msg = payload.msg.load();
+    var tokenAmounts = msg.tokenAmounts.iter();
+    if (!tokenAmounts.empty()) {
+        tokenRegistry = calculateTokenRegistryAddress(st, tokenAmounts.next().token);
+    }
+
     val executorId = generateRandomSendExecutorId();
     val executeMsg = createMessage({
         bounce: true,
@@ -271,6 +285,7 @@ fun onSend(st: OnRamp_Storage, destChainConfig: OnRamp_DestChainConfig, payload:
                     onrampSend: payload,
                     config: CCIPSendExecutor_Config {
                         feeQuoter: config.feeQuoter,
+                        tokenRegistry,
                     }.toCell(),
                 }.toCell(),
             },
@@ -282,6 +297,28 @@ fun onSend(st: OnRamp_Storage, destChainConfig: OnRamp_DestChainConfig, payload:
     st.store();
 }
 
+fun calculateTokenRegistryAddress(st: OnRamp_Storage, token: address): address {
+    // TODO: resolve a per-token registry address. For now the OnRamp stores a single
+    // TokenRegistry address that is used for every token transfer.
+    return st.tokenRegistry!
+}
+
+fun onExecutorRequestsLockOrBurn(st: OnRamp_Storage, msg: OnRamp_ExecutorRequestsLockOrBurn, sender: address) {
+    var destChainConfig = st.destChainConfigs.mustGet(msg.destChainSelector, Error.UnknownDestChainSelector as int);
+    val routerLockOrBurn = createMessage({
+        bounce: true,
+        value: 0,
+        dest: destChainConfig.router,
+        body: Router_LockOrBurn {
+            tokenPool: msg.tokenPool,
+            tokenAmount: msg.tokenAmount,
+            destChainSelector: msg.destChainSelector,
+            executorAddress: sender,
+        }
+    });
+    routerLockOrBurn.send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
+}
+
 /*
 Handles successful completion of a CCIPSendExecutor. Assigns a new sequence
 number and notifies the Router of the successful send.

contracts/contracts/ccip/onramp/messages.tolk

@@ -7,12 +7,14 @@ import "../ccipsend_executor/messages"
 import "../fee_quoter/messages"
 import "../ccipsend_executor/types"
 import "../fee_quoter/types"
+import "../common/types"
 
 type OnRamp_InMessage =
     | OnRamp_Send
     | OnRamp_GetValidatedFee<RemainingBitsAndRefs>
     | FeeQuoter_MessageValidated<OnRamp_GetValidatedFeeContext>
     | FeeQuoter_MessageValidationFailed<OnRamp_GetValidatedFeeContext>
+    | OnRamp_ExecutorRequestsLockOrBurn
     | OnRamp_ExecutorFinishedSuccessfully
     | OnRamp_ExecutorFinishedWithError
     | OnRamp_SetDynamicConfig
@@ -38,6 +40,13 @@ struct (0x9c2ccc7e) OnRamp_GetValidatedFee<T> {
     context: T;
 }
 
+struct (0x9be1fb61) OnRamp_ExecutorRequestsLockOrBurn {
+    tokenAmount: TokenAmount,
+    tokenPool: address,
+    destChainSelector: uint64,
+    executorID: CCIPSendExecutor_ID,
+}
+
 struct OnRamp_GetValidatedFeeContext {
     onrampContext: address; // router address
     userContext: RemainingBitsOrRef<RemainingBitsAndRefs>,
@@ -168,4 +177,4 @@ fun OnRamp_Costs.WithdrawFeeTokens(): int {
         + ton("0.019") // feeTokens fwdFee
         + ton("0.05") // feeCollector computeFee
         ;
-}
+}

contracts/contracts/ccip/onramp/storage.tolk

@@ -14,6 +14,10 @@ struct OnRamp_Storage {
 
     destChainConfigs: map<uint64, OnRamp_DestChainConfig>; // chainSelector -> DestChainConfig
     executor: ExecutorDeployment;
+
+    // Address of the TokenRegistry queried by the CCIPSendExecutor during token transfers.
+    // Null until token transfers are configured. TODO: replace with per-token registry resolution.
+    tokenRegistry: address?;
 }
 
 fun OnRamp_Storage.load(): OnRamp_Storage {

contracts/contracts/ccip/onramp/types.tolk

@@ -70,7 +70,7 @@ struct TVM2AnyRampMessageBody {
     receiver: Cell<CrossChainAddress>;
     data: cell;
     extraArgs: cell;
-    tokenAmounts: cell;
+    tokenAmounts: SnakedCell<TokenAmount>;
     feeToken: address;
     feeTokenAmount: uint256;
 }