Skip to content

Commit be0161b

Browse files
committed
Update tests and pool modules to handle v1 and v2 ccip receive callback upgrade
1 parent 5e02abf commit be0161b

15 files changed

Lines changed: 650 additions & 45 deletions

File tree

bindings/ccip/token_admin_registry/token_admin_registry.go

Lines changed: 31 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

contracts/ccip/ccip/sources/receiver_dispatcher.move

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,18 +11,13 @@ module ccip::receiver_dispatcher {
1111
) {
1212
auth::assert_is_allowed_offramp(signer::address_of(caller));
1313

14-
let dispatch_metadata =
15-
receiver_registry::start_receive(receiver_address, message);
16-
dispatchable_fungible_asset::derived_supply(dispatch_metadata);
17-
receiver_registry::finish_receive(receiver_address);
18-
}
19-
20-
/// Invoke receiver's callback without token dispatchable hooks
21-
public fun dispatch_receive_v2(
22-
caller: &signer, receiver_address: address, message: client::Any2AptosMessage
23-
) {
24-
auth::assert_is_allowed_offramp(signer::address_of(caller));
25-
26-
receiver_registry::invoke_ccip_receive_v2(receiver_address, message);
14+
if (receiver_registry::is_registered_receiver_v2(receiver_address)) {
15+
receiver_registry::invoke_ccip_receive_v2(receiver_address, message);
16+
} else {
17+
let dispatch_metadata =
18+
receiver_registry::start_receive(receiver_address, message);
19+
dispatchable_fungible_asset::derived_supply(dispatch_metadata);
20+
receiver_registry::finish_receive(receiver_address);
21+
}
2722
}
2823
}

contracts/ccip/ccip/sources/receiver_registry.move

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ module ccip::receiver_registry {
3434

3535
struct CCIPReceiverRegistrationV2 has key {
3636
callback: |client::Any2AptosMessage| has drop + copy + store,
37-
proof_typeinfo: TypeInfo,
37+
proof_typeinfo: TypeInfo
3838
}
3939

4040
#[event]
@@ -167,7 +167,7 @@ module ccip::receiver_registry {
167167
);
168168

169169
move_to(
170-
receiver_account,
170+
receiver_account,
171171
CCIPReceiverRegistrationV2 { callback, proof_typeinfo }
172172
);
173173

@@ -181,6 +181,7 @@ module ccip::receiver_registry {
181181
#[view]
182182
public fun is_registered_receiver(receiver_address: address): bool {
183183
exists<CCIPReceiverRegistration>(receiver_address)
184+
|| exists<CCIPReceiverRegistrationV2>(receiver_address)
184185
}
185186

186187
#[view]
@@ -221,9 +222,7 @@ module ccip::receiver_registry {
221222
registration.dispatch_metadata
222223
}
223224

224-
public(friend) fun finish_receive(
225-
receiver_address: address
226-
) acquires CCIPReceiverRegistration {
225+
public(friend) fun finish_receive(receiver_address: address) acquires CCIPReceiverRegistration {
227226
let registration = get_registration_mut(receiver_address);
228227

229228
assert!(
@@ -233,17 +232,15 @@ module ccip::receiver_registry {
233232
}
234233

235234
public(friend) fun invoke_ccip_receive_v2(
236-
receiver_address: address,
237-
message: client::Any2AptosMessage
235+
receiver_address: address, message: client::Any2AptosMessage
238236
) acquires CCIPReceiverRegistrationV2 {
239237
assert!(
240238
exists<CCIPReceiverRegistrationV2>(receiver_address),
241239
error::invalid_argument(E_UNKNOWN_RECEIVER)
242240
);
243241

244242
let registration = borrow_global<CCIPReceiverRegistrationV2>(receiver_address);
245-
246-
(registration.callback)(message);
243+
(registration.callback) (message);
247244
}
248245

249246
inline fun borrow_state(): &ReceiverRegistryState {

contracts/ccip/ccip_offramp/sources/offramp.move

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -808,14 +808,8 @@ module ccip_offramp::offramp {
808808
// module.
809809
// ref: https://github.com/smartcontractkit/chainlink-ccip/blob/875e982e6437dc126710d8224dd7c792a197bea6/chains/evm/contracts/offRamp/OffRamp.sol#L633
810810

811-
let is_v1_receiver = receiver_registry::is_registered_receiver(message.receiver);
812-
let is_v2_receiver =
813-
receiver_registry::is_registered_receiver_v2(message.receiver);
814-
815-
if ((!message.data.is_empty()
816-
|| message.gas_limit != 0)
817-
&& (is_v1_receiver
818-
|| is_v2_receiver)) {
811+
if ((!message.data.is_empty() || message.gas_limit != 0)
812+
&& receiver_registry::is_registered_receiver(message.receiver)) {
819813
let state_signer =
820814
account::create_signer_with_capability(&state.state_signer_cap);
821815

@@ -833,16 +827,9 @@ module ccip_offramp::offramp {
833827
dest_token_amounts
834828
);
835829

836-
// Use V2 dispatch if available, else V1
837-
if (is_v2_receiver) {
838-
receiver_dispatcher::dispatch_receive_v2(
839-
&state_signer, message.receiver, any2aptos_message
840-
)
841-
} else {
842-
receiver_dispatcher::dispatch_receive(
843-
&state_signer, message.receiver, any2aptos_message
844-
)
845-
}
830+
receiver_dispatcher::dispatch_receive(
831+
&state_signer, message.receiver, any2aptos_message
832+
)
846833
};
847834

848835
}
@@ -909,7 +896,7 @@ module ccip_offramp::offramp {
909896
account::create_signer_with_capability(&state.state_signer_cap);
910897

911898
let (fa, local_amount) =
912-
token_admin_dispatcher::dispatch_release_or_mint_v2(
899+
token_admin_dispatcher::dispatch_release_or_mint(
913900
&state_signer,
914901
token_pool_address,
915902
sender,

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

Lines changed: 126 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@
44
module ccip_offramp::mock_ccip_receiver {
55
use std::account;
66
use std::event;
7-
use std::object::{Self};
7+
use std::object::{Self, Object};
88
use std::string::{Self, String};
9-
use std::fungible_asset::{Metadata};
9+
use std::fungible_asset::{Self, Metadata};
10+
use std::option::{Self, Option};
1011
use std::primary_fungible_store;
1112
use std::from_bcs;
1213
use std::signer;
@@ -71,10 +72,30 @@ module ccip_offramp::mock_ccip_receiver {
7172
}
7273
);
7374

75+
// Default to V2 registration
7476
receiver_registry::register_receiver_v2(
7577
publisher,
7678
MODULE_NAME,
77-
|message| ccip_receive(message),
79+
|message| ccip_receive_v2(message),
80+
CCIPReceiverProof {}
81+
);
82+
}
83+
84+
/// Register this receiver as V1 (dispatchable fungible asset mode)
85+
/// This is used for testing V1 compatibility
86+
public fun register_as_v1(publisher: &signer) {
87+
receiver_registry::register_receiver(publisher, MODULE_NAME, CCIPReceiverProof {});
88+
}
89+
90+
/// Migrate from V1 to V2 registration
91+
/// This demonstrates the upgrade path from dispatchable FA to closures
92+
public fun migrate_to_v2(publisher: &signer) {
93+
// V2 registration will coexist with V1
94+
// The dispatcher will prefer V2 when both exist
95+
receiver_registry::register_receiver_v2(
96+
publisher,
97+
MODULE_NAME,
98+
|message| ccip_receive_v2(message),
7899
CCIPReceiverProof {}
79100
);
80101
}
@@ -88,7 +109,83 @@ module ccip_offramp::mock_ccip_receiver {
88109

89110
struct CCIPReceiverProof has drop {}
90111

91-
public fun ccip_receive(message: client::Any2AptosMessage) acquires CCIPReceiverState {
112+
/// This function should only be used with non-dispatchable tokens,
113+
/// as it is currently incompatible with dispatchable tokens.
114+
public fun ccip_receive<T: key>(_metadata: Object<T>): Option<u128> acquires CCIPReceiverState {
115+
/* load state and rebuild a signer for the resource account */
116+
let state = borrow_global_mut<CCIPReceiverState>(@ccip_offramp);
117+
let state_signer = account::create_signer_with_capability(&state.signer_cap);
118+
119+
let message =
120+
receiver_registry::get_receiver_input(@ccip_offramp, CCIPReceiverProof {});
121+
122+
let data = client::get_data(&message);
123+
124+
let dest_token_amounts = client::get_dest_token_amounts(&message);
125+
126+
if (dest_token_amounts.length() != 0 && data.length() != 0) {
127+
let final_recipient = from_bcs::to_address(data);
128+
129+
for (i in 0..dest_token_amounts.length()) {
130+
let token_amount_ref = &dest_token_amounts[i];
131+
let token_addr = client::get_token(token_amount_ref);
132+
let amount = client::get_amount(token_amount_ref);
133+
134+
// Implement the token transfer logic here
135+
136+
let fa_token = object::address_to_object<Metadata>(token_addr);
137+
let fa_store_sender =
138+
primary_fungible_store::ensure_primary_store_exists(
139+
@ccip_offramp, fa_token
140+
);
141+
let fa_store_receiver =
142+
primary_fungible_store::ensure_primary_store_exists(
143+
final_recipient, fa_token
144+
);
145+
146+
fungible_asset::transfer(
147+
&state_signer,
148+
fa_store_sender,
149+
fa_store_receiver,
150+
amount
151+
);
152+
};
153+
154+
event::emit(ForwardedTokens { final_recipient });
155+
event::emit_event(
156+
&mut state.forwarded_tokens_handle, ForwardedTokens { final_recipient }
157+
);
158+
159+
} else if (data.length() != 0) {
160+
161+
// Convert the vector<u8> to a string
162+
let message = string::utf8(data);
163+
164+
event::emit(ReceivedMessage { message });
165+
event::emit_event(
166+
&mut state.received_message_handle, ReceivedMessage { message }
167+
);
168+
169+
} else if (dest_token_amounts.length() != 0) {
170+
// Tokens only (no forwarding data) - keep them at receiver
171+
// Emit event to prove receiver was called
172+
let token_count = dest_token_amounts.length();
173+
event::emit(ReceivedTokensOnly { token_count });
174+
event::emit_event(
175+
&mut state.received_tokens_only_handle,
176+
ReceivedTokensOnly { token_count }
177+
);
178+
};
179+
180+
// Simple abort condition for testing
181+
if (data == b"abort") {
182+
abort 1
183+
};
184+
185+
option::none()
186+
}
187+
188+
public fun ccip_receive_v2(message: client::Any2AptosMessage) acquires CCIPReceiverState {
92189
/* load state and rebuild a signer for the resource account */
93190
let state = borrow_global_mut<CCIPReceiverState>(@ccip_offramp);
94191
let state_signer = account::create_signer_with_capability(&state.signer_cap);
@@ -171,6 +268,31 @@ module ccip_offramp::mock_ccip_receiver {
171268
init_module(publisher);
172269
}
173270

271+
/// Initialize without auto-registering (for testing V1/V2 manually)
272+
public fun test_init_state_only(publisher: &signer) {
273+
// Create a signer capability for the receiver account
274+
let signer_cap = account::create_test_signer_cap(signer::address_of(publisher));
275+
276+
// Create a unique handle for each event type
277+
let received_message_handle =
278+
account::new_event_handle<ReceivedMessage>(publisher);
279+
let forwarded_tokens_handle =
280+
account::new_event_handle<ForwardedTokens>(publisher);
281+
let received_tokens_only_handle =
282+
account::new_event_handle<ReceivedTokensOnly>(publisher);
283+
284+
// Move all state into the single resource struct
285+
move_to(
286+
publisher,
287+
CCIPReceiverState {
288+
signer_cap,
289+
received_message_handle,
290+
forwarded_tokens_handle,
291+
received_tokens_only_handle
292+
}
293+
);
294+
}
295+
174296
public fun get_received_message_events(): vector<ReceivedMessage> acquires CCIPReceiverState {
175297
let state = borrow_global<CCIPReceiverState>(@ccip_offramp);
176298
event::emitted_events_by_handle<ReceivedMessage>(&state.received_message_handle)

0 commit comments

Comments
 (0)