[NONEVM-2617] Token Registry / OnRamp flow

contracts/contracts/ccip/ccipsend_executor/contract.tolk

@@ -2,6 +2,8 @@ import "../types"
 import "messages"
 import "storage"
 import "../fee_quoter/messages"
+import "../token_registry/messages"
+import "../token_registry/storage"
 
 tolk 1.1
 
@@ -19,6 +21,11 @@ fun onInternalMessage(in: InMessage) {
             // TODO match refunds from TokenPool
             this.onJettonTransferFromOnRamp(msg, in.senderAddress);
         }
+        TokenRegistry_TokenPoolInfo => {
+            var this = CCIPSendExecutor<CCIPSendExecutor_State_OnGoingTokenRegistryQuery>.load();
+            assert (in.senderAddress == this.state.load().tokenRegistry) throw 0xFFFF; // TODO propper error
+            this.onTokenPoolInfo(msg);
+        }
         FeeQuoter_MessageValidated<Metadata> => {
             var this = CCIPSendExecutor<CCIPSendExecutor_State_OnGoingFeeValidation>.load();
             this.onMessageValidated(msg, in.senderAddress);
@@ -31,7 +38,7 @@ fun onInternalMessage(in: InMessage) {
 
 fun init(onramp_send: OnRampSend, config: CCIPSendExecutor_Config): CCIPSendExecutor<CCIPSendExecutor_State_Initialized> {
     val st = lazy CCIPSendExecutor_InitialData.fromCell(contract.getData());
-    return CCIPSendExecutor<CCIPSendExecutor_State_Initialized> {
+    val this = CCIPSendExecutor<CCIPSendExecutor_State_Initialized> {
         messageId: st.messageId,
         onrampSend: onramp_send,
         addresses: CCIPSendExecutor_Addresses {
@@ -42,6 +49,7 @@ fun init(onramp_send: OnRampSend, config: CCIPSendExecutor_Config): CCIPSendExec
             tokenRegistry: config.tokenRegistry,
         }.toCell(),
     };
+    return this;
 }
 
 fun CCIPSendExecutor<CCIPSendExecutor_State_Initialized>.onExecute(mutate self, onrampJettonWallet: address?) {
@@ -63,6 +71,7 @@ fun CCIPSendExecutor<CCIPSendExecutor_State_Initialized>.onExecute(mutate self,
 
 fun CCIPSendExecutor<CCIPSendExecutor_State_Initialized>.withdrawJettons(mutate self, onrampJettonWallet: address) {
     val ccipSend = lazy self.onrampSend.msg.load();
+    val state = lazy self.state.load();
 
     val withdrawMsg = createMessage({
         bounce: true,
@@ -80,24 +89,51 @@ fun CCIPSendExecutor<CCIPSendExecutor_State_Initialized>.withdrawJettons(mutate
         onrampSend: self.onrampSend,
         addresses: self.addresses,
         state: CCIPSendExecutor_State_WaitingForJettons {
-            tokenRegistry: self.state.load().tokenRegistry!,
+            tokenRegistry: state.tokenRegistry!,
         }.toCell(),
     };
     newState.store();
 }
 
 fun CCIPSendExecutor<CCIPSendExecutor_State_WaitingForJettons>.onJettonTransferFromOnRamp(mutate self, msg: JettonTransferNotification, jettonWallet: address) {
-    // TODO query token registry for token pool
-    getValidatedFee(self.addresses.load().feeQuoter, self.onrampSend); // temporary
+    val tokenRegistry = self.state.load().tokenRegistry;
+    val queryMsg = createMessage({
+        bounce: true,
+        value: 0,
+        dest: tokenRegistry,
+        body: TokenRegistry_GetTokenInfo {
+            queryId: 0, // TODO: what id to use?
+            ccipSend: self.onrampSend.msg,
+            executorJettonWallet: jettonWallet,
+        }
+    });
+    queryMsg.send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
+    val newState = CCIPSendExecutor<CCIPSendExecutor_State_OnGoingTokenRegistryQuery> {
+        messageId: self.messageId,
+        onrampSend: self.onrampSend,
+        addresses: self.addresses,
+        state: CCIPSendExecutor_State_OnGoingTokenRegistryQuery {
+            tokenRegistry,
+            jettonWallet: jettonWallet,
+        }.toCell(),
+    };
+    newState.store();
+}
+
+fun CCIPSendExecutor<CCIPSendExecutor_State_OnGoingTokenRegistryQuery>.onTokenPoolInfo(mutate self, tokenPoolInfo: TokenRegistry_TokenPoolInfo) {
+    val addresses = lazy self.addresses.load();
+    getValidatedFee(addresses.feeQuoter, self.onrampSend);
+
+    val state = self.state.load();
     val newState = CCIPSendExecutor<CCIPSendExecutor_State_OnGoingFeeValidation> {
         messageId: self.messageId,
         onrampSend: self.onrampSend,
         addresses: self.addresses,
-        state: CCIPSendExecutor_State_OnGoingFeeValidation { // TODO change to OnGoingTokenRegistryQuery
+        state: CCIPSendExecutor_State_OnGoingFeeValidation {
             pendingJettonLock: CCIPSendExecutor_PendingJettonLock {
-                tokenRegistry: self.state.load().tokenRegistry,
-                jettonWallet: jettonWallet,
-                tokenPool: createAddressNone(), // TODO
+                tokenRegistry: state.tokenRegistry,
+                jettonWallet: state.jettonWallet,
+                tokenPool: tokenPoolInfo.tokenPool,
             },
         }.toCell(),
     };

contracts/contracts/ccip/ccipsend_executor/errors.tolk

@@ -1,3 +1,3 @@
-tolk 1.0
+tolk 1.1
 
 const CCIPSEND_EXECUTOR_ERROR_STATE_NOT_EXPECTED = 500;

contracts/contracts/ccip/ccipsend_executor/messages.tolk

@@ -1,6 +1,7 @@
 import "../types"
 import "storage"
 import "../fee_quoter/messages"
+import "../token_registry/messages"
 
 tolk 1.1
 
@@ -20,5 +21,5 @@ type CCIPSendExecutor_InMessage =
     | CCIPSendExecutor_Execute
     | FeeQuoter_MessageValidated<Metadata>
     | JettonTransferNotification
-    // TODO | TokenRegistry_TokenPoolInfo
+    | TokenRegistry_TokenPoolInfo
     // TODO Implement interrupt message that allows to cancel the pending CCIP send and return funds to the sender. This should be callable only after a certain timeout. 

contracts/contracts/ccip/ccipsend_executor/storage.tolk

@@ -24,9 +24,9 @@ struct CCIPSendExecutor_Addresses {
 type CCIPSendExecutor_State =
     | Cell<CCIPSendExecutor_State_Initialized>
     | Cell<CCIPSendExecutor_State_WaitingForJettons>
-    // TODO | Cell<CCIPSendExecutor_State_OnGoingTokenRegistryQuery>
+    | Cell<CCIPSendExecutor_State_OnGoingTokenRegistryQuery>
     | Cell<CCIPSendExecutor_State_OnGoingFeeValidation>
-    // TODO | Cell<CCIPSendExecutor_State_OnGoingLockOrBurn>
+    | Cell<CCIPSendExecutor_State_OnGoingLockOrBurn>
 
 struct CCIPSendExecutor_State_Initialized {
     tokenRegistry: address?,
@@ -36,11 +36,10 @@ struct CCIPSendExecutor_State_WaitingForJettons {
     tokenRegistry: address,
 }
 
-// TODO
-// struct CCIPSendExecutor_State_OnGoingTokenRegistryQuery {
-//     tokenRegistry: address,
-//     jettonWallet: address,
-// }
+struct CCIPSendExecutor_State_OnGoingTokenRegistryQuery {
+    tokenRegistry: address,
+    jettonWallet: address,
+}
 
 struct CCIPSendExecutor_State_OnGoingFeeValidation {
     pendingJettonLock: CCIPSendExecutor_PendingJettonLock?
@@ -52,9 +51,8 @@ struct CCIPSendExecutor_PendingJettonLock {
     tokenPool: address,
 }
 
-// TODO
-// struct CCIPSendExecutor_State_OnGoingLockOrBurn {
-// }
+struct CCIPSendExecutor_State_OnGoingLockOrBurn {
+}
 
 // Casted state
 
@@ -78,7 +76,7 @@ fun CCIPSendExecutor<T>.load(): CCIPSendExecutor<T> {
 
 fun CCIPSendExecutor<T>.store(self) {
     contract.setData(CCIPSendExecutor_Data {
-        messageId: 0,
+        messageId: self.messageId,
         onrampSend: self.onrampSend,
         addresses: self.addresses,
         state: self.state,

contracts/contracts/ccip/onramp.tolk

@@ -8,6 +8,7 @@ import "ccipsend_executor/storage"
 import "ccipsend_executor/messages"
 import "../lib/jetton/messages"
 import "../lib/jetton/messages_extended"
+import "token_registry/storage"
 
 const CCIP_MESSAGE_SENT_TOPIC: int = stringCrc32("CCIPMessageSent"); 
 
@@ -33,6 +34,7 @@ struct OnRamp_Data {
 
     destChainConfigs: map<uint64, DestChainConfig>; // chainSelector -> DestChainConfig
     executor_code: cell; // code for CCIPSendExecutor
+    token_registry_code: cell; // code for TokenRegistry
     currentMessageId: uint224;
     // TODO track last destroyed executor+gaps to allow owner to reclaim funds?
 }
@@ -66,6 +68,7 @@ type OnRamp_InMessage =
     | SetDynamicConfig
     | OnRampUpdateDestChainConfigs
     | UpdateAllowlists
+    | ReturnExcessesBack
     ;
     // TODO | UpdateExecutorCode
 
@@ -96,6 +99,9 @@ fun onInternalMessage(in: InMessage) {
             applyAllowlistUpdates(mutate st, msg.updates);
             st.store();
         }
+        ReturnExcessesBack => {
+            return;
+        }
         else => {
             // ignore empty messages, "wrong opcode" for others
             assert (in.body.isEmpty()) throw 0xFFFF
@@ -130,7 +136,6 @@ fun send(payload: OnRampSend, sender: address, jettonWallet: address? = null) {
     if (!tokenAmounts.empty()) {
         tokenRegistry = calculateTokenRegistryAddress(st, tokenAmounts.next().token);
     }
-
     val executeMsg = createMessage({
         bounce: true,
         value: ton("0.05"), // TODO: calculate
@@ -178,8 +183,17 @@ fun onTransferNotification(notification: JettonTransferNotification, jettonWalle
 }
 
 fun calculateTokenRegistryAddress(st: OnRamp_Data, token: address): address {
-    // TODO: calculate token registry address
-    return createAddressNone()
+    return address.fromValidBuilder(AutoDeployAddress {
+        workchain: BASECHAIN,
+        stateInit: ContractState {
+            code: st.token_registry_code,
+            data: TokenRegistry_Storage {
+                onramp: contract.getAddress(),
+                minterAddress: token,
+            }.toCell(),
+        },
+        toShard: null,
+    }.buildAddress())
 }
 
 fun handleWithdrawJettons(payload: OnRamp_WithdrawJettons, executor: address) {

contracts/contracts/ccip/token_registry/contract.tolk

@@ -0,0 +1,75 @@
+import "messages"
+import "storage"
+import "../../lib/jetton/jetton_client"
+import "../../lib/jetton/utils"
+
+tolk 1.1;
+
+fun onInternalMessage(in: InMessage) {
+    val msg = lazy TokenRegistry_InMessage.fromSlice(in.body);
+
+    match (msg) {
+        TokenRegistry_SetInfo => {
+            var st = TokenRegistry_Storage.load();
+            TokenRegistry{
+                data: TokenRegistry_InitializedStorage {
+                    onramp: st.onramp,
+                    minterAddress: st.minterAddress,
+                    info: msg.info,
+                },
+            }.store();
+        }
+        TokenRegistry_GetTokenInfo => {
+            val this = TokenRegistry.load();
+            this.handleGetTokenInfo(msg, in.senderAddress);
+        }
+        else => {
+            assert (in.body.isEmpty()) throw 10; // TODO propper error
+        }
+    }
+}
+
+struct TokenRegistry {
+    data: TokenRegistry_InitializedStorage
+}
+
+fun TokenRegistry.load(): TokenRegistry {
+    return TokenRegistry {
+        data: TokenRegistry_InitializedStorage.load()
+    };
+}
+
+fun TokenRegistry.store(self) {
+    return self.data.store();
+}
+
+fun TokenRegistry.handleGetTokenInfo(self, msg: TokenRegistry_GetTokenInfo, sender: address) {
+    if (!self.data.info.enabled) {
+        // TODO: send message blocking transaction
+        debug.printString("TokenRegistry: Registry is disabled");
+        throw 501;
+        return;
+    }
+
+    if (msg.executorJettonWallet != calculateUserJettonWalletAddress(
+        sender,
+        self.data.minterAddress,
+        self.data.info.walletCode,
+    )) {
+        // TODO: Send message blocking transaction
+        debug.printString("TokenRegistry: Invalid executor jetton wallet");
+        throw 502;
+        return;
+    }
+
+    val response = createMessage({
+        dest: sender,
+        value: 0,
+        bounce: true,
+        body: TokenRegistry_TokenPoolInfo {
+            queryId: msg.queryId,
+            tokenPool: self.data.info.tokenPool,
+        }
+    });
+    response.send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
+}

contracts/contracts/ccip/token_registry/messages.tolk

@@ -0,0 +1,29 @@
+import "storage"
+import "../types"
+
+tolk 1.1
+
+type TokenRegistry_InMessage = 
+    | TokenRegistry_SetInfo
+    | TokenRegistry_GetTokenInfo
+
+
+// crc32('TokenRegistry_SetInfo')
+struct (0x22E22393)TokenRegistry_SetInfo {
+    queryId: uint64
+    info: TokenRegistry_TokenInfo
+}
+
+// crc32('TokenRegistry_GetTokenInfo')
+struct (0xDD5D5127)TokenRegistry_GetTokenInfo {
+    queryId: uint64
+    ccipSend: Cell<CCIPSend>
+    executorJettonWallet: address
+}
+
+// Outgoing messages
+
+struct (0xB8F88A81)TokenRegistry_TokenPoolInfo {
+    queryId: uint64,
+    tokenPool: address,
+}

contracts/contracts/ccip/token_registry/storage.tolk

@@ -0,0 +1,42 @@
+tolk 1.1
+
+struct TokenRegistry_Storage {
+    onramp: address // TODO this could be MCMS?
+    minterAddress: address
+    info: TokenRegistry_TokenInfo? = null
+    // mintAuthority: address? TODO add mint authority for when the token is delegated / self served token pools
+}
+
+struct TokenRegistry_TokenInfo {
+    tokenPool: address
+    walletCode: cell
+    enabled: bool
+}
+
+struct TokenRegistry_InitializedStorage {
+    onramp: address
+    minterAddress: address
+    info: TokenRegistry_TokenInfo,
+}
+
+fun TokenRegistry_Storage.load(): TokenRegistry_Storage {
+    return TokenRegistry_Storage.fromCell(contract.getData());
+}
+
+fun TokenRegistry_InitializedStorage.load(): TokenRegistry_InitializedStorage {
+    val st = TokenRegistry_Storage.fromCell(contract.getData());
+    assert (st.info != null) throw 0xFFFF; // TODO propper error
+    return TokenRegistry_InitializedStorage {
+        onramp: st.onramp,
+        minterAddress: st.minterAddress,
+        info: st.info,
+    };
+}
+
+fun TokenRegistry_InitializedStorage.store(self) {
+    contract.setData(TokenRegistry_Storage{
+        onramp: self.onramp,
+        minterAddress: self.minterAddress,
+        info: self.info,
+    }.toCell());
+}