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

Author: JohnChangUK

bindings/ccip/token_admin_registry/token_admin_registry.go

@@ -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_dispatcher.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/sources/receiver_registry.move

@@ -16,7 +16,6 @@ module ccip::token_admin_dispatcher {
     ): (vector<u8>, vector<u8>) {
         auth::assert_is_allowed_onramp(signer::address_of(caller));
 
-        // Check for V2 first (preferred - uses closures)
         if (token_admin_registry::has_token_pool_config(token_pool_address)) {
             token_admin_registry::lock_or_burn_v2(
                 caller,
@@ -93,55 +92,5 @@ module ccip::token_admin_dispatcher {
             (fa, destination_amount)
         }
     }
-
-    // ============================================
-    // V2 Closure-Based Dispatch Functions
-    // ============================================
-
-    public fun dispatch_lock_or_burn_v2(
-        caller: &signer,
-        token_pool_address: address,
-        fa: FungibleAsset,
-        sender: address,
-        remote_chain_selector: u64,
-        receiver: vector<u8>
-    ): (vector<u8>, vector<u8>) {
-        token_admin_registry::lock_or_burn_v2(
-            caller,
-            token_pool_address,
-            fa,
-            sender,
-            remote_chain_selector,
-            receiver
-        )
-    }
-
-    public fun dispatch_release_or_mint_v2(
-        caller: &signer,
-        token_pool_address: address,
-        sender: vector<u8>,
-        receiver: address,
-        source_amount: u256,
-        local_token: address,
-        remote_chain_selector: u64,
-        source_pool_address: vector<u8>,
-        source_pool_data: vector<u8>,
-        offchain_token_data: vector<u8>
-    ): (FungibleAsset, u64) {
-        let (fa, destination_amount) =
-            token_admin_registry::release_or_mint_v2(
-                caller,
-                token_pool_address,
-                sender,
-                receiver,
-                source_amount,
-                local_token,
-                remote_chain_selector,
-                source_pool_address,
-                source_pool_data,
-                offchain_token_data
-            );
-
-        (fa, destination_amount)
-    }
 }
+

contracts/ccip/ccip/sources/token_admin_dispatcher.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,