Function closures for CCIP token pool hooks

[JohnChangUK]

Function Closures Migration for CCIP Token Pools and Receivers

Migrates CCIP token pool and receiver infrastructure to use function closures, removing the need for dispatchable_fungible_asset pattern on token pools and receivers.

V1 Reentrancy Problem

When Receiver contracts try to transfer tokens within PTTs, if they call dispatchable_fungible_asset::withdraw/deposit(), this causes reentrancy as receivers in V1 are invoked via dynamic dispatch.

With closures, there is no need to invoke hooks via dynamic dispatch.

V2 Receiver

// V2 receiver - full freedom to transfer tokens
public fun ccip_receive_v2(message: Any2AptosMessage) {
    let tokens = client::get_dest_token_amounts(&message);

    for (token in tokens) {
        // Can call ` primary_fungible_store/dispatchable_fungible_asset` module
        primary_fungible_store::transfer(recipient, token, amount);
    }
}

---

How to Register with V2

Token Pool Registration

Complete example from burn_mint_token_pool:

module burn_mint_token_pool::burn_mint_token_pool {
    // V2 callback implementations
    public fun lock_or_burn_v2(
        fa: FungibleAsset,
        input: token_admin_registry::LockOrBurnInputV1
    ): (vector<u8>, vector<u8>) {
        let pool = borrow_pool_mut();
        let amount = fungible_asset::amount(&fa);

        // Validate
        let dest_token_address = token_pool::validate_lock_or_burn(
            &mut pool.token_pool_state, &fa, &input, amount
        );

        // Burn token
        fungible_asset::burn(pool.burn_ref.borrow(), fa);

        // Return destination info
        (dest_token_address, token_pool::encode_local_decimals(&pool.token_pool_state))
    }

    public fun release_or_mint_v2(
        input: token_admin_registry::ReleaseOrMintInputV1
    ): (FungibleAsset, u64) {
        let pool = borrow_pool_mut();
        let amount = token_admin_registry::get_release_or_mint_amount(&input);

        // Validate
        token_pool::validate_release_or_mint(&mut pool.token_pool_state, &input, amount);

        // Mint token
        let fa = fungible_asset::mint(pool.mint_ref.borrow(), amount);

        // Return token and amount
        (fa, amount)
    }

    // Registration in init_module
    fun init_module(publisher: &signer) {
        // ... initialization code ...

        // V1 registration (backward compatibility)
        token_admin_registry::register_pool(
            publisher,
            b"burn_mint_token_pool",
            @burn_mint_local_token,
            CallbackProof {}
        );

        // V2 registration with closures
        let lock_or_burn_closure = |fa, input| lock_or_burn_v2(fa, input);
        let release_or_mint_closure = |input| release_or_mint_v2(input);

        token_admin_registry::register_pool_v2(
            publisher,
            b"burn_mint_token_pool",
            @burn_mint_local_token,
            lock_or_burn_closure,
            release_or_mint_closure,
            CallbackProof {}
        );
    }
}

Receiver Registration

module mock_ccip_receiver::mock_ccip_receiver {
    // V2 callback implementation
    public fun ccip_receive_v2(message: client::Any2AptosMessage) {
        let tokens = client::get_dest_token_amounts(&message);
        let data = client::get_data(&message);

        // Handle token transfers
        let i = 0;
        while (i < tokens.length()) {
            let token_transfer = tokens.borrow(i);
            let token_obj = client::get_dest_token(token_transfer);
            let amount = client::get_dest_token_amount(token_transfer);

            // Forward tokens to final recipient
            if (!data.is_empty()) {
                let recipient = decode_address_from_data(&data);
                primary_fungible_store::transfer(
                    &receiver_signer,
                    token_obj,
                    recipient,
                    amount
                );
            };

            i = i + 1;
        };

        // Emit event
        event::emit(ReceivedTokensOnly { token_count: tokens.length() });
    }

    // Registration in init_module
    fun init_module(publisher: &signer) {
        // ... initialization code ...

        // V1 registration (backward compatibility)
        receiver_registry::register_receiver(
            publisher,
            b"mock_ccip_receiver",
            CallbackProof {}
        );

        // V2 registration with closure
        let callback = |message| ccip_receive_v2(message);

        receiver_registry::register_receiver_v2(
            publisher,
            b"mock_ccip_receiver",
            callback,
            CallbackProof {}
        );
    }
}

Key Points

Closure Signatures:

• Token pools need TWO closures: lock_or_burn and release_or_mint
• Receivers need ONE closure: ccip_receive
• Closures must have has drop + copy + store abilities

Callback Functions:

• Receive parameters directly (no global state retrieval)
• Return values directly as tuples
• Can freely transfer tokens (no module lock)

---

Backward Compatibility

Dual Registration

All pools support V1 and V2 simultaneously:

fun init_module(publisher: &signer) {
    // V1 registration (backward compatibility)
    token_admin_registry::register_pool(publisher, ...);

    // V2 registration (new functionality)
    let lock_or_burn = |fa, input| lock_or_burn_v2(fa, input);
    let release_or_mint = |input| release_or_mint_v2(input);
    token_admin_registry::register_pool_v2(publisher, lock_or_burn, release_or_mint, ...);
}

Automatic Dispatcher Routing

Dispatcher automatically prefers V2 when available:

if (has_token_pool_config(pool_address)) {
    lock_or_burn_v2(...)  // V2 path - fast, no lock
} else {
    // V1 fallback - slow, module lock
}

Migration Path

1. Deploy upgrade module: upgrade_v2::init_module() adds V2 to existing pool
2. Pool now supports both V1 and V2
3. Dispatcher automatically uses V2
4. Zero downtime - each pool upgraded independently

---

Technical Requirements

REQUIRED: Move 2.2

aptos move compile --dev --language-version 2.2
aptos move test --dev --language-version 2.2

---

[JohnChangUK]

Tests for both V1 and V2 execution in the SAME test, make sure both routes are compatible.