diff --git a/controller/core/src/instructions/rebalance/mod.rs b/controller/core/src/instructions/rebalance/mod.rs index cf88740d..3e57405c 100644 --- a/controller/core/src/instructions/rebalance/mod.rs +++ b/controller/core/src/instructions/rebalance/mod.rs @@ -1,2 +1,3 @@ pub mod end; +pub mod set_rebal_auth; pub mod start; diff --git a/controller/core/src/instructions/rebalance/set_rebal_auth.rs b/controller/core/src/instructions/rebalance/set_rebal_auth.rs new file mode 100644 index 00000000..92b5b71d --- /dev/null +++ b/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 { + /// 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 SetRebalAuthIxAccs { + #[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; + +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; + +pub const SET_REBAL_AUTH_IX_DATA_LEN: usize = SetRebalAuthIxData::DATA_LEN; diff --git a/controller/program/src/instructions/rebalance/mod.rs b/controller/program/src/instructions/rebalance/mod.rs index cf88740d..3e57405c 100644 --- a/controller/program/src/instructions/rebalance/mod.rs +++ b/controller/program/src/instructions/rebalance/mod.rs @@ -1,2 +1,3 @@ pub mod end; +pub mod set_rebal_auth; pub mod start; diff --git a/controller/program/src/instructions/rebalance/set_rebal_auth.rs b/controller/program/src/instructions/rebalance/set_rebal_auth.rs new file mode 100644 index 00000000..acd2b679 --- /dev/null +++ b/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>; + +#[inline] +pub fn set_rebal_auth_accs_checked<'acc>( + abr: &Abr, + accs: &[AccountHandle<'acc>], +) -> Result, 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(()) +} diff --git a/controller/program/src/lib.rs b/controller/program/src/lib.rs index 87cca57a..91530f8c 100644 --- a/controller/program/src/lib.rs +++ b/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()), } } diff --git a/controller/program/tests/tests/rebalance/mod.rs b/controller/program/tests/tests/rebalance/mod.rs index 189f791c..1b63cc3c 100644 --- a/controller/program/tests/tests/rebalance/mod.rs +++ b/controller/program/tests/tests/rebalance/mod.rs @@ -1,2 +1,3 @@ mod chain; +mod set_rebal_auth; mod test_utils; diff --git a/controller/program/tests/tests/rebalance/set_rebal_auth.rs b/controller/program/tests/tests/rebalance/set_rebal_auth.rs new file mode 100644 index 00000000..554ea8c6 --- /dev/null +++ b/controller/program/tests/tests/rebalance/set_rebal_auth.rs @@ -0,0 +1,313 @@ +use inf1_ctl_jiminy::{ + accounts::pool_state::{PoolState, PoolStatePacked}, + err::Inf1CtlErr, + instructions::rebalance::set_rebal_auth::{ + NewSetRebalAuthIxAccsBuilder, SetRebalAuthIxData, SetRebalAuthIxKeysOwned, + SET_REBAL_AUTH_IX_ACCS_IDX_NEW, SET_REBAL_AUTH_IX_ACCS_IDX_SIGNER, + SET_REBAL_AUTH_IX_IS_SIGNER, SET_REBAL_AUTH_IX_IS_WRITER, + }, + keys::POOL_STATE_ID, + program_err::Inf1CtlCustomProgErr, + ID, +}; +use inf1_test_utils::{ + acc_bef_aft, any_normal_pk, any_pool_state, assert_diffs_pool_state, assert_jiminy_prog_err, + dedup_accounts, gen_pool_state, keys_signer_writable_to_metas, mock_sys_acc, + pool_state_account, silence_mollusk_logs, AnyPoolStateArgs, Diff, DiffsPoolStateArgs, + GenPoolStateArgs, PkAccountTup, PoolStateBools, PoolStatePks, +}; +use jiminy_cpi::program_error::{ProgramError, MISSING_REQUIRED_SIGNATURE}; +use mollusk_svm::result::{InstructionResult, ProgramResult}; +use proptest::{prelude::*, strategy::Union}; +use solana_instruction::Instruction; +use solana_pubkey::Pubkey; + +use crate::common::SVM; + +fn set_rebal_auth_ix(keys: SetRebalAuthIxKeysOwned) -> Instruction { + let accounts = keys_signer_writable_to_metas( + keys.0.iter(), + SET_REBAL_AUTH_IX_IS_SIGNER.0.iter(), + SET_REBAL_AUTH_IX_IS_WRITER.0.iter(), + ); + Instruction { + program_id: Pubkey::new_from_array(ID), + accounts, + data: SetRebalAuthIxData::as_buf().into(), + } +} + +fn set_rebal_auth_test_accs(keys: SetRebalAuthIxKeysOwned, pool: PoolState) -> Vec { + // dont care abt lamports, shouldnt affect anything + const LAMPORTS: u64 = 1_000_000_000; + let accs = NewSetRebalAuthIxAccsBuilder::start() + .with_signer(mock_sys_acc(LAMPORTS)) + .with_new(mock_sys_acc(LAMPORTS)) + .with_pool_state(pool_state_account(pool)) + .build(); + let mut res = keys.0.into_iter().map(Into::into).zip(accs.0).collect(); + dedup_accounts(&mut res); + res +} + +/// Returns `pool_state.rebalance_auth` at the end of ix +fn set_rebal_auth_test( + ix: &Instruction, + bef: &[PkAccountTup], + expected_err: Option>, +) -> [u8; 32] { + let InstructionResult { + program_result, + resulting_accounts: aft, + .. + } = SVM.with(|svm| svm.process_instruction(ix, bef)); + + let [pool_state_bef, pool_state_aft] = acc_bef_aft(&POOL_STATE_ID.into(), bef, &aft).map(|a| { + PoolStatePacked::of_acc_data(&a.data) + .unwrap() + .into_pool_state() + }); + + let old_rebal_auth = pool_state_bef.rebalance_authority; + let expected_new_rebal_auth = ix.accounts[SET_REBAL_AUTH_IX_ACCS_IDX_NEW].pubkey; + + match expected_err { + None => { + assert_eq!(program_result, ProgramResult::Success); + assert_diffs_pool_state( + &DiffsPoolStateArgs { + pks: PoolStatePks::default().with_rebalance_authority(Diff::Changed( + old_rebal_auth, + expected_new_rebal_auth.to_bytes(), + )), + ..Default::default() + }, + &pool_state_bef, + &pool_state_aft, + ); + } + Some(e) => { + assert_jiminy_prog_err(&program_result, e); + } + } + + pool_state_aft.rebalance_authority +} + +#[test] +fn admin_set_rebal_auth_test_correct_basic() { + let [admin, new_rebal_auth] = core::array::from_fn(|i| [u8::try_from(i).unwrap(); 32]); + let pool = gen_pool_state(GenPoolStateArgs { + pks: PoolStatePks::default().with_admin(admin), + ..Default::default() + }); + let keys = NewSetRebalAuthIxAccsBuilder::start() + .with_new(new_rebal_auth) + .with_signer(admin) + .with_pool_state(POOL_STATE_ID) + .build(); + let ret = set_rebal_auth_test( + &set_rebal_auth_ix(keys), + &set_rebal_auth_test_accs(keys, pool), + Option::::None, + ); + assert_eq!(ret, new_rebal_auth); +} + +/// generates (new_rebal_auth, pool_state) +fn correct_strat_params() -> impl Strategy { + ( + any_normal_pk(), + any_pool_state(AnyPoolStateArgs { + bools: PoolStateBools::normal(), + ..Default::default() + }), + ) +} + +fn admin_correct_strat() -> impl Strategy)> { + correct_strat_params() + .prop_map(|(new_rebal_auth, ps)| { + ( + NewSetRebalAuthIxAccsBuilder::start() + .with_new(new_rebal_auth) + .with_signer(ps.admin) + .with_pool_state(POOL_STATE_ID) + .build(), + ps, + ) + }) + .prop_map(|(k, ps)| (set_rebal_auth_ix(k), set_rebal_auth_test_accs(k, ps))) +} + +fn rebal_auth_correct_strat() -> impl Strategy)> { + correct_strat_params() + .prop_map(|(new_rebal_auth, ps)| { + ( + NewSetRebalAuthIxAccsBuilder::start() + .with_new(new_rebal_auth) + .with_signer(ps.rebalance_authority) + .with_pool_state(POOL_STATE_ID) + .build(), + ps, + ) + }) + .prop_map(|(k, ps)| (set_rebal_auth_ix(k), set_rebal_auth_test_accs(k, ps))) +} + +fn unauthorized_strat() -> impl Strategy)> { + correct_strat_params() + .prop_flat_map(|(new_rebal_auth, ps)| { + ( + any::<[u8; 32]>().prop_filter("", move |pk| { + *pk != ps.admin && *pk != ps.rebalance_authority + }), + Just(new_rebal_auth), + Just(ps), + ) + }) + .prop_map(|(unauthorized_signer, new_admin, ps)| { + ( + NewSetRebalAuthIxAccsBuilder::start() + .with_new(new_admin) + .with_signer(unauthorized_signer) + .with_pool_state(POOL_STATE_ID) + .build(), + ps, + ) + }) + .prop_map(|(k, ps)| (set_rebal_auth_ix(k), set_rebal_auth_test_accs(k, ps))) +} + +fn missing_sig_strat() -> impl Strategy)> { + Union::new([ + admin_correct_strat().boxed(), + rebal_auth_correct_strat().boxed(), + ]) + .prop_map(|(mut ix, accs)| { + ix.accounts[SET_REBAL_AUTH_IX_ACCS_IDX_SIGNER].is_signer = false; + (ix, accs) + }) +} + +fn disabled_strat() -> impl Strategy)> { + ( + any_normal_pk(), + any_pool_state(AnyPoolStateArgs { + bools: PoolStateBools::normal().with_is_disabled(Some(Just(true).boxed())), + ..Default::default() + }), + ) + .prop_flat_map(|(new_rebal_auth, ps)| { + ( + Union::new([ps.rebalance_authority, ps.admin].map(|signer| { + Just( + NewSetRebalAuthIxAccsBuilder::start() + .with_new(new_rebal_auth) + .with_signer(signer) + .with_pool_state(POOL_STATE_ID) + .build(), + ) + })), + Just(ps), + ) + }) + .prop_map(|(k, ps)| (set_rebal_auth_ix(k), set_rebal_auth_test_accs(k, ps))) +} + +fn rebalancing_strat() -> impl Strategy)> { + ( + any_normal_pk(), + any_pool_state(AnyPoolStateArgs { + bools: PoolStateBools::normal().with_is_rebalancing(Some(Just(true).boxed())), + ..Default::default() + }), + ) + .prop_flat_map(|(new_rebal_auth, ps)| { + ( + Union::new([ps.rebalance_authority, ps.admin].map(|signer| { + Just( + NewSetRebalAuthIxAccsBuilder::start() + .with_new(new_rebal_auth) + .with_signer(signer) + .with_pool_state(POOL_STATE_ID) + .build(), + ) + })), + Just(ps), + ) + }) + .prop_map(|(k, ps)| (set_rebal_auth_ix(k), set_rebal_auth_test_accs(k, ps))) +} + +proptest! { + #[test] + fn admin_set_rebal_auth_correct_pt( + (ix, bef) in admin_correct_strat(), + ) { + silence_mollusk_logs(); + set_rebal_auth_test(&ix, &bef, Option::::None); + } +} + +proptest! { + #[test] + fn rebal_auth_set_rebal_auth_correct_pt( + (ix, bef) in admin_correct_strat(), + ) { + silence_mollusk_logs(); + set_rebal_auth_test(&ix, &bef, Option::::None); + } +} + +proptest! { + #[test] + fn set_rebal_auth_unauthorized_pt( + (ix, bef) in unauthorized_strat(), + ) { + silence_mollusk_logs(); + set_rebal_auth_test( + &ix, + &bef, + Some(Inf1CtlCustomProgErr(Inf1CtlErr::UnauthorizedSetRebalanceAuthoritySigner)) + ); + } +} + +proptest! { + #[test] + fn set_rebal_auth_missing_sig_pt( + (ix, bef) in missing_sig_strat(), + ) { + silence_mollusk_logs(); + set_rebal_auth_test(&ix, &bef, Some(MISSING_REQUIRED_SIGNATURE)); + } +} + +proptest! { + #[test] + fn set_rebal_auth_disabled_pt( + (ix, bef) in disabled_strat(), + ) { + silence_mollusk_logs(); + set_rebal_auth_test( + &ix, + &bef, + Some(Inf1CtlCustomProgErr(Inf1CtlErr::PoolDisabled)) + ); + } +} + +proptest! { + #[test] + fn set_rebal_auth_rebalancing_pt( + (ix, bef) in rebalancing_strat(), + ) { + silence_mollusk_logs(); + set_rebal_auth_test( + &ix, + &bef, + Some(Inf1CtlCustomProgErr(Inf1CtlErr::PoolRebalancing)) + ); + } +}