[LST-348] Reimpl SetRebalanceAuthority

controller/core/src/instructions/rebalance/mod.rs

@@ -1,2 +1,3 @@
 pub mod end;
+pub mod set_rebal_auth;
 pub mod start;

controller/core/src/instructions/rebalance/set_rebal_auth.rs

@@ -0,0 +1,47 @@
+use generic_array_struct::generic_array_struct;
+
+use crate::instructions::generic::DiscmOnlyIxData;
+
+// Accounts
+
+#[generic_array_struct(builder pub)]
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
+#[repr(transparent)]
+pub struct SetRebalAuthIxAccs<T> {
+    /// The signer setting the new rebalance auth.
+    /// Can either be pool admin or current rebalance auth.
+    pub signer: T,
+
+    /// New rebalance auth to set to
+    pub new: T,
+
+    /// The pool's state singleton PDA
+    pub pool_state: T,
+}
+
+impl<T: Copy> SetRebalAuthIxAccs<T> {
+    #[inline]
+    pub const fn memset(val: T) -> Self {
+        Self([val; SET_REBAL_AUTH_IX_ACCS_LEN])
+    }
+}
+
+pub type SetRebalAuthIxKeys<'a> = SetRebalAuthIxAccs<&'a [u8; 32]>;
+
+pub type SetRebalAuthIxKeysOwned = SetRebalAuthIxAccs<[u8; 32]>;
+
+pub type SetRebalAuthIxAccFlags = SetRebalAuthIxAccs<bool>;
+
+pub const SET_REBAL_AUTH_IX_IS_WRITER: SetRebalAuthIxAccFlags =
+    SetRebalAuthIxAccFlags::memset(false).const_with_pool_state(true);
+
+pub const SET_REBAL_AUTH_IX_IS_SIGNER: SetRebalAuthIxAccFlags =
+    SetRebalAuthIxAccFlags::memset(false).const_with_signer(true);
+
+// Data
+
+pub const SET_REBAL_AUTH_IX_DISCM: u8 = 21;
+
+pub type SetRebalAuthIxData = DiscmOnlyIxData<SET_REBAL_AUTH_IX_DISCM>;
+
+pub const SET_REBAL_AUTH_IX_DATA_LEN: usize = SetRebalAuthIxData::DATA_LEN;

controller/program/src/instructions/rebalance/mod.rs

@@ -1,2 +1,3 @@
 pub mod end;
+pub mod set_rebal_auth;
 pub mod start;

controller/program/src/instructions/rebalance/set_rebal_auth.rs

@@ -0,0 +1,62 @@
+use inf1_ctl_jiminy::{
+    account_utils::{pool_state_checked, pool_state_checked_mut},
+    err::Inf1CtlErr,
+    instructions::rebalance::set_rebal_auth::{
+        NewSetRebalAuthIxAccsBuilder, SetRebalAuthIxAccs, SET_REBAL_AUTH_IX_IS_SIGNER,
+    },
+    keys::POOL_STATE_ID,
+    program_err::Inf1CtlCustomProgErr,
+};
+use jiminy_cpi::{
+    account::{Abr, AccountHandle},
+    program_error::{ProgramError, NOT_ENOUGH_ACCOUNT_KEYS},
+};
+
+use crate::verify::{verify_not_rebalancing_and_not_disabled, verify_pks, verify_signers};
+
+type SetRebalAuthIxAccounts<'acc> = SetRebalAuthIxAccs<AccountHandle<'acc>>;
+
+#[inline]
+pub fn set_rebal_auth_accs_checked<'acc>(
+    abr: &Abr,
+    accs: &[AccountHandle<'acc>],
+) -> Result<SetRebalAuthIxAccounts<'acc>, ProgramError> {
+    let accs = accs.first_chunk().ok_or(NOT_ENOUGH_ACCOUNT_KEYS)?;
+    let accs = SetRebalAuthIxAccs(*accs);
+
+    let expected_pks = NewSetRebalAuthIxAccsBuilder::start()
+        .with_pool_state(&POOL_STATE_ID)
+        // Free: check either rebal auth or pool admin below
+        .with_signer(abr.get(*accs.signer()).key())
+        // Free: signer is free to set new auth to whatever pk as pleased
+        .with_new(abr.get(*accs.new()).key())
+        .build();
+    verify_pks(abr, &accs.0, &expected_pks.0)?;
+
+    verify_signers(abr, &accs.0, &SET_REBAL_AUTH_IX_IS_SIGNER.0)?;
+
+    let pool = pool_state_checked(abr.get(*accs.pool_state()))?;
+
+    verify_not_rebalancing_and_not_disabled(pool)?;
+
+    let signer_pk = abr.get(*accs.signer()).key();
+
+    if *signer_pk != pool.rebalance_authority && *signer_pk != pool.admin {
+        return Err(
+            Inf1CtlCustomProgErr(Inf1CtlErr::UnauthorizedSetRebalanceAuthoritySigner).into(),
+        );
+    }
+
+    Ok(accs)
+}
+
+#[inline]
+pub fn process_set_rebal_auth(
+    abr: &mut Abr,
+    accs: SetRebalAuthIxAccounts,
+) -> Result<(), ProgramError> {
+    let new_rebal_auth = *abr.get(*accs.new()).key();
+    let pool = pool_state_checked_mut(abr.get_mut(*accs.pool_state()))?;
+    pool.rebalance_authority = new_rebal_auth;
+    Ok(())
+}

controller/program/src/lib.rs

@@ -26,6 +26,7 @@ use inf1_ctl_jiminy::instructions::{
     },
     rebalance::{
         end::END_REBALANCE_IX_DISCM,
+        set_rebal_auth::SET_REBAL_AUTH_IX_DISCM,
         start::{StartRebalanceIxData, START_REBALANCE_IX_DISCM},
     },
     swap::{exact_in::SWAP_EXACT_IN_IX_DISCM, exact_out::SWAP_EXACT_OUT_IX_DISCM, IxData},
@@ -70,7 +71,11 @@ use crate::instructions::{
         },
         withdraw_protocol_fee::{process_withdraw_protocol_fees, withdraw_protocol_fees_checked},
     },
-    rebalance::{end::process_end_rebalance, start::process_start_rebalance},
+    rebalance::{
+        end::process_end_rebalance,
+        set_rebal_auth::{process_set_rebal_auth, set_rebal_auth_accs_checked},
+        start::process_start_rebalance,
+    },
     swap::{process_swap_exact_in, process_swap_exact_out},
     sync_sol_value::process_sync_sol_value,
 };
@@ -249,6 +254,11 @@ fn process_ix(
             sol_log("EndRebalance");
             process_end_rebalance(abr, accounts, cpi)
         }
+        (&SET_REBAL_AUTH_IX_DISCM, _) => {
+            sol_log("SetRebalAuth");
+            let accs = set_rebal_auth_accs_checked(abr, accounts)?;
+            process_set_rebal_auth(abr, accs)
+        }
         _ => Err(INVALID_INSTRUCTION_DATA.into()),
     }
 }

controller/program/tests/tests/rebalance/mod.rs

@@ -1,2 +1,3 @@
 mod chain;
+mod set_rebal_auth;
 mod test_utils;