Update tests and pool modules to handle v1 and v2 ccip receive callback upgrade

Author: JohnChangUK

contracts/ccip/ccip/sources/receiver_dispatcher.move

@@ -11,18 +11,13 @@ module ccip::receiver_dispatcher {
     ) {
         auth::assert_is_allowed_offramp(signer::address_of(caller));
 
-        let dispatch_metadata =
-            receiver_registry::start_receive(receiver_address, message);
-        dispatchable_fungible_asset::derived_supply(dispatch_metadata);
-        receiver_registry::finish_receive(receiver_address);
-    }
-
-    /// Invoke receiver's callback without token dispatchable hooks
-    public fun dispatch_receive_v2(
-        caller: &signer, receiver_address: address, message: client::Any2AptosMessage
-    ) {
-        auth::assert_is_allowed_offramp(signer::address_of(caller));
-
-        receiver_registry::invoke_ccip_receive_v2(receiver_address, message);
+        if (receiver_registry::is_registered_receiver_v2(receiver_address)) {
+            receiver_registry::invoke_ccip_receive_v2(receiver_address, message);
+        } else {
+            let dispatch_metadata =
+                receiver_registry::start_receive(receiver_address, message);
+            dispatchable_fungible_asset::derived_supply(dispatch_metadata);
+            receiver_registry::finish_receive(receiver_address);
+        }
     }
 }

contracts/ccip/ccip/sources/receiver_registry.move

@@ -34,7 +34,7 @@ module ccip::receiver_registry {
 
     struct CCIPReceiverRegistrationV2 has key {
         callback: |client::Any2AptosMessage| has drop + copy + store,
-        proof_typeinfo: TypeInfo,
+        proof_typeinfo: TypeInfo
     }
 
     #[event]
@@ -167,7 +167,7 @@ module ccip::receiver_registry {
         );
 
         move_to(
-            receiver_account, 
+            receiver_account,
             CCIPReceiverRegistrationV2 { callback, proof_typeinfo }
         );
 
@@ -181,6 +181,7 @@ module ccip::receiver_registry {
     #[view]
     public fun is_registered_receiver(receiver_address: address): bool {
         exists<CCIPReceiverRegistration>(receiver_address)
+            || exists<CCIPReceiverRegistrationV2>(receiver_address)
     }
 
     #[view]
@@ -221,9 +222,7 @@ module ccip::receiver_registry {
         registration.dispatch_metadata
     }
 
-    public(friend) fun finish_receive(
-        receiver_address: address
-    ) acquires CCIPReceiverRegistration {
+    public(friend) fun finish_receive(receiver_address: address) acquires CCIPReceiverRegistration {
         let registration = get_registration_mut(receiver_address);
 
         assert!(
@@ -233,17 +232,15 @@ module ccip::receiver_registry {
     }
 
     public(friend) fun invoke_ccip_receive_v2(
-        receiver_address: address,
-        message: client::Any2AptosMessage
+        receiver_address: address, message: client::Any2AptosMessage
     ) acquires CCIPReceiverRegistrationV2 {
         assert!(
             exists<CCIPReceiverRegistrationV2>(receiver_address),
             error::invalid_argument(E_UNKNOWN_RECEIVER)
         );
 
         let registration = borrow_global<CCIPReceiverRegistrationV2>(receiver_address);
-
-        (registration.callback)(message);
+        (registration.callback) (message);
     }
 
     inline fun borrow_state(): &ReceiverRegistryState {

contracts/ccip/ccip_offramp/sources/offramp.move

@@ -808,14 +808,8 @@ module ccip_offramp::offramp {
         // module.
         // ref: https://github.com/smartcontractkit/chainlink-ccip/blob/875e982e6437dc126710d8224dd7c792a197bea6/chains/evm/contracts/offRamp/OffRamp.sol#L633
 
-        let is_v1_receiver = receiver_registry::is_registered_receiver(message.receiver);
-        let is_v2_receiver =
-            receiver_registry::is_registered_receiver_v2(message.receiver);
-
-        if ((!message.data.is_empty()
-            || message.gas_limit != 0)
-            && (is_v1_receiver
-                || is_v2_receiver)) {
+        if ((!message.data.is_empty() || message.gas_limit != 0)
+            && receiver_registry::is_registered_receiver(message.receiver)) {
             let state_signer =
                 account::create_signer_with_capability(&state.state_signer_cap);
 
@@ -833,16 +827,9 @@ module ccip_offramp::offramp {
                     dest_token_amounts
                 );
 
-            // Use V2 dispatch if available, else V1
-            if (is_v2_receiver) {
-                receiver_dispatcher::dispatch_receive_v2(
-                    &state_signer, message.receiver, any2aptos_message
-                )
-            } else {
-                receiver_dispatcher::dispatch_receive(
-                    &state_signer, message.receiver, any2aptos_message
-                )
-            }
+            receiver_dispatcher::dispatch_receive(
+                &state_signer, message.receiver, any2aptos_message
+            )
         };
 
     }
@@ -909,7 +896,7 @@ module ccip_offramp::offramp {
             account::create_signer_with_capability(&state.state_signer_cap);
 
         let (fa, local_amount) =
-            token_admin_dispatcher::dispatch_release_or_mint_v2(
+            token_admin_dispatcher::dispatch_release_or_mint(
                 &state_signer,
                 token_pool_address,
                 sender,

contracts/ccip/ccip_offramp/tests/mock/mock_ccip_receiver.move

@@ -4,9 +4,10 @@
 module ccip_offramp::mock_ccip_receiver {
     use std::account;
     use std::event;
-    use std::object::{Self};
+    use std::object::{Self, Object};
     use std::string::{Self, String};
-    use std::fungible_asset::{Metadata};
+    use std::fungible_asset::{Self, Metadata};
+    use std::option::{Self, Option};
     use std::primary_fungible_store;
     use std::from_bcs;
     use std::signer;
@@ -71,10 +72,30 @@ module ccip_offramp::mock_ccip_receiver {
             }
         );
 
+        // Default to V2 registration
         receiver_registry::register_receiver_v2(
             publisher,
             MODULE_NAME,
-            |message| ccip_receive(message),
+            |message| ccip_receive_v2(message),
+            CCIPReceiverProof {}
+        );
+    }
+
+    /// Register this receiver as V1 (dispatchable fungible asset mode)
+    /// This is used for testing V1 compatibility
+    public fun register_as_v1(publisher: &signer) {
+        receiver_registry::register_receiver(publisher, MODULE_NAME, CCIPReceiverProof {});
+    }
+
+    /// Migrate from V1 to V2 registration
+    /// This demonstrates the upgrade path from dispatchable FA to closures
+    public fun migrate_to_v2(publisher: &signer) {
+        // V2 registration will coexist with V1
+        // The dispatcher will prefer V2 when both exist
+        receiver_registry::register_receiver_v2(
+            publisher,
+            MODULE_NAME,
+            |message| ccip_receive_v2(message),
             CCIPReceiverProof {}
         );
     }
@@ -88,7 +109,83 @@ module ccip_offramp::mock_ccip_receiver {
 
     struct CCIPReceiverProof has drop {}
 
-    public fun ccip_receive(message: client::Any2AptosMessage) acquires CCIPReceiverState {
+    /// This function should only be used with non-dispatchable tokens,
+    /// as it is currently incompatible with dispatchable tokens.
+    public fun ccip_receive<T: key>(_metadata: Object<T>): Option<u128> acquires CCIPReceiverState {
+        /* load state and rebuild a signer for the resource account */
+        let state = borrow_global_mut<CCIPReceiverState>(@ccip_offramp);
+        let state_signer = account::create_signer_with_capability(&state.signer_cap);
+
+        let message =
+            receiver_registry::get_receiver_input(@ccip_offramp, CCIPReceiverProof {});
+
+        let data = client::get_data(&message);
+
+        let dest_token_amounts = client::get_dest_token_amounts(&message);
+
+        if (dest_token_amounts.length() != 0 && data.length() != 0) {
+            let final_recipient = from_bcs::to_address(data);
+
+            for (i in 0..dest_token_amounts.length()) {
+                let token_amount_ref = &dest_token_amounts[i];
+                let token_addr = client::get_token(token_amount_ref);
+                let amount = client::get_amount(token_amount_ref);
+
+                // Implement the token transfer logic here
+
+                let fa_token = object::address_to_object<Metadata>(token_addr);
+                let fa_store_sender =
+                    primary_fungible_store::ensure_primary_store_exists(
+                        @ccip_offramp, fa_token
+                    );
+                let fa_store_receiver =
+                    primary_fungible_store::ensure_primary_store_exists(
+                        final_recipient, fa_token
+                    );
+
+                fungible_asset::transfer(
+                    &state_signer,
+                    fa_store_sender,
+                    fa_store_receiver,
+                    amount
+                );
+            };
+
+            event::emit(ForwardedTokens { final_recipient });
+            event::emit_event(
+                &mut state.forwarded_tokens_handle, ForwardedTokens { final_recipient }
+            );
+
+        } else if (data.length() != 0) {
+
+            // Convert the vector<u8> to a string
+            let message = string::utf8(data);
+
+            event::emit(ReceivedMessage { message });
+            event::emit_event(
+                &mut state.received_message_handle, ReceivedMessage { message }
+            );
+
+        } else if (dest_token_amounts.length() != 0) {
+            // Tokens only (no forwarding data) - keep them at receiver
+            // Emit event to prove receiver was called
+            let token_count = dest_token_amounts.length();
+            event::emit(ReceivedTokensOnly { token_count });
+            event::emit_event(
+                &mut state.received_tokens_only_handle,
+                ReceivedTokensOnly { token_count }
+            );
+        };
+
+        // Simple abort condition for testing
+        if (data == b"abort") {
+            abort 1
+        };
+
+        option::none()
+    }
+
+    public fun ccip_receive_v2(message: client::Any2AptosMessage) acquires CCIPReceiverState {
         /* load state and rebuild a signer for the resource account */
         let state = borrow_global_mut<CCIPReceiverState>(@ccip_offramp);
         let state_signer = account::create_signer_with_capability(&state.signer_cap);
@@ -171,6 +268,31 @@ module ccip_offramp::mock_ccip_receiver {
         init_module(publisher);
     }
 
+    /// Initialize without auto-registering (for testing V1/V2 manually)
+    public fun test_init_state_only(publisher: &signer) {
+        // Create a signer capability for the receiver account
+        let signer_cap = account::create_test_signer_cap(signer::address_of(publisher));
+
+        // Create a unique handle for each event type
+        let received_message_handle =
+            account::new_event_handle<ReceivedMessage>(publisher);
+        let forwarded_tokens_handle =
+            account::new_event_handle<ForwardedTokens>(publisher);
+        let received_tokens_only_handle =
+            account::new_event_handle<ReceivedTokensOnly>(publisher);
+
+        // Move all state into the single resource struct
+        move_to(
+            publisher,
+            CCIPReceiverState {
+                signer_cap,
+                received_message_handle,
+                forwarded_tokens_handle,
+                received_tokens_only_handle
+            }
+        );
+    }
+
     public fun get_received_message_events(): vector<ReceivedMessage> acquires CCIPReceiverState {
         let state = borrow_global<CCIPReceiverState>(@ccip_offramp);
         event::emitted_events_by_handle<ReceivedMessage>(&state.received_message_handle)

contracts/ccip/ccip_token_pools/burn_mint_token_pool/sources/burn_mint_token_pool.move

@@ -85,6 +85,8 @@ module burn_mint_token_pool::burn_mint_token_pool {
         let lock_or_burn_closure = |fa, input| lock_or_burn_v2(fa, input);
         let release_or_mint_closure = |input| release_or_mint_v2(input);
 
+        // If the contract has already been deployed with V1 and needs to be upgraded to V2,
+        // create a new module
         token_admin_registry::register_pool_v2(
             publisher,
             token_pool_module_name,
@@ -810,4 +812,61 @@ module burn_mint_token_pool::burn_mint_token_pool {
             &borrow_global<BurnMintTokenPoolState>(state).token_pool_state
         )
     }
+
+    #[test_only]
+    /// Used for registering the pool with V2 closure-based callbacks.
+    public fun create_callback_proof(): CallbackProof {
+        CallbackProof {}
+    }
+
+    #[test_only]
+    public fun test_init_v1(publisher: &signer) {
+        // register the pool on deployment, because in the case of object code deployment,
+        // this is the only time we have a signer ref to @ccip_burn_mint_pool.
+        assert!(
+            object::object_exists<Metadata>(@burn_mint_local_token),
+            error::invalid_argument(E_INVALID_FUNGIBLE_ASSET)
+        );
+        let metadata = object::address_to_object<Metadata>(@burn_mint_local_token);
+
+        // create an Account on the object for event handles.
+        account::create_account_if_does_not_exist(@burn_mint_token_pool);
+
+        // the name of this module. if incorrect, callbacks will fail to be registered and
+        // register_pool will revert.
+        let token_pool_module_name = b"burn_mint_token_pool";
+
+        // Register the entrypoint with mcms
+        if (@mcms_register_entrypoints == @0x1) {
+            register_mcms_entrypoint(publisher, token_pool_module_name);
+        };
+
+        token_admin_registry::register_pool(
+            publisher,
+            token_pool_module_name,
+            @burn_mint_local_token,
+            CallbackProof {}
+        );
+
+        // create a resource account to be the owner of the primary FungibleStore we will use.
+        let (store_signer, store_signer_cap) =
+            account::create_resource_account(publisher, STORE_OBJECT_SEED);
+
+        // make sure this is a valid fungible asset that is primary fungible store enabled,
+        // ie. created with primary_fungible_store::create_primary_store_enabled_fungible_asset
+        primary_fungible_store::ensure_primary_store_exists(
+            signer::address_of(&store_signer), metadata
+        );
+
+        move_to(
+            publisher,
+            BurnMintTokenPoolDeployment {
+                store_signer_cap,
+                ownable_state: ownable::new(&store_signer, @burn_mint_token_pool),
+                token_pool_state: token_pool::initialize(
+                    &store_signer, @burn_mint_local_token, vector[]
+                )
+            }
+        );
+    }
 }