[ReceiveExecutor] ref: store data in state machine

contracts/contracts/ccip/offramp/contract.tolk

@@ -850,19 +850,18 @@ fun onExecuteValidated(st: Storage, msg: OffRamp_ExecuteValidated, sender: addre
 
     val executorAddress: AutoDeployAddress = getReceiverExecutorDeployAddress(st, execId);
 
-    val initExecuteMsg = ReceiveExecutor_InitExecute {
+    var initExecuteMsg = ReceiveExecutor_InitExecute {
         gasOverride: msg.gasOverride,
         root: sender,
         sequenceNumber: message.header.sequenceNumber,
         sourceChainSelector: message.header.sourceChainSelector,
         messageId: message.header.messageId,
     };
 
-    var tokenAdminRegistry: Cell<address>?;
-
     // TODO: This is precarious. We should standarize a pattern to handle errors. Maybe all handlers should be wrapped in a try/catch
     try {
-        tokenAdminRegistry = getTokenAdminRegistry(st, message.tokenAmounts, message.header.sourceChainSelector);
+        val tokenAdminRegistry = getTokenAdminRegistry(st, message.tokenAmounts, message.header.sourceChainSelector);
+        initExecuteMsg.tokenAdminRegistry = tokenAdminRegistry;
     } catch {
         return messageExecutionFailed(initExecuteMsg);
     }
@@ -882,7 +881,6 @@ fun onExecuteValidated(st: Storage, msg: OffRamp_ExecuteValidated, sender: addre
                         message: message.toCell(),
                         root: sender,
                         execId,
-                        tokenAdminRegistry,
                     }.toCell(),
                 }
             },

contracts/contracts/ccip/receive_executor/contract.tolk

@@ -39,43 +39,52 @@ fun onInternalMessage(in: InMessage) {
         ReceiveExecutor_Confirm => {
             var st = ReceiveExecutor_Storage.load();
             assert(in.senderAddress == st.owner, ReceiveExecutor_Error.Unauthorized);
+            assert (st.state.messageExecution == ReceiveExecutor_MessageExecutionState.Execute) throw ReceiveExecutor_Error.UpdatingStateOfNonExecutedMessage;
             onCCIPReceiveConfirm(mutate st, msg);
             st.store();
         }
         // Sender must be the TokenAdminRegistry
         TokenRegistry_ReturnTokenInfo => {
             var st = ReceiveExecutor_Storage.load();
-            assert(st.tokenAdminRegistry != null, ReceiveExecutor_Error.Unauthorized);
-            val tokenAdminRegistry = st.tokenAdminRegistry.load();
-            assert(in.senderAddress == tokenAdminRegistry, ReceiveExecutor_Error.Unauthorized);
-            onTokenAdminRegistryResponse(mutate st, msg);
+            val ttInfo = st.state.tokenTransfer.loadOpt();
+            assert (ttInfo != null && ttInfo.state is ReceiveExecutor_TokenTransferState_TokenAdminRegistryQuery) throw ReceiveExecutor_Error.TokenAdminRegistryUnexpectedResponse;
+            assert(in.senderAddress == ttInfo.tokenAdminRegistry, ReceiveExecutor_Error.Unauthorized);
+            onTokenAdminRegistryResponse(mutate st, msg, ttInfo);
             st.store();
         }
         // Sender must be the TokenPool
         TokenPool_ReleaseOrMintFinished => {
             var st = ReceiveExecutor_Storage.load();
-            assert(st.tokenPool != null && in.senderAddress == st.tokenPool.load(), ReceiveExecutor_Error.Unauthorized);
-            onReleaseOrMintFinished(mutate st, msg);
+            val ttInfo = st.state.tokenTransfer.loadOpt();
+            assert (ttInfo != null && ttInfo.state is ReceiveExecutor_TokenTransferState_ReleaseOrMint) throw ReceiveExecutor_Error.TokenPoolUnexpectedResponse;
+            assert(in.senderAddress == ttInfo.state.tokenPool, ReceiveExecutor_Error.Unauthorized);
+            onReleaseOrMintFinished(mutate st, ttInfo, msg);
             st.store();
         }
         // Sender must be the TokenPool
         TokenPool_ReleaseOrMintFailure => {
             var st = ReceiveExecutor_Storage.load();
-            assert(st.tokenPool != null && in.senderAddress == st.tokenPool.load(), ReceiveExecutor_Error.Unauthorized);
-            onReleaseOrMintFailure(mutate st, msg);
+            val ttInfo = st.state.tokenTransfer.loadOpt();
+            assert (ttInfo != null && ttInfo.state is ReceiveExecutor_TokenTransferState_ReleaseOrMint) throw ReceiveExecutor_Error.TokenPoolUnexpectedResponse;
+            assert(in.senderAddress == ttInfo.state.tokenPool, ReceiveExecutor_Error.Unauthorized);
+            // SAFETY: TokenPool_ReleaseOrMintFailure.errorCode is uint16, so casting to int32 is safe and will not overflow.
+            onReleaseOrMintFailure(mutate st, ttInfo, ttInfo.state.tokenPool, msg.errorCode as int32);
             st.store();
         }
         // Sender must be the owner (OffRamp)
         ReleaseOrMint_ReleaseOrMintBounced => {
             var st = ReceiveExecutor_Storage.load();
             assert(in.senderAddress == st.owner, ReceiveExecutor_Error.Unauthorized);
-            onReleaseOrMintBounced(mutate st, msg);
+            val ttInfo = st.state.tokenTransfer.loadOpt();
+            assert (ttInfo != null && ttInfo.state is ReceiveExecutor_TokenTransferState_ReleaseOrMint) throw ReceiveExecutor_Error.TokenPoolUnexpectedResponse;
+            onReleaseOrMintFailure(mutate st, ttInfo, ttInfo.state.tokenPool, msg.exitCode);
             st.store();
         }
         // Sender must be the owner (OffRamp)
         ReceiveExecutor_Bounced => {
             var st = ReceiveExecutor_Storage.load();
             assert(in.senderAddress == st.owner, ReceiveExecutor_Error.Unauthorized);
+            assert (st.state.messageExecution == ReceiveExecutor_MessageExecutionState.Execute) throw ReceiveExecutor_Error.UpdatingStateOfNonExecutedMessage;
             onCCIPReceiveBounced(mutate st, msg);
             st.store();
         }
@@ -86,66 +95,68 @@ fun onInternalMessage(in: InMessage) {
     }
 }
 
-fun executeOrTransfer(mutate st: ReceiveExecutor_Storage, msg: ReceiveExecutor_InitExecute) {
-     if (st.tokenAdminRegistry != null) { // This message has a token transfer
-        // st.shouldOverrideGas = msg.gasOverride; // TODO
-        val ccipSend = lazy st.message.load();
-        // TODO data and token transfer not suported
-        assert (ccipSend.data.beginParse().isEmpty()) throw ReceiveExecutor_Error.PTTNotSupported;
-
-        return queryTokenAdminRegistry(mutate st, st.tokenAdminRegistry!.load());
-    }
-    execute(mutate st, msg.gasOverride);
-}
-
-fun onInitExecute(mutate st: ReceiveExecutor_Storage, msg: ReceiveExecutor_InitExecute) {
-    match (st.state) {
-        ReceiveExecutor_MessageState.Untouched => executeOrTransfer(mutate st, msg),
-        ReceiveExecutor_MessageState.TokenAdminRegistryQueryFailed => executeOrTransfer(mutate st, msg),
-        ReceiveExecutor_MessageState.ExecuteFailed => execute(mutate st, msg.gasOverride),
-        ReceiveExecutor_MessageState.TokenTransferFailed => {
-            val rampMsg = lazy st.message.load();
-            assert (rampMsg.tokenAmounts != null) throw ReceiveExecutor_Error.NoTokenAmountsInMessage;
-            val token = rampMsg.tokenAmounts.getTokenTransfer(ReceiveExecutor_Error.UnsupportedNumberOfTokens);
-            // TODO put conditional data such as tokenPool address in st.state
-            val tokenPool = st.tokenPool!.load();
-            releaseOrMint(mutate st, rampMsg, tokenPool, token);
+fun onInitExecute(mutate st: ReceiveExecutor_Storage, msg: ReceiveExecutor_InitExecute) { 
+    if (msg.tokenAdminRegistry != null) { // Message has token transfer
+        {
+            // st.shouldOverrideGas = msg.gasOverride; // TODO
+            val ccipSend = lazy st.message.load();
+            // TODO data and token transfer not supported
+            assert (ccipSend.data.beginParse().isEmpty()) throw ReceiveExecutor_Error.PTTNotSupported;
         }
-        ReceiveExecutor_MessageState.TokenAdminRegistryQuery => throw ReceiveExecutor_Error.ExecutionAlreadyInProgress,
-        ReceiveExecutor_MessageState.TokenTransfer => throw ReceiveExecutor_Error.ExecutionAlreadyInProgress,
-        ReceiveExecutor_MessageState.Execute => throw ReceiveExecutor_Error.ExecutionAlreadyInProgress,
-        ReceiveExecutor_MessageState.Success =>  throw ReceiveExecutor_Error.MessageAlreadyExecuted,
+        var ttInfo = st.state.tokenTransfer.loadOpt();
+        ttInfo = ttInfo.init(msg.tokenAdminRegistry.load());
+        if (ttInfo.state !is ReceiveExecutor_TokenTransferState_Success) {
+            return match (ttInfo.state) {
+                ReceiveExecutor_TokenTransferState_Untouched => queryTokenAdminRegistry(mutate st, ttInfo),
+                ReceiveExecutor_TokenTransferState_TokenAdminRegistryQueryFailed => queryTokenAdminRegistry(mutate st, ttInfo),
+                ReceiveExecutor_TokenTransferState_ReleaseOrMintFailed => {
+                    val rampMsg = lazy st.message.load();
+                    assert (rampMsg.tokenAmounts != null) throw ReceiveExecutor_Error.NoTokenAmountsInMessage;
+                    val token = rampMsg.tokenAmounts.getTokenTransfer(ReceiveExecutor_Error.UnsupportedNumberOfTokens);
+                    releaseOrMint(mutate st, ttInfo, ttInfo.state.tokenPool, rampMsg,token);
+                }
+                ReceiveExecutor_TokenTransferState_TokenAdminRegistryQuery => throw ReceiveExecutor_Error.ExecutionAlreadyInProgress,
+                ReceiveExecutor_TokenTransferState_ReleaseOrMint => throw ReceiveExecutor_Error.ExecutionAlreadyInProgress,
+            }
+        }
+    }
+    match (st.state.messageExecution) {
+        ReceiveExecutor_MessageExecutionState.Untouched => execute(mutate st, msg.gasOverride),
+        ReceiveExecutor_MessageExecutionState.Execute => throw ReceiveExecutor_Error.ExecutionAlreadyInProgress,
+        ReceiveExecutor_MessageExecutionState.ExecuteFailed => execute(mutate st, msg.gasOverride),
+        ReceiveExecutor_MessageExecutionState.Success =>  throw ReceiveExecutor_Error.MessageAlreadyExecuted,
     }
 }
 
-fun queryTokenAdminRegistry(mutate st: ReceiveExecutor_Storage, tokenAdminRegistry: address) {
-    st.state = ReceiveExecutor_MessageState.TokenAdminRegistryQuery;
+fun queryTokenAdminRegistry(mutate st: ReceiveExecutor_Storage, ttInfo: ReceiveExecutor_TokenTransferInfo) {
+    ttInfo.state = ReceiveExecutor_TokenTransferState_TokenAdminRegistryQuery {};
+    st.state.tokenTransfer = ttInfo.toCell();
     createMessage(
         {
             bounce: true,
             value: 0,
-            dest: tokenAdminRegistry,
+            dest: ttInfo.tokenAdminRegistry,
             body: TokenRegistry_GetTokenInfo {
             }
         }
     ).send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
 }
 
-fun onTokenAdminRegistryResponse(mutate st: ReceiveExecutor_Storage, msg: TokenRegistry_ReturnTokenInfo) {
-    assert(st.state == ReceiveExecutor_MessageState.TokenAdminRegistryQuery, ReceiveExecutor_Error.TokenAdminRegistryUnexpectedResponse);
+fun onTokenAdminRegistryResponse(mutate st: ReceiveExecutor_Storage, msg: TokenRegistry_ReturnTokenInfo, ttInfo: ReceiveExecutor_TokenTransferInfo) {
     val rampMsg = lazy st.message.load();
     assert (rampMsg.tokenAmounts != null) throw ReceiveExecutor_Error.NoTokenAmountsInMessage;
     val token = rampMsg.tokenAmounts.getTokenTransfer(ReceiveExecutor_Error.UnsupportedNumberOfTokens);
     if (msg.tokenPool == null) {
-        st.state = ReceiveExecutor_MessageState.TokenAdminRegistryQueryFailed;
+        ttInfo.state = ReceiveExecutor_TokenTransferState_TokenAdminRegistryQueryFailed {};
+        st.state.tokenTransfer = ttInfo.toCell();
         return endExecutionWithFailure(mutate st, rampMsg.header, ReceiveExecutor_Error.TokenNotEnabledInTokenRegistry);
     }
-    st.tokenPool = msg.tokenPool!.toCell();
-    releaseOrMint(mutate st, rampMsg, msg.tokenPool!, token);
+    releaseOrMint(mutate st, ttInfo, msg.tokenPool!, rampMsg, token!);
 }
 
-fun releaseOrMint(mutate st: ReceiveExecutor_Storage, rampMsg: Any2TVMRampMessage, tokenPool: address, tokenAmount: Any2TVMTokenTransfer) {
-    st.state = ReceiveExecutor_MessageState.TokenTransfer;
+fun releaseOrMint(mutate st: ReceiveExecutor_Storage, ttInfo: ReceiveExecutor_TokenTransferInfo, tokenPool: address, rampMsg: Any2TVMRampMessage, tokenAmount: Any2TVMTokenTransfer) {
+    ttInfo.state = ReceiveExecutor_TokenTransferState_ReleaseOrMint { tokenPool };
+    st.state.tokenTransfer = ttInfo.toCell();
     val releseOrMint = createMessage(
         {
             bounce: true,
@@ -179,7 +190,7 @@ fun releaseOrMint(mutate st: ReceiveExecutor_Storage, rampMsg: Any2TVMRampMessag
 
 fun execute(mutate st: ReceiveExecutor_Storage, gasOverride: coins?) {
     st.lastExecutionTimestamp = blockchain.now();
-    st.state = ReceiveExecutor_MessageState.Execute;
+    st.state.messageExecution = ReceiveExecutor_MessageExecutionState.Execute;
     createMessage(
         {
             bounce: true,
@@ -196,8 +207,6 @@ fun execute(mutate st: ReceiveExecutor_Storage, gasOverride: coins?) {
 
 // Finalizes message execution by returning all balance to the OffRamp and freezing the contract
 fun onCCIPReceiveConfirm(mutate st: ReceiveExecutor_Storage, msg: ReceiveExecutor_Confirm) {
-    assert(st.state == ReceiveExecutor_MessageState.Execute, ReceiveExecutor_Error.UpdatingStateOfNonExecutedMessage);
-
     val ccipMessage = lazy st.message.load();
 
     assert (msg.receiver == ccipMessage.receiver, ReceiveExecutor_Error.NotificationFromInvalidReceiver);
@@ -206,14 +215,12 @@ fun onCCIPReceiveConfirm(mutate st: ReceiveExecutor_Storage, msg: ReceiveExecuto
 }
 
 fun onCCIPReceiveBounced(mutate st: ReceiveExecutor_Storage, msg: ReceiveExecutor_Bounced){
-    assert(st.state == ReceiveExecutor_MessageState.Execute, ReceiveExecutor_Error.UpdatingStateOfNonExecutedMessage);
-
     val ccipMessage = lazy st.message.load();
 
     assert (msg.receiver == ccipMessage.receiver, ReceiveExecutor_Error.NotificationFromInvalidReceiver);
 
-    st.state = ReceiveExecutor_MessageState.ExecuteFailed;
-    endExecutionWithFailure(mutate st, ccipMessage.header, );
+    st.state.messageExecution = ReceiveExecutor_MessageExecutionState.ExecuteFailed;
+    endExecutionWithFailure(mutate st, ccipMessage.header);
 }
 
 fun endExecutionWithFailure(mutate st: ReceiveExecutor_Storage, messageHeader: RampMessageHeader, _errorCode: uint16 = 0) {
@@ -231,9 +238,9 @@ fun endExecutionWithFailure(mutate st: ReceiveExecutor_Storage, messageHeader: R
     notifyFailure.send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
 }
 
-fun onReleaseOrMintFinished(mutate st: ReceiveExecutor_Storage, msg: TokenPool_ReleaseOrMintFinished) {
-    assert(st.state == ReceiveExecutor_MessageState.TokenTransfer, ReceiveExecutor_Error.TokenPoolUnexpectedResponse);
-
+fun onReleaseOrMintFinished(mutate st: ReceiveExecutor_Storage, ttInfo: ReceiveExecutor_TokenTransferInfo, _msg: TokenPool_ReleaseOrMintFinished) {
+    ttInfo.state = ReceiveExecutor_TokenTransferState_Success {};
+    st.state.tokenTransfer = ttInfo.toCell();
     val ccipMessage = lazy st.message.load();
     if (ccipMessage.data.beginParse().isEmpty()) {
         return endExecutionSuccessfully(mutate st, ccipMessage.header);
@@ -244,7 +251,7 @@ fun onReleaseOrMintFinished(mutate st: ReceiveExecutor_Storage, msg: TokenPool_R
 }
 
 fun endExecutionSuccessfully(mutate st: ReceiveExecutor_Storage, messageHeader: RampMessageHeader) {
-    st.state = ReceiveExecutor_MessageState.Success;
+    st.state.messageExecution = ReceiveExecutor_MessageExecutionState.Success;
 
     val notifySuccess = createMessage({
         bounce: true,
@@ -259,20 +266,10 @@ fun endExecutionSuccessfully(mutate st: ReceiveExecutor_Storage, messageHeader:
     notifySuccess.send(SEND_MODE_CARRY_ALL_BALANCE);
 }
 
-// TBD do we want to forward the TokenPool_ReleaseOrMintFailure.errorCode ?
-fun onReleaseOrMintFailure(mutate st: ReceiveExecutor_Storage, _msg: TokenPool_ReleaseOrMintFailure) {
-    assert(st.state == ReceiveExecutor_MessageState.TokenTransfer, ReceiveExecutor_Error.TokenPoolUnexpectedResponse);
-
-    st.state = ReceiveExecutor_MessageState.TokenTransferFailed;
-    val ccipMessage = lazy st.message.load();
-    endExecutionWithFailure(mutate st, ccipMessage.header);
-}
-
-// TBD do we want to forward the ReleaseOrMint_ReleaseOrMintBounced.exitCode ?
-fun onReleaseOrMintBounced(mutate st: ReceiveExecutor_Storage, _msg: ReleaseOrMint_ReleaseOrMintBounced) {
-    assert(st.state == ReceiveExecutor_MessageState.TokenTransfer, ReceiveExecutor_Error.TokenPoolUnexpectedResponse);
-
-    st.state = ReceiveExecutor_MessageState.TokenTransferFailed;
+// TBD do we want to forward the exit code ?
+fun onReleaseOrMintFailure(mutate st: ReceiveExecutor_Storage, ttInfo: ReceiveExecutor_TokenTransferInfo, tokenPool: address, _error: int32) {
+    ttInfo.state = ReceiveExecutor_TokenTransferState_ReleaseOrMintFailed { tokenPool };
+    st.state.tokenTransfer = ttInfo.toCell();
     val ccipMessage = lazy st.message.load();
     endExecutionWithFailure(mutate st, ccipMessage.header);
 }

contracts/contracts/ccip/receive_executor/messages.tolk

@@ -20,6 +20,7 @@ struct (0x64cd2fd2) ReceiveExecutor_InitExecute {
     sequenceNumber: uint64;
     sourceChainSelector: uint64;
     messageId: uint256;
+    tokenAdminRegistry: Cell<address>? = null;
 }
 
 // Sent by the off-ramp to the receive executor when a release or mint operation has reverted.

contracts/contracts/ccip/receive_executor/storage.tolk

@@ -7,10 +7,12 @@ struct ReceiveExecutor_Storage {
     message: Cell<Any2TVMRampMessage>;
     root: address;
     execId: uint192;
-    state: ReceiveExecutor_MessageState = ReceiveExecutor_MessageState.Untouched;
+    state: ReceiveExecutor_State = ReceiveExecutor_State {
+        // tokenTransfer is initialized it in the handler `onInitExecute` if it is null.
+        tokenTransfer: null,
+        messageExecution: ReceiveExecutor_MessageExecutionState.Untouched,
+    };
     lastExecutionTimestamp: uint64 = 0;
-    tokenAdminRegistry: Cell<address>? = null;
-    tokenPool: Cell<address>? = null; // TODO move this to state attribute
 }
 
 fun ReceiveExecutor_Storage.load(): ReceiveExecutor_Storage {

contracts/contracts/ccip/receive_executor/types.tolk

@@ -5,13 +5,49 @@ const ReceiveExecutor_FACILITY_ID = ReceiveExecutor_FACILITY_NAME.crc32() % 640
 // NOTE: inlined getFacilityId because compiler doesn't allow custom compile time functions
 
 // State stores the general state machine indicating which phase we're in
-enum ReceiveExecutor_MessageState {
-    Untouched = 0
-    TokenAdminRegistryQuery
-    TokenAdminRegistryQueryFailed
-    TokenTransfer
+struct ReceiveExecutor_State {
+    tokenTransfer: Cell<ReceiveExecutor_TokenTransferInfo>?
+    messageExecution: ReceiveExecutor_MessageExecutionState
+}
+
+struct ReceiveExecutor_TokenTransferInfo {
+    tokenAdminRegistry: address
+    state : ReceiveExecutor_TokenTransferState = ReceiveExecutor_TokenTransferState_Untouched {};
+}
+
+fun ReceiveExecutor_TokenTransferInfo?.init(self, tokenAdminRegistry: address): ReceiveExecutor_TokenTransferInfo {
+    if (self != null) {
+        return self
+    }
+    return ReceiveExecutor_TokenTransferInfo { tokenAdminRegistry };
+}
+
+type ReceiveExecutor_TokenTransferState = 
+    | ReceiveExecutor_TokenTransferState_Success
+    | ReceiveExecutor_TokenTransferState_Pending;
+
+struct ReceiveExecutor_TokenTransferState_Success {}
+
+type ReceiveExecutor_TokenTransferState_Pending = 
+    | ReceiveExecutor_TokenTransferState_Untouched
+    | ReceiveExecutor_TokenTransferState_TokenAdminRegistryQuery
+    | ReceiveExecutor_TokenTransferState_TokenAdminRegistryQueryFailed
+    | ReceiveExecutor_TokenTransferState_ReleaseOrMint
+    | ReceiveExecutor_TokenTransferState_ReleaseOrMintFailed;
+
+struct ReceiveExecutor_TokenTransferState_Untouched {}
+struct ReceiveExecutor_TokenTransferState_TokenAdminRegistryQuery {}
+struct ReceiveExecutor_TokenTransferState_TokenAdminRegistryQueryFailed {}
+struct ReceiveExecutor_TokenTransferState_ReleaseOrMint {
+    tokenPool: address;
+}
+struct ReceiveExecutor_TokenTransferState_ReleaseOrMintFailed {
+    tokenPool: address;
+}
+
+enum ReceiveExecutor_MessageExecutionState {
+    Untouched
     Execute
     ExecuteFailed
-    TokenTransferFailed
     Success
-}
+}

contracts/contracts/lib/utils.tolk

@@ -248,6 +248,13 @@ fun T?.toOptCell(self): Cell<T>? {
     return self.toCell();
 }
 
+fun Cell<T>?.loadOpt(self): T? {
+    if (self == null) {
+        return null;
+    }
+    return self.load();
+}
+
 // This type works as a wrapper for RemainingBitsAndRefs that implements
 // automatic serialization/deserialization. It determines whether to store the
 // data directly as bits or as a reference to a cell, depending on available