#![allow(clippy::integer_arithmetic)]
mod arg_parser;
mod args;
mod stake_accounts;
use {
crate::{
arg_parser::parse_args,
args::{
resolve_command, AuthorizeArgs, Command, MoveArgs, NewArgs, RebaseArgs, SetLockupArgs,
},
},
solana_cli_config::Config,
solana_client::{client_error::ClientError, rpc_client::RpcClient},
solana_sdk::{
message::Message,
native_token::lamports_to_sol,
pubkey::Pubkey,
signature::{unique_signers, Signature, Signer},
signers::Signers,
stake::{instruction::LockupArgs, state::Lockup},
transaction::Transaction,
},
solana_stake_program::stake_state,
std::{env, error::Error},
};
fn get_balance_at(client: &RpcClient, pubkey: &Pubkey, i: usize) -> Result<u64, ClientError> {
let address = stake_accounts::derive_stake_account_address(pubkey, i);
client.get_balance(&address)
}
fn count_stake_accounts(client: &RpcClient, base_pubkey: &Pubkey) -> Result<usize, ClientError> {
let mut i = 0;
while get_balance_at(client, base_pubkey, i)? > 0 {
i += 1;
}
Ok(i)
}
fn get_balances(
client: &RpcClient,
addresses: Vec<Pubkey>,
) -> Result<Vec<(Pubkey, u64)>, ClientError> {
addresses
.into_iter()
.map(|pubkey| client.get_balance(&pubkey).map(|bal| (pubkey, bal)))
.collect()
}
fn get_lockup(client: &RpcClient, address: &Pubkey) -> Result<Lockup, ClientError> {
client
.get_account(address)
.map(|account| stake_state::lockup_from(&account).unwrap())
}
fn get_lockups(
client: &RpcClient,
addresses: Vec<Pubkey>,
) -> Result<Vec<(Pubkey, Lockup)>, ClientError> {
addresses
.into_iter()
.map(|pubkey| get_lockup(client, &pubkey).map(|bal| (pubkey, bal)))
.collect()
}
fn process_new_stake_account(
client: &RpcClient,
args: &NewArgs<Pubkey, Box<dyn Signer>>,
) -> Result<Signature, ClientError> {
let message = stake_accounts::new_stake_account(
&args.fee_payer.pubkey(),
&args.funding_keypair.pubkey(),
&args.base_keypair.pubkey(),
args.lamports,
&args.stake_authority,
&args.withdraw_authority,
&Pubkey::default(),
args.index,
);
let signers = unique_signers(vec![
&*args.fee_payer,
&*args.funding_keypair,
&*args.base_keypair,
]);
let signature = send_and_confirm_message(client, message, &signers, false)?;
Ok(signature)
}
fn process_authorize_stake_accounts(
client: &RpcClient,
args: &AuthorizeArgs<Pubkey, Box<dyn Signer>>,
) -> Result<(), ClientError> {
let messages = stake_accounts::authorize_stake_accounts(
&args.fee_payer.pubkey(),
&args.base_pubkey,
&args.stake_authority.pubkey(),
&args.withdraw_authority.pubkey(),
&args.new_stake_authority,
&args.new_withdraw_authority,
args.num_accounts,
);
let signers = unique_signers(vec![
&*args.fee_payer,
&*args.stake_authority,
&*args.withdraw_authority,
]);
send_and_confirm_messages(client, messages, &signers, false)?;
Ok(())
}
fn process_lockup_stake_accounts(
client: &RpcClient,
args: &SetLockupArgs<Pubkey, Box<dyn Signer>>,
) -> Result<(), ClientError> {
let addresses =
stake_accounts::derive_stake_account_addresses(&args.base_pubkey, args.num_accounts);
let existing_lockups = get_lockups(client, addresses)?;
let lockup = LockupArgs {
epoch: args.lockup_epoch,
unix_timestamp: args.lockup_date,
custodian: args.new_custodian,
};
let messages = stake_accounts::lockup_stake_accounts(
&args.fee_payer.pubkey(),
&args.custodian.pubkey(),
&lockup,
&existing_lockups,
args.unlock_years,
);
if messages.is_empty() {
eprintln!("No work to do");
return Ok(());
}
let signers = unique_signers(vec![&*args.fee_payer, &*args.custodian]);
send_and_confirm_messages(client, messages, &signers, args.no_wait)?;
Ok(())
}
fn process_rebase_stake_accounts(
client: &RpcClient,
args: &RebaseArgs<Pubkey, Box<dyn Signer>>,
) -> Result<(), ClientError> {
let addresses =
stake_accounts::derive_stake_account_addresses(&args.base_pubkey, args.num_accounts);
let balances = get_balances(client, addresses)?;
let messages = stake_accounts::rebase_stake_accounts(
&args.fee_payer.pubkey(),
&args.new_base_keypair.pubkey(),
&args.stake_authority.pubkey(),
&balances,
);
if messages.is_empty() {
eprintln!("No accounts found");
return Ok(());
}
let signers = unique_signers(vec![
&*args.fee_payer,
&*args.new_base_keypair,
&*args.stake_authority,
]);
send_and_confirm_messages(client, messages, &signers, false)?;
Ok(())
}
fn process_move_stake_accounts(
client: &RpcClient,
move_args: &MoveArgs<Pubkey, Box<dyn Signer>>,
) -> Result<(), ClientError> {
let authorize_args = &move_args.authorize_args;
let args = &move_args.rebase_args;
let addresses =
stake_accounts::derive_stake_account_addresses(&args.base_pubkey, args.num_accounts);
let balances = get_balances(client, addresses)?;
let messages = stake_accounts::move_stake_accounts(
&args.fee_payer.pubkey(),
&args.new_base_keypair.pubkey(),
&args.stake_authority.pubkey(),
&authorize_args.withdraw_authority.pubkey(),
&authorize_args.new_stake_authority,
&authorize_args.new_withdraw_authority,
&balances,
);
if messages.is_empty() {
eprintln!("No accounts found");
return Ok(());
}
let signers = unique_signers(vec![
&*args.fee_payer,
&*args.new_base_keypair,
&*args.stake_authority,
&*authorize_args.withdraw_authority,
]);
send_and_confirm_messages(client, messages, &signers, false)?;
Ok(())
}
fn send_and_confirm_message<S: Signers>(
client: &RpcClient,
message: Message,
signers: &S,
no_wait: bool,
) -> Result<Signature, ClientError> {
let mut transaction = Transaction::new_unsigned(message);
let blockhash = client.get_new_latest_blockhash(&transaction.message().recent_blockhash)?;
transaction.try_sign(signers, blockhash)?;
if no_wait {
client.send_transaction(&transaction)
} else {
client.send_and_confirm_transaction_with_spinner(&transaction)
}
}
fn send_and_confirm_messages<S: Signers>(
client: &RpcClient,
messages: Vec<Message>,
signers: &S,
no_wait: bool,
) -> Result<Vec<Signature>, ClientError> {
let mut signatures = vec![];
for message in messages {
let signature = send_and_confirm_message(client, message, signers, no_wait)?;
signatures.push(signature);
println!("{}", signature);
}
Ok(signatures)
}
fn main() -> Result<(), Box<dyn Error>> {
let command_args = parse_args(env::args_os());
let config = Config::load(&command_args.config_file).unwrap_or_default();
let json_rpc_url = command_args.url.unwrap_or(config.json_rpc_url);
let client = RpcClient::new(json_rpc_url);
match resolve_command(&command_args.command)? {
Command::New(args) => {
process_new_stake_account(&client, &args)?;
}
Command::Count(args) => {
let num_accounts = count_stake_accounts(&client, &args.base_pubkey)?;
println!("{}", num_accounts);
}
Command::Addresses(args) => {
let addresses = stake_accounts::derive_stake_account_addresses(
&args.base_pubkey,
args.num_accounts,
);
for address in addresses {
println!("{:?}", address);
}
}
Command::Balance(args) => {
let addresses = stake_accounts::derive_stake_account_addresses(
&args.base_pubkey,
args.num_accounts,
);
let balances = get_balances(&client, addresses)?;
let lamports: u64 = balances.into_iter().map(|(_, bal)| bal).sum();
let sol = lamports_to_sol(lamports);
println!("{} SOL", sol);
}
Command::Authorize(args) => {
process_authorize_stake_accounts(&client, &args)?;
}
Command::SetLockup(args) => {
process_lockup_stake_accounts(&client, &args)?;
}
Command::Rebase(args) => {
process_rebase_stake_accounts(&client, &args)?;
}
Command::Move(args) => {
process_move_stake_accounts(&client, &args)?;
}
}
Ok(())
}