From 84ee7cb704a7e3ff4fabba68e4de5573a3a3a61f Mon Sep 17 00:00:00 2001 From: Samuel Vanderwaal Date: Wed, 27 Dec 2023 20:15:00 -0900 Subject: [PATCH] add spl airdrop (#307) * add spl airdrop * bump to Jib 0.5.0; include priority fee option * fix Cargo toml file --- Cargo.lock | 4 +- Cargo.toml | 4 +- src/airdrop/mod.rs | 127 +++----------------- src/airdrop/sol.rs | 109 +++++++++++++++++ src/airdrop/spl.rs | 234 +++++++++++++++++++++++++++++++++++++ src/opt.rs | 30 +++++ src/process_subcommands.rs | 23 +++- 7 files changed, 412 insertions(+), 119 deletions(-) create mode 100644 src/airdrop/sol.rs create mode 100644 src/airdrop/spl.rs diff --git a/Cargo.lock b/Cargo.lock index 3d7e75d..0cd580d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1799,9 +1799,9 @@ checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c" [[package]] name = "jib" -version = "0.4.1" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e66683adc02208fbff00788190de136dec6bb61c254f5814eed5d56178b7330" +checksum = "38488cc7a8673de9e2a39c073d8dad9055327f00a336f44bdd0aee413eed39b8" dependencies = [ "bincode", "indicatif 0.17.7", diff --git a/Cargo.toml b/Cargo.toml index 23e3ee1..7965a9d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,12 +19,13 @@ futures = "0.3.29" glob = "0.3.1" indexmap = { version = "1.9.3", features = ["serde"] } indicatif = { version = "0.16.2", features = ["rayon"] } -jib = "0.4.1" +jib = "0.5.0" lazy_static = "1.4.0" log = "0.4.20" metaboss_lib = "0.17.0" mpl-token-metadata = { version = "3.2.3", features = ["serde"] } num_cpus = "1.16.0" +once_cell = "1.19.0" phf = { version = "0.10", features = ["macros"] } ratelimit = "0.4.4" rayon = "1.8.0" @@ -45,7 +46,6 @@ structopt = "0.3.26" thiserror = "1.0.51" tokio = "1.35.1" regex = "1.10.2" -once_cell = "1.19.0" [features] diff --git a/src/airdrop/mod.rs b/src/airdrop/mod.rs index edf5ab0..df23acd 100644 --- a/src/airdrop/mod.rs +++ b/src/airdrop/mod.rs @@ -1,21 +1,18 @@ -use std::{collections::HashMap, fs::File, path::PathBuf, str::FromStr}; +pub mod sol; +pub mod spl; +pub use sol::*; +pub use spl::*; -use anyhow::Result; -use jib::{Jib, Network}; -use log::debug; -use serde::{Deserialize, Serialize}; -use solana_client::rpc_client::RpcClient; -use solana_sdk::{pubkey::Pubkey, signer::Signer}; +pub use std::{collections::HashMap, fs::File, path::PathBuf, str::FromStr}; -use crate::update::{parse_keypair, parse_solana_config}; +pub use anyhow::Result; +pub use jib::{Jib, Network}; +pub use log::debug; +pub use serde::{Deserialize, Serialize}; +pub use solana_client::rpc_client::RpcClient; +pub use solana_sdk::{pubkey::Pubkey, signer::Signer}; -pub struct AirdropSolArgs { - pub client: RpcClient, - pub keypair: Option, - pub network: Network, - pub recipient_list: Option, - pub cache_file: Option, -} +pub use crate::update::{parse_keypair, parse_solana_config}; #[derive(Debug, Clone, Deserialize, Serialize)] struct FailedTransaction { @@ -30,102 +27,4 @@ struct Recipient { amount: u64, } -pub async fn airdrop_sol(args: AirdropSolArgs) -> Result<()> { - let solana_opts = parse_solana_config(); - let keypair = parse_keypair(args.keypair, solana_opts); - - let mut jib = Jib::new(vec![keypair], args.network)?; - - let mut instructions = vec![]; - - if args.recipient_list.is_some() && args.cache_file.is_some() { - eprintln!("Cannot provide both a recipient list and a cache file."); - std::process::exit(1); - } - - // Get the current time as yyyy-mm-dd-hh-mm-ss - let now = chrono::Local::now(); - let timestamp = now.format("%Y-%m-%d-%H-%M-%S").to_string(); - - let mut cache_file_name = format!("mb-cache-airdrop-{timestamp}.json"); - let successful_tx_file_name = format!("mb-successful-airdrops-{timestamp}.json"); - - let mut airdrop_list: HashMap = if let Some(list_file) = args.recipient_list { - serde_json::from_reader(File::open(list_file)?)? - } else if let Some(cache_file) = args.cache_file { - cache_file_name = PathBuf::from(cache_file.clone()) - .file_name() - .unwrap() - .to_str() - .unwrap() - .to_string(); - - let failed_txes: Vec = serde_json::from_reader(File::open(cache_file)?)?; - failed_txes - .iter() - .flat_map(|f| f.recipients.clone()) - .collect() - } else { - eprintln!("No recipient list or cache file provided."); - std::process::exit(1); - }; - - for (address, amount) in &airdrop_list { - let pubkey = match Pubkey::from_str(address) { - Ok(pubkey) => pubkey, - Err(_) => { - eprintln!("Invalid address: {}, skipping...", address); - continue; - } - }; - - instructions.push(solana_sdk::system_instruction::transfer( - &jib.payer().pubkey(), - &pubkey, - *amount, - )); - } - - jib.set_instructions(instructions); - let results = jib.hoist()?; - - if results.iter().any(|r| r.is_failure()) { - println!("Some transactions failed. Check the {cache_file_name} cache file for details."); - } - - let mut successes = vec![]; - let mut failures = vec![]; - - results.iter().for_each(|r| { - if r.is_failure() { - let tx = r.transaction().unwrap(); // Transactions exist on failures. - let account_keys = tx.message().account_keys.clone(); - let transaction_accounts = account_keys.iter().map(|k| k.to_string()).collect(); - - // All accounts except the first and last are recipients. - let recipients: HashMap = account_keys[1..account_keys.len() - 1] - .iter() - .map(|pubkey| pubkey.to_string()) - .map(|a| airdrop_list.remove_entry(&a).expect("Recipient not found")) - .collect(); - - failures.push(FailedTransaction { - transaction_accounts, - recipients, - error: r.error().unwrap(), // Errors exist on failures. - }) - } else { - debug!("Transaction successful: {}", r.signature().unwrap()); // Signatures exist on successes. - successes.push(r.signature().unwrap()); // Signatures exist on successes. - } - }); - - // Write cache file and successful transactions. - let successful_tx_file = std::fs::File::create(successful_tx_file_name)?; - serde_json::to_writer_pretty(successful_tx_file, &successes)?; - - let cache_file = std::fs::File::create(cache_file_name)?; - serde_json::to_writer_pretty(cache_file, &failures)?; - - Ok(()) -} +pub const PRIORITY_FEE: u64 = 25_000; diff --git a/src/airdrop/sol.rs b/src/airdrop/sol.rs new file mode 100644 index 0000000..c687728 --- /dev/null +++ b/src/airdrop/sol.rs @@ -0,0 +1,109 @@ +use super::*; + +pub struct AirdropSolArgs { + pub client: RpcClient, + pub keypair: Option, + pub network: Network, + pub recipient_list: Option, + pub cache_file: Option, +} + +pub async fn airdrop_sol(args: AirdropSolArgs) -> Result<()> { + let solana_opts = parse_solana_config(); + let keypair = parse_keypair(args.keypair, solana_opts); + + let mut jib = Jib::new(vec![keypair], args.network)?; + + let mut instructions = vec![]; + + if args.recipient_list.is_some() && args.cache_file.is_some() { + eprintln!("Cannot provide both a recipient list and a cache file."); + std::process::exit(1); + } + + // Get the current time as yyyy-mm-dd-hh-mm-ss + let now = chrono::Local::now(); + let timestamp = now.format("%Y-%m-%d-%H-%M-%S").to_string(); + + let mut cache_file_name = format!("mb-cache-airdrop-{timestamp}.json"); + let successful_tx_file_name = format!("mb-successful-airdrops-{timestamp}.json"); + + let mut airdrop_list: HashMap = if let Some(list_file) = args.recipient_list { + serde_json::from_reader(File::open(list_file)?)? + } else if let Some(cache_file) = args.cache_file { + cache_file_name = PathBuf::from(cache_file.clone()) + .file_name() + .unwrap() + .to_str() + .unwrap() + .to_string(); + + let failed_txes: Vec = serde_json::from_reader(File::open(cache_file)?)?; + failed_txes + .iter() + .flat_map(|f| f.recipients.clone()) + .collect() + } else { + eprintln!("No recipient list or cache file provided."); + std::process::exit(1); + }; + + for (address, amount) in &airdrop_list { + let pubkey = match Pubkey::from_str(address) { + Ok(pubkey) => pubkey, + Err(_) => { + eprintln!("Invalid address: {}, skipping...", address); + continue; + } + }; + + instructions.push(solana_sdk::system_instruction::transfer( + &jib.payer().pubkey(), + &pubkey, + *amount, + )); + } + + jib.set_instructions(instructions); + let results = jib.hoist()?; + + if results.iter().any(|r| r.is_failure()) { + println!("Some transactions failed. Check the {cache_file_name} cache file for details."); + } + + let mut successes = vec![]; + let mut failures = vec![]; + + results.iter().for_each(|r| { + if r.is_failure() { + let tx = r.transaction().unwrap(); // Transactions exist on failures. + let account_keys = tx.message().account_keys.clone(); + let transaction_accounts = account_keys.iter().map(|k| k.to_string()).collect(); + + // All accounts except the first and last are recipients. + let recipients: HashMap = account_keys[1..account_keys.len() - 1] + .iter() + .map(|pubkey| pubkey.to_string()) + .map(|a| airdrop_list.remove_entry(&a).expect("Recipient not found")) + .collect(); + + failures.push(FailedTransaction { + transaction_accounts, + recipients, + error: r.error().unwrap(), // Errors exist on failures. + }) + } else { + debug!("Transaction successful: {}", r.signature().unwrap()); // Signatures exist on successes. + successes.push(r.signature().unwrap()); // Signatures exist on successes. + } + }); + + // Write cache file and successful transactions. + let successful_tx_file = std::fs::File::create(successful_tx_file_name)?; + serde_json::to_writer_pretty(successful_tx_file, &successes)?; + + let cache_file = std::fs::File::create(cache_file_name)?; + serde_json::to_writer_pretty(cache_file, &failures)?; + + Ok(()) +} diff --git a/src/airdrop/spl.rs b/src/airdrop/spl.rs new file mode 100644 index 0000000..07e1ef8 --- /dev/null +++ b/src/airdrop/spl.rs @@ -0,0 +1,234 @@ +use borsh::{BorshDeserialize, BorshSerialize}; +use metaboss_lib::transaction::send_and_confirm_tx; +use solana_program::{ + instruction::{AccountMeta, Instruction}, + program_pack::Pack, + pubkey, system_program, +}; +use spl_associated_token_account::get_associated_token_address; +use spl_token::instruction::transfer_checked; + +use super::*; + +pub struct AirdropSplArgs { + pub client: RpcClient, + pub keypair: Option, + pub network: Network, + pub recipient_list: Option, + pub cache_file: Option, + pub mint: Pubkey, + pub mint_tokens: bool, + pub boost: bool, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +struct SplRecipient { + address: String, + ata: String, + amount: u64, +} + +type Recipient = String; +type Ata = String; + +pub async fn airdrop_spl(args: AirdropSplArgs) -> Result<()> { + let solana_opts = parse_solana_config(); + let keypair = parse_keypair(args.keypair, solana_opts); + + let mut jib = Jib::new(vec![keypair], args.network)?; + let mut instructions = vec![]; + + let mut recipients_lookup: HashMap = HashMap::new(); + + let source_ata = get_associated_token_address(&jib.payer().pubkey(), &args.mint); + + let mint_account = + spl_token::state::Mint::unpack(jib.rpc_client().get_account(&args.mint)?.data.as_slice())?; + let decimals = mint_account.decimals; + + if args.recipient_list.is_some() && args.cache_file.is_some() { + eprintln!("Cannot provide both a recipient list and a cache file."); + std::process::exit(1); + } + + // Get the current time as yyyy-mm-dd-hh-mm-ss + let now = chrono::Local::now(); + let timestamp = now.format("%Y-%m-%d-%H-%M-%S").to_string(); + + let mut cache_file_name = format!("mb-cache-airdrop-{timestamp}.json"); + let successful_tx_file_name = format!("mb-successful-airdrops-{timestamp}.json"); + + let airdrop_list: HashMap = if let Some(list_file) = args.recipient_list { + serde_json::from_reader(File::open(list_file)?)? + } else if let Some(cache_file) = args.cache_file { + cache_file_name = PathBuf::from(cache_file.clone()) + .file_name() + .unwrap() + .to_str() + .unwrap() + .to_string(); + + let failed_txes: Vec = serde_json::from_reader(File::open(cache_file)?)?; + failed_txes + .iter() + .flat_map(|f| f.recipients.clone()) + .collect() + } else { + eprintln!("No recipient list or cache file provided."); + std::process::exit(1); + }; + + if args.mint_tokens { + let total_tokens = airdrop_list.values().sum::(); + let total_tokens_native_units = total_tokens * 10u64.pow(decimals as u32); + + let mint_tokens_ix = spl_token::instruction::mint_to( + &spl_token::ID, + &args.mint, + &source_ata, + &jib.payer().pubkey(), + &[], + total_tokens_native_units, + )?; + send_and_confirm_tx(&args.client, &[jib.payer()], &[mint_tokens_ix])?; + } + + if args.boost { + jib.set_priority_fee(PRIORITY_FEE); + } + + for (address, amount) in &airdrop_list { + let amount_native_units = amount * 10u64.pow(decimals as u32); + + let pubkey = match Pubkey::from_str(address) { + Ok(pubkey) => pubkey, + Err(_) => { + eprintln!("Invalid address: {}, skipping...", address); + continue; + } + }; + + let destination_ata = get_associated_token_address(&pubkey, &args.mint); + + recipients_lookup.insert(destination_ata.to_string(), pubkey.to_string()); + + instructions.push(create_token_if_missing_instruction( + &jib.payer().pubkey(), + &destination_ata, + &args.mint, + &pubkey, + &destination_ata, + )); + + instructions.push(transfer_checked( + &spl_token::ID, + &source_ata, + &args.mint, + &destination_ata, + &jib.payer().pubkey(), + &[], + amount_native_units, + decimals, + )?); + } + + jib.set_instructions(instructions); + let results = jib.hoist()?; + + if results.iter().any(|r| r.is_failure()) { + println!("Some transactions failed. Check the {cache_file_name} cache file for details."); + } + + let mut successes = vec![]; + let mut failures = vec![]; + + results.iter().for_each(|r| { + if r.is_failure() { + let tx = r.transaction().unwrap(); // Transactions exist on failures. + let account_keys = tx.message().account_keys.clone(); + let transaction_accounts = account_keys.iter().map(|k| k.to_string()).collect(); + + // We iterate over all account keys and check if they are in the recipients lookup to find the + // pubkey associated with any ATAs. Then we use the pubkey to find the address and amount pair + // to build the airdrop list from the failures so they can be retried. + let recipients: HashMap = account_keys + .iter() + .map(|p| p.to_string()) + .filter_map(|s| { + recipients_lookup.get(&s).and_then(|a| { + airdrop_list + .get_key_value(a) + .map(|(address, amount)| (address.clone(), *amount)) + }) + }) + .collect(); + + failures.push(FailedTransaction { + transaction_accounts, + recipients, + error: r.error().unwrap(), // Errors exist on failures. + }) + } else { + debug!("Transaction successful: {}", r.signature().unwrap()); // Signatures exist on successes. + successes.push(r.signature().unwrap()); // Signatures exist on successes. + } + }); + + // Write cache file and successful transactions. + let successful_tx_file = std::fs::File::create(successful_tx_file_name)?; + serde_json::to_writer_pretty(successful_tx_file, &successes)?; + + let cache_file = std::fs::File::create(cache_file_name)?; + serde_json::to_writer_pretty(cache_file, &failures)?; + + Ok(()) +} + +const MPL_TOOLBOX_ID: Pubkey = pubkey!("TokExjvjJmhKaRBShsBAsbSvEWMA1AgUNK7ps4SAc2p"); + +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +#[rustfmt::skip] +pub enum TokenExtrasInstruction { + /// Creates a new associated token account for the given mint and owner, if and only if + /// the given token account does not exists and the token account is the same as the + /// associated token account. That way, clients can ensure that, after this instruction, + /// the token account will exists. + /// + /// Notice this instruction asks for both the token account and the associated token account (ATA) + /// These may or may not be the same account. Here are all the possible cases: + /// + /// - Token exists and Token is ATA: Instruction succeeds. + /// - Token exists and Token is not ATA: Instruction succeeds. + /// - Token does not exist and Token is ATA: Instruction creates the ATA account and succeeds. + /// - Token does not exist and Token is not ATA: Instruction fails as we cannot create a + /// non-ATA account without it being a signer. + /// + /// Note that additional checks are made to ensure that the token account provided + /// matches the mint account and owner account provided. + CreateTokenIfMissing, +} + +fn create_token_if_missing_instruction( + payer: &Pubkey, + token: &Pubkey, + mint: &Pubkey, + owner: &Pubkey, + ata: &Pubkey, +) -> Instruction { + Instruction { + program_id: MPL_TOOLBOX_ID, + accounts: vec![ + AccountMeta::new(*payer, true), + AccountMeta::new_readonly(*token, false), + AccountMeta::new_readonly(*mint, false), + AccountMeta::new_readonly(*owner, false), + AccountMeta::new(*ata, false), + AccountMeta::new_readonly(system_program::id(), false), + AccountMeta::new_readonly(spl_token::id(), false), + AccountMeta::new_readonly(spl_associated_token_account::id(), false), + ], + data: TokenExtrasInstruction::CreateTokenIfMissing + .try_to_vec() + .unwrap(), + } +} diff --git a/src/opt.rs b/src/opt.rs index e0bf99a..c62d1a0 100644 --- a/src/opt.rs +++ b/src/opt.rs @@ -179,6 +179,36 @@ pub enum AirdropSubcommands { #[structopt(short, long)] cache_file: Option, }, + /// Airdrop SPL tokens + #[structopt(name = "spl")] + Spl { + /// Path to the owner keypair file + #[structopt(short, long)] + keypair: Option, + + /// Network cluster to use, defaults to devnet + #[structopt(short, long, default_value = "devnet")] + network: String, + + /// Path to the mint list file + #[structopt(short = "L", long)] + recipient_list: Option, + + /// Cache file + #[structopt(short, long)] + cache_file: Option, + + /// Mint from the SPL token mint + #[structopt(short, long)] + mint: Pubkey, + + #[structopt(long)] + mint_tokens: bool, + + /// Boost the transactions w/ priority fees + #[structopt(long)] + boost: bool, + }, } #[derive(Debug, StructOpt)] diff --git a/src/process_subcommands.rs b/src/process_subcommands.rs index 4322854..1f5b780 100644 --- a/src/process_subcommands.rs +++ b/src/process_subcommands.rs @@ -9,7 +9,7 @@ use metaboss_lib::decode::{ use mpl_token_metadata::types::Data; use solana_client::{nonblocking::rpc_client::RpcClient as AsyncRpcClient, rpc_client::RpcClient}; -use crate::airdrop::{airdrop_sol, AirdropSolArgs}; +use crate::airdrop::*; use crate::burn::*; use crate::collections::{ approve_delegate, check_collection_items, get_collection_items, migrate_collection, @@ -203,6 +203,27 @@ pub async fn process_airdrop(client: RpcClient, commands: AirdropSubcommands) -> }) .await } + AirdropSubcommands::Spl { + keypair, + recipient_list, + network, + cache_file, + mint, + mint_tokens, + boost, + } => { + airdrop_spl(AirdropSplArgs { + client, + keypair, + network: Network::from_str(&network).unwrap(), + recipient_list, + cache_file, + mint, + mint_tokens, + boost, + }) + .await + } } }