Added set-staking-auth and set-owner commands to stake pool CLI (#827)

* Added set-staking-auth and set-owner commands to stake pool CLI

* Error handling refactored, fixed new fee owner setting, fixed set owner arg group settings

* Added check for fee account mint to stake pool initialization and set owner methods

stake-pool/cli/src/main.rs

@@ -1,6 +1,6 @@
 use clap::{
     crate_description, crate_name, crate_version, value_t, value_t_or_exit, App, AppSettings, Arg,
-    SubCommand,
+    ArgGroup, SubCommand,
 };
 use solana_account_decoder::UiAccountEncoding;
 use solana_clap_utils::{
@@ -25,8 +25,8 @@ use solana_sdk::{
 };
 use spl_stake_pool::{
     instruction::{
-        claim, deposit, initialize as initialize_pool, withdraw, Fee as PoolFee,
-        InitArgs as PoolInitArgs,
+        claim, deposit, initialize as initialize_pool, set_owner, set_staking_authority, withdraw,
+        Fee as PoolFee, InitArgs as PoolInitArgs,
     },
     processor::Processor as PoolProcessor,
     stake::authorize as authorize_stake,
@@ -361,17 +361,17 @@ fn command_list(config: &Config, pool: &Pubkey) -> CommandResult {
     let accounts = get_authority_accounts(config, &pool_withdraw_authority);
 
     if accounts.is_empty() {
-        println!("No accounts found.")
-    } else {
-        let mut total_balance: u64 = 0;
-        for (pubkey, account) in accounts {
-            let balance = account.lamports;
-            total_balance += balance;
-            println!("{}\t{} SOL", pubkey, lamports_to_sol(balance));
-        }
-        println!("Total: {} SOL", lamports_to_sol(total_balance));
+        return Err("No accounts found.".to_string().into());
     }
 
+    let mut total_balance: u64 = 0;
+    for (pubkey, account) in accounts {
+        let balance = account.lamports;
+        total_balance += balance;
+        println!("{}\t{} SOL", pubkey, lamports_to_sol(balance));
+    }
+    println!("Total: {} SOL", lamports_to_sol(total_balance));
+
     Ok(None)
 }
 
@@ -419,25 +419,23 @@ fn command_withdraw(
         TokenAccount::unpack_from_slice(account_data.as_slice()).unwrap();
 
     if account_data.mint != pool_data.pool_mint {
-        println!("Wrong token account.");
-        return Ok(None);
+        return Err("Wrong token account.".to_string().into());
     }
 
     // Check burn_from balance
     let max_withdraw_amount = pool_tokens_to_stake_amount(&pool_data, account_data.amount);
     if max_withdraw_amount < amount {
-        println!(
+        return Err(format!(
             "Not enough token balance to withdraw {} SOL.\nMaximum withdraw amount is {} SOL.",
             lamports_to_sol(amount),
             lamports_to_sol(max_withdraw_amount)
-        );
-        return Ok(None);
+        )
+        .into());
     }
 
     let mut accounts = get_authority_accounts(config, &pool_withdraw_authority);
     if accounts.is_empty() {
-        println!("No accounts found.");
-        return Ok(None);
+        return Err("No accounts found.".to_string().into());
     }
     // Sort from lowest to highest balance
     accounts.sort_by(|a, b| a.1.lamports.cmp(&b.1.lamports));
@@ -555,11 +553,109 @@ fn command_withdraw(
         return Ok(Some(transaction));
     }
 
-    println!(
+    Err(format!(
         "No stake accounts found in this pool with enough balance to withdraw {} SOL.",
         lamports_to_sol(amount)
+    )
+    .into())
+}
+
+fn command_set_staking_auth(
+    config: &Config,
+    pool: &Pubkey,
+    stake_account: &Pubkey,
+    new_staker: &Pubkey,
+) -> CommandResult {
+    let pool_data = config.rpc_client.get_account_data(&pool)?;
+    let pool_data: StakePool = PoolState::deserialize(pool_data.as_slice())
+        .unwrap()
+        .stake_pool()
+        .unwrap();
+
+    let pool_withdraw_authority: Pubkey = PoolProcessor::authority_id(
+        &spl_stake_pool::id(),
+        pool,
+        PoolProcessor::AUTHORITY_WITHDRAW,
+        pool_data.withdraw_bump_seed,
+    )
+    .unwrap();
+
+    let mut transaction = Transaction::new_with_payer(
+        &[set_staking_authority(
+            &spl_stake_pool::id(),
+            &pool,
+            &config.owner.pubkey(),
+            &pool_withdraw_authority,
+            &stake_account,
+            &new_staker,
+            &stake_program_id(),
+        )?],
+        Some(&config.fee_payer.pubkey()),
     );
-    Ok(None)
+
+    let (recent_blockhash, fee_calculator) = config.rpc_client.get_recent_blockhash()?;
+    check_fee_payer_balance(config, fee_calculator.calculate_fee(&transaction.message()))?;
+    let mut signers = vec![config.fee_payer.as_ref(), config.owner.as_ref()];
+    unique_signers!(signers);
+    transaction.sign(&signers, recent_blockhash);
+    Ok(Some(transaction))
+}
+
+fn command_set_owner(
+    config: &Config,
+    pool: &Pubkey,
+    new_owner: &Option<Pubkey>,
+    new_fee_receiver: &Option<Pubkey>,
+) -> CommandResult {
+    let pool_data = config.rpc_client.get_account_data(&pool)?;
+    let pool_data: StakePool = PoolState::deserialize(pool_data.as_slice())
+        .unwrap()
+        .stake_pool()
+        .unwrap();
+
+    // If new accounts are missing in the arguments use the old ones
+    let new_owner: Pubkey = match new_owner {
+        None => pool_data.owner,
+        Some(value) => *value,
+    };
+    let new_fee_receiver: Pubkey = match new_fee_receiver {
+        None => pool_data.owner_fee_account,
+        Some(value) => {
+            // Check for fee receiver being a valid token account and have to same mint as the stake pool
+            let account_data = config.rpc_client.get_account_data(value)?;
+            let account_data: TokenAccount =
+                match TokenAccount::unpack_from_slice(account_data.as_slice()) {
+                    Ok(data) => data,
+                    Err(_) => {
+                        return Err(format!("{} is not a token account", value).into());
+                    }
+                };
+            if account_data.mint != pool_data.pool_mint {
+                return Err("Fee receiver account belongs to a different mint"
+                    .to_string()
+                    .into());
+            }
+            *value
+        }
+    };
+
+    let mut transaction = Transaction::new_with_payer(
+        &[set_owner(
+            &spl_stake_pool::id(),
+            &pool,
+            &config.owner.pubkey(),
+            &new_owner,
+            &new_fee_receiver,
+        )?],
+        Some(&config.fee_payer.pubkey()),
+    );
+
+    let (recent_blockhash, fee_calculator) = config.rpc_client.get_recent_blockhash()?;
+    check_fee_payer_balance(config, fee_calculator.calculate_fee(&transaction.message()))?;
+    let mut signers = vec![config.fee_payer.as_ref(), config.owner.as_ref()];
+    unique_signers!(signers);
+    transaction.sign(&signers, recent_blockhash);
+    Ok(Some(transaction))
 }
 
 fn main() {
@@ -719,6 +815,68 @@ fn main() {
                     .help("Stake account to receive SOL from the stake pool. Defaults to a new stake account."),
             )
         )
+        .subcommand(SubCommand::with_name("set-staking-auth").about("Changes staking authority of one of the accounts from the stake pool.")
+            .arg(
+                Arg::with_name("pool")
+                    .long("pool")
+                    .validator(is_pubkey)
+                    .value_name("ADDRESS")
+                    .takes_value(true)
+                    .required(true)
+                    .help("Stake pool address."),
+            )
+            .arg(
+                Arg::with_name("stake_account")
+                    .long("stake-account")
+                    .validator(is_pubkey)
+                    .value_name("ADDRESS")
+                    .takes_value(true)
+                    .required(true)
+                    .help("Stake account address to change staking authority."),
+            )
+            .arg(
+                Arg::with_name("new_staker")
+                    .long("new-staker")
+                    .validator(is_pubkey)
+                    .value_name("ADDRESS")
+                    .takes_value(true)
+                    .required(true)
+                    .help("Public key of the new staker account."),
+            )
+        )
+        .subcommand(SubCommand::with_name("set-owner").about("Changes owner or fee receiver account for the stake pool.")
+            .arg(
+                Arg::with_name("pool")
+                    .long("pool")
+                    .validator(is_pubkey)
+                    .value_name("ADDRESS")
+                    .takes_value(true)
+                    .required(true)
+                    .help("Stake pool address."),
+            )
+            .arg(
+                Arg::with_name("new_owner")
+                    .long("new-owner")
+                    .validator(is_pubkey)
+                    .value_name("ADDRESS")
+                    .takes_value(true)
+                    .help("Public key for the new stake pool owner."),
+            )
+            .arg(
+                Arg::with_name("new_fee_receiver")
+                    .long("new-fee-receiver")
+                    .validator(is_pubkey)
+                    .value_name("ADDRESS")
+                    .takes_value(true)
+                    .help("Public key for the new account to set as the stake pool fee receiver."),
+            )
+            .group(ArgGroup::with_name("new_accounts")
+                .arg("new_owner")
+                .arg("new_fee_receiver")
+                .required(true)
+                .multiple(true)
+            )
+        )
         .get_matches();
 
     let mut wallet_manager = None;
@@ -793,6 +951,18 @@ fn main() {
             let stake_receiver: Option<Pubkey> = pubkey_of(arg_matches, "stake_receiver");
             command_withdraw(&config, &pool_account, amount, &burn_from, &stake_receiver)
         }
+        ("set-staking-auth", Some(arg_matches)) => {
+            let pool_account: Pubkey = pubkey_of(arg_matches, "pool").unwrap();
+            let stake_account: Pubkey = pubkey_of(arg_matches, "stake_account").unwrap();
+            let new_staker: Pubkey = pubkey_of(arg_matches, "new_staker").unwrap();
+            command_set_staking_auth(&config, &pool_account, &stake_account, &new_staker)
+        }
+        ("set-owner", Some(arg_matches)) => {
+            let pool_account: Pubkey = pubkey_of(arg_matches, "pool").unwrap();
+            let new_owner: Option<Pubkey> = pubkey_of(arg_matches, "new_owner");
+            let new_fee_receiver: Option<Pubkey> = pubkey_of(arg_matches, "new_fee_receiver");
+            command_set_owner(&config, &pool_account, &new_owner, &new_fee_receiver)
+        }
         _ => unreachable!(),
     }
     .and_then(|transaction| {

stake-pool/program/src/error.rs

@@ -43,6 +43,9 @@ pub enum Error {
     /// Stake pool fee > 1.
     #[error("FeeTooHigh")]
     FeeTooHigh,
+    /// Token account is associated with the wrong mint.
+    #[error("WrongAccountMint")]
+    WrongAccountMint,
 }
 impl From<Error> for ProgramError {
     fn from(e: Error) -> Self {

stake-pool/program/src/processor.rs

@@ -10,7 +10,7 @@ use num_traits::FromPrimitive;
 use solana_program::{
     account_info::next_account_info, account_info::AccountInfo, decode_error::DecodeError,
     entrypoint::ProgramResult, info, program::invoke_signed, program_error::PrintProgramError,
-    program_error::ProgramError, pubkey::Pubkey,
+    program_error::ProgramError, program_pack::Pack, pubkey::Pubkey,
 };
 use std::convert::TryFrom;
 
@@ -184,6 +184,13 @@ impl Processor {
             return Err(Error::FeeTooHigh.into());
         }
 
+        // Check for owner fee account to have proper mint assigned
+        if *pool_mint_info.key
+            != spl_token::state::Account::unpack_from_slice(&owner_fee_info.data.borrow())?.mint
+        {
+            return Err(Error::WrongAccountMint.into());
+        }
+
         let (_, deposit_bump_seed) = Self::find_authority_bump_seed(
             program_id,
             stake_pool_info.key,
@@ -557,6 +564,14 @@ impl Processor {
         if !owner_info.is_signer {
             return Err(Error::InvalidInput.into());
         }
+
+        // Check for owner fee account to have proper mint assigned
+        if stake_pool.pool_mint
+            != spl_token::state::Account::unpack_from_slice(&new_owner_fee_info.data.borrow())?.mint
+        {
+            return Err(Error::WrongAccountMint.into());
+        }
+
         stake_pool.owner = *new_owner_info.key;
         stake_pool.owner_fee_account = *new_owner_fee_info.key;
         State::Init(stake_pool).serialize(&mut stake_pool_info.data.borrow_mut())?;
@@ -612,6 +627,7 @@ impl PrintProgramError for Error {
             Error::InvalidOutput => info!("Error: InvalidOutput"),
             Error::CalculationFailure => info!("Error: CalculationFailure"),
             Error::FeeTooHigh => info!("Error: FeeTooHigh"),
+            Error::WrongAccountMint => info!("Error: WrongAccountMint"),
         }
     }
 }