add set treasury fee instruction

stake-pool/cli/src/main.rs

@@ -3832,6 +3832,51 @@ fn command_set_no_fee_deposit_threshold(
     Ok(())
 }
 
+fn command_set_treasury_fee_account(
+    config: &Config,
+    stake_pool_address: &Pubkey,
+) -> CommandResult {
+    let treasury_keypair = Keypair::new();
+    println!("Creating treasury {}", treasury_keypair.pubkey());
+
+    let token_account_rent_exempt = config
+        .rpc_client
+        .get_minimum_balance_for_rent_exemption(spl_token::state::Account::LEN)?;
+
+    let mut signers = vec![config.fee_payer.as_ref(), config.manager.as_ref(), &treasury_keypair];
+    unique_signers!(signers);
+
+    let stake_pool = get_stake_pool(&config.rpc_client, stake_pool_address)?;
+    let transaction = checked_transaction_with_signers(
+        config,
+        &[
+            // Create treasury account
+            system_instruction::create_account(
+                &config.fee_payer.pubkey(),
+                &treasury_keypair.pubkey(),
+                token_account_rent_exempt,
+                spl_token::state::Account::LEN as u64,
+                &spl_token::id(),
+            ),
+            // Initialize treasury account as token account
+            spl_token::instruction::initialize_account(
+                &spl_token::id(),
+                &treasury_keypair.pubkey(),
+                &stake_pool.pool_mint,
+                &config.manager.pubkey(),
+            )?,
+            spl_stake_pool::instruction::set_treasury_fee_account(
+            &spl_stake_pool::id(),
+            stake_pool_address,
+            &treasury_keypair.pubkey(),
+            &config.manager.pubkey(),
+        )],
+        &signers,
+    )?;
+    send_transaction(config, transaction)?;
+    Ok(())
+}
+
 fn command_list_all_pools(config: &Config) -> CommandResult {
     let all_pools = get_stake_pools(&config.rpc_client)?;
     let cli_stake_pool_vec: Vec<CliStakePool> =
@@ -6520,6 +6565,18 @@ fn main() {
                     .help("No fee deposit threshold."),
             )
         )
+        .subcommand(SubCommand::with_name("set-treasury-fee-account")
+            .about("Set or update treasury-fee-account. Must be signed by the manager.")
+            .arg(
+                Arg::with_name("pool")
+                    .index(1)
+                    .validator(is_pubkey)
+                    .value_name("POOL_ADDRESS")
+                    .takes_value(true)
+                    .required(true)
+                    .help("Stake pool address."),
+            )
+        )
         .subcommand(SubCommand::with_name("dao-strategy-mint-extra-community-tokens")
             .about("Mint extra community tokens for those users who didn't receive them by mistake")
             .arg(
@@ -7126,6 +7183,10 @@ fn main() {
 
             command_set_no_fee_deposit_threshold(&config, &stake_pool_address, threshold)
         }
+        ("set-treasury-fee-account", Some(arg_matches)) => {
+            let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();
+            command_set_treasury_fee_account(&config, &stake_pool_address)
+        }
         ("dao-strategy-mint-extra-community-tokens", Some(arg_matches)) => {
             let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();
             let to = pubkey_of(arg_matches, "to").unwrap();

stake-pool/program/src/instruction.rs

@@ -745,6 +745,15 @@ pub enum StakePoolInstruction {
     //
     // Index 39 
     DaoStrategyDepositSolWithReferrer2(u64),
+
+    /// (Manager only) Treasury fee account
+    ///
+    ///  0. `[w]` StakePool
+    ///  1. `[r]` TreasuryFeeAccount
+    ///  2. `[s]` Manager
+    ///
+    /// userdata: threshold
+    SetTreasuryFeeAccount(),
 }
 
 /// Creates an 'initialize' instruction.
@@ -1642,6 +1651,25 @@ pub fn set_no_fee_deposit_threshold(
     }
 }
 
+/// Creates a 'set treasury fee account' instruction.
+pub fn set_treasury_fee_account(
+    program_id: &Pubkey,
+    stake_pool: &Pubkey,
+    treasury_fee_account: &Pubkey,
+    manager: &Pubkey,
+) -> Instruction {
+    let accounts = vec![
+        AccountMeta::new(*stake_pool, false),
+        AccountMeta::new_readonly(*treasury_fee_account, false),
+        AccountMeta::new_readonly(*manager, true),
+    ];
+    Instruction {
+        program_id: *program_id,
+        accounts,
+        data: StakePoolInstruction::SetTreasuryFeeAccount().try_to_vec().unwrap(),
+    }
+}
+
 /// Creates a 'set staker' instruction.
 pub fn set_staker(
     program_id: &Pubkey,

stake-pool/program/src/processor.rs

@@ -2887,6 +2887,43 @@ impl Processor {
         Ok(())
     }
 
+    /// Changes the StakePool treasury fee account.
+    /// Сan only be performed by the StakePool manager.
+    ///
+    /// Processes [SetTreasuryFeeAccount](enum.Instruction.html).
+    #[inline(never)] // needed to avoid stack size violation
+    fn process_set_treasury_fee_account(
+        program_id: &Pubkey,
+        accounts: &[AccountInfo],
+    ) -> ProgramResult {
+        let account_info_iter = &mut accounts.iter();
+        let stake_pool_info = next_account_info(account_info_iter)?;
+        let treasury_fee_account_info = next_account_info(account_info_iter)?;
+        let manager_info = next_account_info(account_info_iter)?;
+
+        check_account_owner(stake_pool_info, program_id)?;
+        let mut stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data.borrow())?;
+        if !stake_pool.is_valid() {
+            return Err(StakePoolError::InvalidState.into());
+        }
+        stake_pool.check_manager(manager_info)?;
+
+        let treasury_fee = spl_token::state::Account::unpack_from_slice(&treasury_fee_account_info.data.borrow())?;
+        if treasury_fee.owner != *manager_info.key {
+            return Err(ProgramError::IllegalOwner.into());
+        }
+        if stake_pool.pool_mint
+            != treasury_fee
+            .mint
+        {
+            return Err(StakePoolError::WrongAccountMint.into());
+        }
+
+        stake_pool.treasury_fee_account = *treasury_fee_account_info.key;
+        stake_pool.serialize(&mut *stake_pool_info.data.borrow_mut())?;
+        Ok(())
+    }
+
     /// Changes the StakePool staker.
     /// Сan only be performed by the StakePool manager or staker.
     /// 
@@ -5790,7 +5827,11 @@ impl Processor {
             StakePoolInstruction::SetNoFeeDepositThreshold(no_fee_deposit_threshold) => {
                 msg!("Instruction: SetNoFeeDepositThreshold");
                 Self::process_set_no_fee_deposit_threshold(program_id, accounts, no_fee_deposit_threshold)
-            }       
+            }
+            StakePoolInstruction::SetTreasuryFeeAccount() => {
+                msg!("Instruction: SetTreasuryFeeAccount");
+                Self::process_set_treasury_fee_account(program_id, accounts)
+            }
         }
     }
 }