MCMS support LockReleaseTokenPool RebalancerCap

Author: JohnChangUK

bindings/generated/ccip/ccip_token_pools/lock_release_token_pool/lock_release_token_pool.go

@@ -15,6 +15,7 @@ use std::type_name::{Self, TypeName};
 use sui::address;
 use sui::clock::Clock;
 use sui::coin::{Self, Coin, CoinMetadata, TreasuryCap};
+use sui::event;
 use sui::package::UpgradeCap;
 
 #[allow(lint(coin_field))]
@@ -30,13 +31,23 @@ public struct RebalancerCap has key, store {
     id: UID,
 }
 
+public struct RebalancerSet has copy, drop {
+    old_rebalancer_cap_id: ID,
+    new_rebalancer_cap_id: ID,
+}
+
 const CLOCK_ADDRESS: address = @0x6;
 
 const EInvalidArguments: u64 = 1;
 const ETokenPoolBalanceTooLow: u64 = 2;
 const EInvalidOwnerCap: u64 = 3;
 const EInvalidFunction: u64 = 4;
 const EInvalidRebalancerCap: u64 = 5;
+const EMcmsRebalancerAlreadySet: u64 = 6;
+const ERebalancerCapIsInUse: u64 = 7;
+const ERebalancerCapNotTransferredOut: u64 = 8;
+const EInvalidRebalancer: u64 = 9;
+const ERebalancerCapDoesNotExist: u64 = 10;
 
 // ================================================================
 // |                             Init                             |
@@ -515,6 +526,145 @@ public fun get_rebalancer<T>(state: &LockReleaseTokenPoolState<T>): address {
     object::id_to_address(&state.rebalancer_cap_id)
 }
 
+public struct McmsCap has key, store {
+    id: UID,
+    owner_cap: OwnerCap,
+    rebalancer_cap: Option<RebalancerCap>,
+}
+
+fun new_mcms_cap(
+    owner_cap: OwnerCap,
+    rebalancer_cap: Option<RebalancerCap>,
+    ctx: &mut TxContext,
+): McmsCap {
+    McmsCap {
+        id: object::new(ctx),
+        owner_cap,
+        rebalancer_cap,
+    }
+}
+
+public fun mcms_set_rebalancer<T>(
+    state: &mut LockReleaseTokenPoolState<T>,
+    registry: &mut Registry,
+    params: ExecutingCallbackParams,
+    ctx: &mut TxContext,
+) {
+    let (mcms_cap, function, data) = mcms_registry::get_callback_params_with_caps<
+        McmsCallback<T>,
+        McmsCap,
+    >(
+        registry,
+        McmsCallback<T> {},
+        params,
+    );
+    assert!(function == string::utf8(b"set_rebalancer"), EInvalidFunction);
+
+    let mut stream = bcs_stream::new(data);
+    bcs_stream::validate_obj_addrs(
+        vector[object::id_address(state), object::id_address(&mcms_cap.owner_cap)],
+        &mut stream,
+    );
+    let rebalancer = bcs_stream::deserialize_address(&mut stream);
+    bcs_stream::assert_is_consumed(&stream);
+
+    if (rebalancer == mcms_registry::get_multisig_address()) {
+        assert!(mcms_cap.rebalancer_cap.is_none(), EMcmsRebalancerAlreadySet);
+
+        let new_rebalancer_cap = RebalancerCap {
+            id: object::new(ctx),
+        };
+        let new_rebalancer_cap_id = object::id(&new_rebalancer_cap);
+        mcms_cap.rebalancer_cap.fill(new_rebalancer_cap);
+
+        set_rebalancer_cap_id(state, &mcms_cap.owner_cap, new_rebalancer_cap_id, ctx);
+    } else {
+        set_rebalancer(state, &mcms_cap.owner_cap, rebalancer, ctx);
+        // If MCMS previously owned the rebalancer cap, we destroy it now that MCMS no longer owns it
+        if (mcms_cap.rebalancer_cap.is_some()) {
+            destroy_rebalancer_cap(state, mcms_cap.rebalancer_cap.extract(), ctx);
+        }
+    }
+}
+
+public fun set_rebalancer<T>(
+    state: &mut LockReleaseTokenPoolState<T>,
+    owner_cap: &OwnerCap,
+    rebalancer: address,
+    ctx: &mut TxContext,
+) {
+    assert!(rebalancer != mcms_registry::get_multisig_address(), EInvalidRebalancer);
+
+    let rebalancer_cap = RebalancerCap {
+        id: object::new(ctx),
+    };
+
+    // Update `LockReleaseTokenPoolState` before sending to rebalancer address
+    set_rebalancer_cap_id(
+        state,
+        owner_cap,
+        object::id(&rebalancer_cap),
+        ctx,
+    );
+
+    transfer::public_transfer(rebalancer_cap, rebalancer);
+}
+
+/// Only owner can set rebalancer cap id
+fun set_rebalancer_cap_id<T>(
+    state: &mut LockReleaseTokenPoolState<T>,
+    owner_cap: &OwnerCap,
+    new_rebalancer_cap_id: ID,
+    _ctx: &mut TxContext,
+) {
+    assert!(object::id(owner_cap) == ownable::owner_cap_id(&state.ownable_state), EInvalidOwnerCap);
+
+    let old_rebalancer_cap_id = state.rebalancer_cap_id;
+    state.rebalancer_cap_id = new_rebalancer_cap_id;
+
+    event::emit(RebalancerSet { old_rebalancer_cap_id, new_rebalancer_cap_id });
+}
+
+public fun mcms_destroy_mcms_cap<T>(
+    state: &mut LockReleaseTokenPoolState<T>,
+    registry: &mut Registry,
+    params: ExecutingCallbackParams,
+    ctx: &mut TxContext,
+) {
+    let (mcms_cap, function, data) = mcms_registry::get_callback_params_with_caps<
+        McmsCallback<T>,
+        McmsCap,
+    >(
+        registry,
+        McmsCallback<T> {},
+        params,
+    );
+    assert!(function == string::utf8(b"destroy_mcms_cap"), EInvalidFunction);
+    assert!(mcms_cap.rebalancer_cap.is_some(), ERebalancerCapDoesNotExist);
+
+    let mut stream = bcs_stream::new(data);
+    bcs_stream::validate_obj_addrs(
+        vector[object::id_address(state), object::id_address(mcms_cap.rebalancer_cap.borrow())],
+        &mut stream,
+    );
+    bcs_stream::assert_is_consumed(&stream);
+
+    let rebalancer_cap = mcms_cap.rebalancer_cap.extract();
+    destroy_rebalancer_cap(state, rebalancer_cap, ctx);
+}
+
+/// Clean up old rebalancer caps not in use, anyone can call this function to clean up old rebalancer caps not in use.
+public fun destroy_rebalancer_cap<T>(
+    state: &mut LockReleaseTokenPoolState<T>,
+    rebalancer_cap: RebalancerCap,
+    _ctx: &mut TxContext,
+) {
+    assert!(state.rebalancer_cap_id != object::id(&rebalancer_cap), ERebalancerCapIsInUse);
+
+    let RebalancerCap { id } = rebalancer_cap;
+    object::delete(id);
+}
+
 #[test_only]
 public fun create_fake_rebalancer_cap(ctx: &mut TxContext): RebalancerCap {
     RebalancerCap {
@@ -597,17 +747,20 @@ public fun execute_ownership_transfer<T>(
     ownable::execute_ownership_transfer(owner_cap, &mut state.ownable_state, to, ctx);
 }
 
+/// Setting rebalancer cap should be called via `mcms_set_rebalancer`
 public fun execute_ownership_transfer_to_mcms<T>(
     owner_cap: OwnerCap,
     state: &mut LockReleaseTokenPoolState<T>,
     registry: &mut Registry,
     to: address,
     ctx: &mut TxContext,
 ) {
-    ownable::execute_ownership_transfer_to_mcms(
-        owner_cap,
+    assert!(object::id(&owner_cap) == state.ownable_state.owner_cap_id(), EInvalidOwnerCap);
+
+    ownable::execute_ownership_and_cap_transfer_to_mcms(
         &mut state.ownable_state,
         registry,
+        new_mcms_cap(owner_cap, option::none(), ctx),
         to,
         McmsCallback<T> {},
         vector[b"lock_release_token_pool"],
@@ -946,9 +1099,9 @@ public fun mcms_transfer_ownership<T>(
     params: ExecutingCallbackParams,
     ctx: &mut TxContext,
 ) {
-    let (owner_cap, function, data) = mcms_registry::get_callback_params_with_caps<
+    let (mcms_cap, function, data) = mcms_registry::get_callback_params_with_caps<
         McmsCallback<T>,
-        OwnerCap,
+        McmsCap,
     >(
         registry,
         McmsCallback<T> {},
@@ -958,14 +1111,14 @@ public fun mcms_transfer_ownership<T>(
 
     let mut stream = bcs_stream::new(data);
     bcs_stream::validate_obj_addrs(
-        vector[object::id_address(state), object::id_address(owner_cap)],
+        vector[object::id_address(state), object::id_address(&mcms_cap.owner_cap)],
         &mut stream,
     );
 
     let to = bcs_stream::deserialize_address(&mut stream);
     bcs_stream::assert_is_consumed(&stream);
 
-    transfer_ownership(state, owner_cap, to, ctx);
+    transfer_ownership(state, &mcms_cap.owner_cap, to, ctx);
 }
 
 public fun mcms_execute_ownership_transfer<T>(
@@ -974,9 +1127,9 @@ public fun mcms_execute_ownership_transfer<T>(
     params: ExecutingCallbackParams,
     ctx: &mut TxContext,
 ) {
-    let (_owner_cap, function, data) = mcms_registry::get_callback_params_with_caps<
+    let (mcms_cap, function, data) = mcms_registry::get_callback_params_with_caps<
         McmsCallback<T>,
-        OwnerCap,
+        McmsCap,
     >(
         registry,
         McmsCallback<T> {},
@@ -986,14 +1139,22 @@ public fun mcms_execute_ownership_transfer<T>(
 
     let mut stream = bcs_stream::new(data);
     bcs_stream::validate_obj_addrs(
-        vector[object::id_address(_owner_cap), object::id_address(state)],
+        vector[object::id_address(&mcms_cap.owner_cap), object::id_address(state)],
         &mut stream,
     );
 
     let to = bcs_stream::deserialize_address(&mut stream);
     bcs_stream::assert_is_consumed(&stream);
 
-    let owner_cap = mcms_registry::release_cap(registry, McmsCallback<T> {});
+    let McmsCap { id, owner_cap, rebalancer_cap } = mcms_registry::release_cap(
+        registry,
+        McmsCallback<T> {},
+    );
+    assert!(rebalancer_cap.is_none(), ERebalancerCapNotTransferredOut);
+
+    rebalancer_cap.destroy_none();
+    object::delete(id);
+
     execute_ownership_transfer(owner_cap, state, to, ctx);
 }
 
@@ -1078,3 +1239,8 @@ public fun destroy_token_pool<T>(
 
     reserve
 }
+
+#[test_only]
+public fun mcms_rebalancer_cap_address(mcms_cap: &McmsCap): address {
+    object::id_address(mcms_cap.rebalancer_cap.borrow())
+}

bindings/generated/mcms/mcms_registry/mcms_registry.go

@@ -223,6 +223,35 @@ public fun execute_ownership_transfer_to_mcms<T: drop>(
     event::emit(OwnershipTransferred { from: current_owner, to: new_owner });
 }
 
+/// Used for registering a custom Cap type with MCMS, e.g. McmsCap for LockReleaseTokenPool.
+/// Needed as `OwnableState` does not know about the custom Cap type and needs updating.
+/// Calling function should assert the owner cap inside `cap` matches `ownable_state.owner_cap_id`.
+public(package) fun execute_ownership_and_cap_transfer_to_mcms<C: key + store, T: drop>(
+    state: &mut OwnableState,
+    registry: &mut Registry,
+    cap: C,
+    to: address,
+    proof: T,
+    allowed_modules: vector<vector<u8>>,
+    ctx: &mut TxContext,
+) {
+    assert!(state.pending_transfer.is_some(), ENoPendingTransfer);
+
+    let pending_transfer = state.pending_transfer.extract();
+    let current_owner = state.owner;
+
+    assert!(pending_transfer.from == current_owner, EOwnerChanged);
+    assert!(pending_transfer.to == to, EProposedOwnerMismatch);
+    assert!(pending_transfer.accepted, ETransferNotAccepted);
+    assert!(to == mcms_registry::get_multisig_address(), EMustTransferToMcms);
+
+    state.owner = to;
+    state.pending_transfer = option::none();
+
+    mcms_registry::register_entrypoint(registry, proof, cap, allowed_modules, ctx);
+    event::emit(OwnershipTransferred { from: current_owner, to });
+}
+
 public fun destroy(state: OwnableState, owner_cap: OwnerCap, _ctx: &mut TxContext) {
     let OwnableState {
         id: state_id,