poc_framework/
lib.rs

1use crate::solana_sdk::clock::UnixTimestamp;
2use std::convert::TryFrom;
3use std::{
4    collections::{HashMap, HashSet},
5    convert::TryInto,
6    path::Path,
7    time::{SystemTime, UNIX_EPOCH},
8};
9
10use borsh::BorshDeserialize;
11use bpf_loader_upgradeable::UpgradeableLoaderState;
12use itertools::izip;
13use rand::{prelude::StdRng, rngs::OsRng, SeedableRng};
14use serde::de::DeserializeOwned;
15use sha2::{Digest, Sha256};
16use solana_bpf_loader_program::{
17    solana_bpf_loader_deprecated_program, solana_bpf_loader_program,
18    solana_bpf_loader_upgradeable_program,
19};
20use solana_cli_output::display::println_transaction;
21use solana_client::{rpc_client::RpcClient, rpc_config::RpcTransactionConfig};
22use solana_program::{
23    bpf_loader, bpf_loader_upgradeable,
24    hash::Hash,
25    instruction::Instruction,
26    loader_instruction,
27    message::{v0::LoadedAddresses, Message, SanitizedMessage},
28    program_option::COption,
29    program_pack::Pack,
30    pubkey::Pubkey,
31    system_instruction, system_program,
32    sysvar::{self, rent},
33};
34use solana_runtime::{
35    accounts_db::AccountShrinkThreshold,
36    accounts_index::AccountSecondaryIndexes,
37    bank::{
38        Bank, TransactionBalancesSet, TransactionExecutionDetails, TransactionExecutionResult,
39        TransactionResults,
40    },
41    builtins::{Builtin, Builtins},
42    genesis_utils,
43};
44use solana_sdk::{
45    account::{Account, AccountSharedData},
46    commitment_config::CommitmentConfig,
47    genesis_config::GenesisConfig,
48    packet,
49    signature::Keypair,
50    signature::Signer,
51    system_transaction,
52    transaction::{Transaction, VersionedTransaction},
53};
54use solana_transaction_status::{
55    token_balances, ConfirmedTransactionWithStatusMeta, EncodedConfirmedTransactionWithStatusMeta,
56    InnerInstructions, TransactionStatusMeta, TransactionTokenBalance, TransactionWithStatusMeta,
57    UiTransactionEncoding, VersionedTransactionWithStatusMeta,
58};
59use spl_associated_token_account::get_associated_token_address;
60
61pub use bincode;
62pub use borsh;
63pub use serde;
64pub use solana_client;
65pub use solana_program;
66pub use solana_sdk;
67pub use solana_transaction_status;
68pub use spl_associated_token_account;
69pub use spl_memo;
70pub use spl_token;
71
72mod keys;
73mod programs;
74
75/// A generic Environment trait. Provides the possibility of writing generic exploits that work both remote and local, for easy debugging.
76pub trait Environment {
77    /// Returns the keypair used to pay for all transactions. All transaction fees and rent costs are payed for by this keypair.
78    fn payer(&self) -> Keypair;
79    /// Executes the batch of transactions in the right order and waits for them to be confirmed. The execution results are returned.
80    fn execute_transaction(
81        &mut self,
82        txs: Transaction,
83    ) -> EncodedConfirmedTransactionWithStatusMeta;
84    /// Fetch a recent blockhash, for construction of transactions.
85    #[deprecated(since = "0.2.0", note = "Please use `get_latest_blockhash()` instead")]
86    fn get_recent_blockhash(&self) -> Hash {
87        self.get_latest_blockhash()
88    }
89    /// Fetch the latest blockhash, for construction of transactions.
90    fn get_latest_blockhash(&self) -> Hash;
91    /// Fetch the amount of lamports needed for an account of the given size to be rent excempt.
92    fn get_rent_excemption(&self, data: usize) -> u64;
93    /// Fetch an account. None if the account does not exist.
94    fn get_account(&self, pubkey: Pubkey) -> Option<Account>;
95
96    /// Assemble the given instructions into a transaction and sign it. All transactions constructed by this method are signed and payed for by the payer.
97    fn tx_with_instructions(
98        &self,
99        instructions: &[Instruction],
100        signers: &[&Keypair],
101    ) -> Transaction {
102        let payer = self.payer();
103        let mut signer_vec = vec![&payer];
104        signer_vec.extend_from_slice(signers);
105
106        let message = Message::new(instructions, Some(&self.payer().pubkey()));
107        let num_sigs: usize = message.header.num_required_signatures.into();
108        let required_sigs = message.account_keys[..num_sigs]
109            .into_iter()
110            .copied()
111            .collect::<HashSet<_>>();
112        let provided_sigs = signer_vec
113            .iter()
114            .map(|x| x.pubkey())
115            .collect::<HashSet<_>>();
116
117        for key in required_sigs.difference(&provided_sigs) {
118            println!("missing signature from {}", key.to_string());
119        }
120
121        for key in provided_sigs.difference(&required_sigs) {
122            println!("unnecessary signature from {}", key.to_string());
123        }
124
125        Transaction::new(&signer_vec, message, self.get_latest_blockhash())
126    }
127
128    /// Assemble the given instructions into a transaction and sign it. All transactions executed by this method are signed and payed for by the payer.
129    fn execute_as_transaction(
130        &mut self,
131        instructions: &[Instruction],
132        signers: &[&Keypair],
133    ) -> EncodedConfirmedTransactionWithStatusMeta {
134        let tx = self.tx_with_instructions(instructions, signers);
135        return self.execute_transaction(tx);
136    }
137
138    /// Assemble the given instructions into a transaction and sign it. All transactions executed by this method are signed and payed for by the payer.
139    /// Prints the transaction before sending it.
140    fn execute_as_transaction_debug(
141        &mut self,
142        instructions: &[Instruction],
143        signers: &[&Keypair],
144    ) -> EncodedConfirmedTransactionWithStatusMeta {
145        let tx = self.tx_with_instructions(instructions, signers);
146        println!("{:#?}", &tx);
147        return self.execute_transaction(tx);
148    }
149
150    /// Executes a transaction constructing an empty account with the specified amount of space and lamports, owned by the provided program.
151    fn create_account(&mut self, keypair: &Keypair, lamports: u64, space: usize, owner: Pubkey) {
152        self.execute_transaction(system_transaction::create_account(
153            &self.payer(),
154            &keypair,
155            self.get_latest_blockhash(),
156            lamports,
157            space as u64,
158            &owner,
159        ))
160        .assert_success();
161    }
162
163    /// Executes a transaction constructing an empty rent-excempt account with the specified amount of space, owned by the provided program.
164    fn create_account_rent_excempt(&mut self, keypair: &Keypair, space: usize, owner: Pubkey) {
165        self.execute_transaction(system_transaction::create_account(
166            &self.payer(),
167            &keypair,
168            self.get_latest_blockhash(),
169            self.get_rent_excemption(space),
170            space as u64,
171            &owner,
172        ))
173        .assert_success();
174    }
175
176    /// Executes a transaction constructing a token mint. The account needs to be empty and belong to system for this to work.
177    fn create_token_mint(
178        &mut self,
179        mint: &Keypair,
180        authority: Pubkey,
181        freeze_authority: Option<Pubkey>,
182        decimals: u8,
183    ) {
184        self.execute_as_transaction(
185            &[
186                system_instruction::create_account(
187                    &self.payer().pubkey(),
188                    &mint.pubkey(),
189                    self.get_rent_excemption(spl_token::state::Mint::LEN),
190                    spl_token::state::Mint::LEN as u64,
191                    &spl_token::ID,
192                ),
193                spl_token::instruction::initialize_mint(
194                    &spl_token::ID,
195                    &mint.pubkey(),
196                    &authority,
197                    freeze_authority.as_ref(),
198                    decimals,
199                )
200                .unwrap(),
201            ],
202            &[mint],
203        )
204        .assert_success();
205    }
206
207    /// Executes a transaction that mints tokens from a mint to an account belonging to that mint.
208    fn mint_tokens(&mut self, mint: Pubkey, authority: &Keypair, account: Pubkey, amount: u64) {
209        self.execute_as_transaction(
210            &[spl_token::instruction::mint_to(
211                &spl_token::ID,
212                &mint,
213                &account,
214                &authority.pubkey(),
215                &[],
216                amount,
217            )
218            .unwrap()],
219            &[authority],
220        )
221        .assert_success();
222    }
223
224    /// Executes a transaction constructing a token account of the specified mint. The account needs to be empty and belong to system for this to work.
225    /// Prefer to use [create_associated_token_account] if you don't need the provided account to contain the token account.
226    fn create_token_account(&mut self, account: &Keypair, mint: Pubkey) {
227        self.execute_as_transaction(
228            &[
229                system_instruction::create_account(
230                    &self.payer().pubkey(),
231                    &account.pubkey(),
232                    self.get_rent_excemption(spl_token::state::Account::LEN),
233                    spl_token::state::Account::LEN as u64,
234                    &spl_token::ID,
235                ),
236                spl_token::instruction::initialize_account(
237                    &spl_token::ID,
238                    &account.pubkey(),
239                    &mint,
240                    &account.pubkey(),
241                )
242                .unwrap(),
243            ],
244            &[account],
245        )
246        .assert_success();
247    }
248
249    /// Executes a transaction constructing the associated token account of the specified mint belonging to the owner. This will fail if the account already exists.
250    fn create_associated_token_account(&mut self, owner: &Keypair, mint: Pubkey) -> Pubkey {
251        self.execute_as_transaction(
252            &[
253                spl_associated_token_account::instruction::create_associated_token_account(
254                    &self.payer().pubkey(),
255                    &owner.pubkey(),
256                    &mint,
257                ),
258            ],
259            &[],
260        );
261        get_associated_token_address(&owner.pubkey(), &mint)
262    }
263
264    /// Executes a transaction constructing the associated token account of the specified mint belonging to the owner.
265    fn get_or_create_associated_token_account(&mut self, owner: &Keypair, mint: Pubkey) -> Pubkey {
266        let acc = get_associated_token_address(&owner.pubkey(), &mint);
267        if self.get_account(acc).is_none() {
268            self.create_associated_token_account(owner, mint);
269        }
270        acc
271    }
272
273    /// Executes a transaction creating and filling the given account with the given data.
274    /// The account is required to be empty and will be owned by bpf_loader afterwards.
275    fn create_account_with_data(&mut self, account: &Keypair, data: Vec<u8>) {
276        self.execute_transaction(system_transaction::create_account(
277            &self.payer(),
278            account,
279            self.get_latest_blockhash(),
280            self.get_rent_excemption(data.len()),
281            data.len() as u64,
282            &bpf_loader::id(),
283        ))
284        .assert_success();
285
286        let mut offset = 0usize;
287        for chunk in data.chunks(900) {
288            println!("writing bytes {} to {}", offset, offset + chunk.len());
289            self.execute_as_transaction(
290                &[loader_instruction::write(
291                    &account.pubkey(),
292                    &bpf_loader::id(),
293                    offset as u32,
294                    chunk.to_vec(),
295                )],
296                &[account],
297            )
298            .assert_success();
299            offset += chunk.len();
300        }
301    }
302
303    /// Executes a transaction deploying a program from a file if it does not already exist.
304    /// The keypair is derived from the file contents.
305    fn deploy_program<P: AsRef<Path>>(&mut self, program_path: P) -> Pubkey {
306        let data = std::fs::read(program_path).unwrap();
307        let mut hash = Sha256::default();
308        hash.update(&data);
309        let mut rng = StdRng::from_seed(hash.finalize()[..].try_into().unwrap());
310        let keypair = Keypair::generate(&mut rng);
311
312        if self.get_account(keypair.pubkey()).is_none() {
313            self.create_account_with_data(&keypair, data);
314            self.execute_as_transaction(
315                &[loader_instruction::finalize(
316                    &keypair.pubkey(),
317                    &bpf_loader::id(),
318                )],
319                &[&keypair],
320            )
321            .assert_success();
322        }
323
324        keypair.pubkey()
325    }
326
327    /// Gets and unpacks an account. None if the account does not exist.
328    fn get_unpacked_account<T: Pack>(&self, pubkey: Pubkey) -> Option<T> {
329        let acc = self.get_account(pubkey)?;
330        Some(T::unpack_unchecked(&acc.data).unwrap())
331    }
332
333    /// Gets and deserializes an account. None if the account does not exist.
334    fn get_deserialized_account<T: BorshDeserialize>(&self, pubkey: Pubkey) -> Option<T> {
335        let acc = self.get_account(pubkey)?;
336        Some(T::try_from_slice(&acc.data).unwrap())
337    }
338
339    /// Gets and deserializes an account. None if the account does not exist.
340    fn get_serde_deserialized_account<'a, T: DeserializeOwned>(&self, pubkey: Pubkey) -> Option<T> {
341        let acc = self.get_account(pubkey)?;
342        Some(bincode::deserialize(&acc.data).unwrap())
343    }
344}
345
346/// An clean environment that executes transactions locally. Good for testing and debugging.
347/// This environment has the most important SPL programs: spl-token, spl-associated-token-account and spl-memo v1 and v3.
348pub struct LocalEnvironment {
349    bank: Bank,
350    faucet: Keypair,
351}
352
353impl LocalEnvironment {
354    /// Constructs a builder for a local environment
355    pub fn builder() -> LocalEnvironmentBuilder {
356        LocalEnvironmentBuilder::new()
357    }
358
359    /// Constructs a clean local environment.
360    pub fn new() -> LocalEnvironment {
361        Self::builder().build()
362    }
363
364    pub fn bank(&mut self) -> &mut Bank {
365        &mut self.bank
366    }
367
368    /// Advance the bank to the next blockhash.
369    pub fn advance_blockhash(&self) -> Hash {
370        let parent_distance = if self.bank.slot() == 0 {
371            1
372        } else {
373            self.bank.slot() - self.bank.parent_slot()
374        };
375
376        for _ in 0..parent_distance {
377            let last_blockhash = self.bank.last_blockhash();
378            while self.bank.last_blockhash() == last_blockhash {
379                self.bank.register_tick(&Hash::new_unique())
380            }
381        }
382
383        self.get_latest_blockhash()
384    }
385}
386
387impl Environment for LocalEnvironment {
388    fn payer(&self) -> Keypair {
389        clone_keypair(&self.faucet)
390    }
391
392    fn execute_transaction(
393        &mut self,
394        tx: Transaction,
395    ) -> EncodedConfirmedTransactionWithStatusMeta {
396        let len = bincode::serialize(&tx).unwrap().len();
397        if len > packet::PACKET_DATA_SIZE {
398            panic!(
399                "tx {:?} of size {} is {} too large",
400                tx,
401                len,
402                len - packet::PACKET_DATA_SIZE
403            )
404        }
405        let txs = vec![tx];
406
407        let batch = self.bank.prepare_batch_for_tests(txs.clone());
408        let mut mint_decimals = HashMap::new();
409        let tx_pre_token_balances =
410            token_balances::collect_token_balances(&self.bank, &batch, &mut mint_decimals);
411        let slot = self.bank.slot();
412        let mut timings = Default::default();
413        let (
414            TransactionResults {
415                execution_results, ..
416            },
417            TransactionBalancesSet {
418                pre_balances,
419                post_balances,
420                ..
421            },
422        ) = self.bank.load_execute_and_commit_transactions(
423            &batch,
424            usize::MAX,
425            true,
426            true,
427            true,
428            true,
429            &mut timings,
430        );
431
432        let tx_post_token_balances =
433            token_balances::collect_token_balances(&self.bank, &batch, &mut mint_decimals);
434        izip!(
435            txs.iter(),
436            execution_results.into_iter(),
437            pre_balances.into_iter(),
438            post_balances.into_iter(),
439            tx_pre_token_balances.into_iter(),
440            tx_post_token_balances.into_iter(),
441        )
442        .map(
443            |(
444                tx,
445                execution_result,
446                pre_balances,
447                post_balances,
448                pre_token_balances,
449                post_token_balances,
450            ): (
451                &Transaction,
452                TransactionExecutionResult,
453                Vec<u64>,
454                Vec<u64>,
455                Vec<TransactionTokenBalance>,
456                Vec<TransactionTokenBalance>,
457            )| {
458                let fee = self.bank.get_fee_for_message(&SanitizedMessage::try_from(tx.message().clone()).expect("Failed to sanitize transaction"))
459                    .expect("Fee calculation must succeed");
460
461                let (status, inner_instructions, log_messages) = match execution_result {
462                    TransactionExecutionResult::Executed { details: TransactionExecutionDetails { status, inner_instructions, log_messages, .. }, .. } =>
463                        (status, inner_instructions, log_messages),
464                    TransactionExecutionResult::NotExecuted(err) => (Err(err), None, None)
465                };
466
467                let inner_instructions = inner_instructions.map(|inner_instructions| {
468                    inner_instructions
469                        .into_iter()
470                        .enumerate()
471                        .map(|(index, instructions)| InnerInstructions {
472                            index: index as u8,
473                            instructions,
474                        })
475                        .filter(|i| !i.instructions.is_empty())
476                        .collect()
477                });
478
479                let tx_status_meta = TransactionStatusMeta {
480                    status,
481                    fee,
482                    pre_balances,
483                    post_balances,
484                    pre_token_balances: Some(pre_token_balances),
485                    post_token_balances: Some(post_token_balances),
486                    inner_instructions,
487                    log_messages,
488                    rewards: None,
489                    loaded_addresses: LoadedAddresses {
490                        writable: vec![], // TODO
491                        readonly: vec![], // TODO
492                    },
493                    return_data: None
494                };
495
496                ConfirmedTransactionWithStatusMeta {
497                    slot,
498                    tx_with_meta: TransactionWithStatusMeta::Complete(VersionedTransactionWithStatusMeta {
499                        transaction: VersionedTransaction::from(tx.clone()),
500                        meta: tx_status_meta,
501                    }),
502                    block_time: Some(
503                        SystemTime::now()
504                            .duration_since(UNIX_EPOCH)
505                            .unwrap()
506                            .as_secs()
507                            .try_into()
508                            .unwrap(),
509                    ),
510                }
511                .encode(UiTransactionEncoding::Binary, None)
512                .expect("Failed to encode transaction")
513            },
514        )
515        .next().expect("transaction could not be executed. Enable debug logging to get more information on why")
516    }
517
518    fn get_latest_blockhash(&self) -> Hash {
519        self.bank.last_blockhash()
520    }
521
522    fn get_rent_excemption(&self, data: usize) -> u64 {
523        self.bank.get_minimum_balance_for_rent_exemption(data)
524    }
525
526    fn get_account(&self, pubkey: Pubkey) -> Option<Account> {
527        self.bank.get_account(&pubkey).map(|acc| acc.into())
528    }
529}
530
531pub struct LocalEnvironmentBuilder {
532    config: GenesisConfig,
533    faucet: Keypair,
534}
535
536impl LocalEnvironmentBuilder {
537    fn new() -> Self {
538        let faucet = random_keypair();
539        let mut config = GenesisConfig::new(
540            &[(
541                faucet.pubkey(),
542                AccountSharedData::new(1u64 << 48, 0, &system_program::id()),
543            )],
544            &[],
545        );
546        genesis_utils::activate_all_features(&mut config);
547
548        let mut builder = LocalEnvironmentBuilder { faucet, config };
549        builder.add_account_with_data(
550            spl_associated_token_account::ID,
551            bpf_loader::ID,
552            programs::SPL_ASSOCIATED_TOKEN,
553            true,
554        );
555        builder.add_account_with_data(
556            "Memo1UhkJRfHyvLMcVucJwxXeuD728EqVDDwQDxFMNo"
557                .parse()
558                .unwrap(),
559            bpf_loader::ID,
560            programs::SPL_MEMO1,
561            true,
562        );
563        builder.add_account_with_data(spl_memo::ID, bpf_loader::ID, programs::SPL_MEMO3, true);
564        builder.add_account_with_data(spl_token::ID, bpf_loader::ID, programs::SPL_TOKEN, true);
565        builder.add_account_with_lamports(rent::ID, sysvar::ID, 1);
566        builder
567    }
568
569    /// Sets the creation time of the network
570    pub fn set_creation_time(&mut self, unix_timestamp: UnixTimestamp) -> &mut Self {
571        self.config.creation_time = unix_timestamp as UnixTimestamp;
572        self
573    }
574
575    /// Adds the account into the environment.
576    pub fn add_account(&mut self, pubkey: Pubkey, account: Account) -> &mut Self {
577        self.config.add_account(pubkey, account.into());
578        self
579    }
580
581    /// Reads the program from the path and add it at the address into the environment.
582    pub fn add_program<P: AsRef<Path>>(&mut self, pubkey: Pubkey, path: P) -> &mut Self {
583        self.add_account_with_data(pubkey, bpf_loader::ID, &std::fs::read(path).unwrap(), true);
584        self
585    }
586
587    // Adds a rent-excempt account into the environment.
588    pub fn add_account_with_data(
589        &mut self,
590        pubkey: Pubkey,
591        owner: Pubkey,
592        data: &[u8],
593        executable: bool,
594    ) -> &mut Self {
595        self.add_account(
596            pubkey,
597            Account {
598                lamports: self.config.rent.minimum_balance(data.len()),
599                data: data.to_vec(),
600                executable,
601                owner,
602                rent_epoch: 0,
603            },
604        )
605    }
606
607    // Adds an account with the given balance into the environment.
608    pub fn add_account_with_lamports(
609        &mut self,
610        pubkey: Pubkey,
611        owner: Pubkey,
612        lamports: u64,
613    ) -> &mut Self {
614        self.add_account(
615            pubkey,
616            Account {
617                lamports,
618                data: vec![],
619                executable: false,
620                owner,
621                rent_epoch: 0,
622            },
623        )
624    }
625
626    // Adds a rent-excempt account into the environment.
627    pub fn add_account_with_packable<P: Pack>(
628        &mut self,
629        pubkey: Pubkey,
630        owner: Pubkey,
631        data: P,
632    ) -> &mut Self {
633        let data = {
634            let mut buf = vec![0u8; P::LEN];
635            data.pack_into_slice(&mut buf[..]);
636            buf
637        };
638        self.add_account_with_data(pubkey, owner, &data, false)
639    }
640
641    // Add a token-mint into the environment.
642    pub fn add_token_mint(
643        &mut self,
644        pubkey: Pubkey,
645        mint_authority: Option<Pubkey>,
646        supply: u64,
647        decimals: u8,
648        freeze_authority: Option<Pubkey>,
649    ) -> &mut Self {
650        self.add_account_with_packable(
651            pubkey,
652            spl_token::ID,
653            spl_token::state::Mint {
654                mint_authority: COption::from(mint_authority.map(|c| c.clone())),
655                supply,
656                decimals,
657                is_initialized: true,
658                freeze_authority: COption::from(freeze_authority.map(|c| c.clone())),
659            },
660        )
661    }
662
663    // Add a token-account into the environment.
664    pub fn add_account_with_tokens(
665        &mut self,
666        pubkey: Pubkey,
667        mint: Pubkey,
668        owner: Pubkey,
669        amount: u64,
670    ) -> &mut Self {
671        self.add_account_with_packable(
672            pubkey,
673            spl_token::ID,
674            spl_token::state::Account {
675                mint,
676                owner,
677                amount,
678                delegate: COption::None,
679                state: spl_token::state::AccountState::Initialized,
680                is_native: COption::None,
681                delegated_amount: 0,
682                close_authority: COption::None,
683            },
684        )
685    }
686
687    // Add the associated token-account into the environment.
688    pub fn add_associated_account_with_tokens(
689        &mut self,
690        owner: Pubkey,
691        mint: Pubkey,
692        amount: u64,
693    ) -> &mut Self {
694        self.add_account_with_packable(
695            get_associated_token_address(&owner, &mint),
696            spl_token::ID,
697            spl_token::state::Account {
698                mint,
699                owner,
700                amount,
701                delegate: COption::None,
702                state: spl_token::state::AccountState::Initialized,
703                is_native: COption::None,
704                delegated_amount: 0,
705                close_authority: COption::None,
706            },
707        )
708    }
709
710    /// Clone an account from a cluster using the given rpc client. Use [clone_upgradable_program_from_cluster] if you want to clone a upgradable program, as this requires multiple accounts.
711    pub fn clone_account_from_cluster(&mut self, pubkey: Pubkey, client: &RpcClient) -> &mut Self {
712        println!("Loading account {} from cluster", pubkey);
713        let account = client
714            .get_account(&pubkey)
715            .expect("couldn't retrieve account");
716        self.add_account(
717            pubkey,
718            Account {
719                lamports: account.lamports,
720                data: account.data,
721                executable: account.executable,
722                owner: account.owner,
723                rent_epoch: 0,
724            },
725        )
726    }
727
728    /// Clone multiple accounts from a cluster using the given rpc client.
729    pub fn clone_accounts_from_cluster(
730        &mut self,
731        pubkeys: &[Pubkey],
732        client: &RpcClient,
733    ) -> &mut Self {
734        for &pubkey in pubkeys {
735            self.clone_account_from_cluster(pubkey, client);
736        }
737        self
738    }
739
740    /// Clones all accounts required to execute the given executable program from the cluster, using the given rpc client.
741    pub fn clone_upgradable_program_from_cluster(
742        &mut self,
743        client: &RpcClient,
744        pubkey: Pubkey,
745    ) -> &mut Self {
746        println!("Loading upgradable program {} from cluster", pubkey);
747        let account = client
748            .get_account(&pubkey)
749            .expect("couldn't retrieve account");
750        let upgradable: UpgradeableLoaderState = account.deserialize_data().unwrap();
751        if let UpgradeableLoaderState::Program {
752            programdata_address,
753        } = upgradable
754        {
755            self.add_account(pubkey, account);
756            self.clone_account_from_cluster(programdata_address, client);
757        } else {
758            panic!("Account is not an upgradable program")
759        }
760        self
761    }
762
763    /// Finalizes the environment.
764    pub fn build(&mut self) -> LocalEnvironment {
765        let tmpdir = Path::new("/tmp/");
766
767        let bank = Bank::new_with_paths(
768            &self.config,
769            vec![tmpdir.to_path_buf()],
770            None,
771            Some(&Builtins {
772                genesis_builtins: [
773                    solana_bpf_loader_upgradeable_program!(),
774                    solana_bpf_loader_program!(),
775                    solana_bpf_loader_deprecated_program!(),
776                ]
777                .iter()
778                .map(|p| Builtin::new(&p.0, p.1, p.2))
779                .collect(),
780                feature_transitions: vec![],
781            }),
782            AccountSecondaryIndexes {
783                keys: None,
784                indexes: HashSet::new(),
785            },
786            false,
787            AccountShrinkThreshold::default(),
788            false,
789            None,
790            None,
791        );
792
793        let env = LocalEnvironment {
794            bank,
795            faucet: clone_keypair(&self.faucet),
796        };
797        env.advance_blockhash();
798
799        env
800    }
801}
802
803/// A remote environment on a cluster. Interacts with the cluster using RPC.
804pub struct RemoteEnvironment {
805    client: RpcClient,
806    payer: Keypair,
807}
808
809impl RemoteEnvironment {
810    /// Contruct a new remote environment. The payer keypair is expected to have enough funds to fund all transactions.
811    pub fn new(client: RpcClient, payer: Keypair) -> Self {
812        RemoteEnvironment { client, payer }
813    }
814
815    /// Construct a new remote environment, airdropping lamports from the given airdrop endpoint up to the given account. Use this on devnet and testnet.
816    pub fn new_with_airdrop(client: RpcClient, payer: Keypair, lamports: u64) -> Self {
817        let env = RemoteEnvironment { client, payer };
818        env.airdrop(env.payer().pubkey(), lamports);
819        env
820    }
821
822    /// Airdrop lamports up to the given balance to the account.
823    pub fn airdrop(&self, account: Pubkey, lamports: u64) {
824        if self.client.get_balance(&account).expect("get balance") < lamports {
825            println!("Requesting airdrop...");
826            let blockhash = self.client.get_latest_blockhash().unwrap();
827            let sig = self
828                .client
829                .request_airdrop_with_blockhash(&account, lamports, &blockhash)
830                .unwrap();
831            self.client
832                .confirm_transaction_with_spinner(&sig, &blockhash, CommitmentConfig::confirmed())
833                .unwrap();
834        }
835    }
836}
837
838impl Environment for RemoteEnvironment {
839    fn payer(&self) -> Keypair {
840        clone_keypair(&self.payer)
841    }
842
843    fn execute_transaction(
844        &mut self,
845        tx: Transaction,
846    ) -> EncodedConfirmedTransactionWithStatusMeta {
847        let sig = match self.client.send_and_confirm_transaction(&tx) {
848            Err(e) => panic!("{:#?}", e),
849            Ok(sig) => sig,
850        };
851        self.client
852            .get_transaction_with_config(
853                &sig,
854                RpcTransactionConfig {
855                    encoding: Some(UiTransactionEncoding::Binary),
856                    commitment: Some(CommitmentConfig::confirmed()),
857                    ..RpcTransactionConfig::default()
858                },
859            )
860            .unwrap()
861    }
862
863    fn get_latest_blockhash(&self) -> Hash {
864        self.client.get_latest_blockhash().unwrap()
865    }
866
867    fn get_rent_excemption(&self, data: usize) -> u64 {
868        self.client
869            .get_minimum_balance_for_rent_exemption(data)
870            .unwrap()
871    }
872
873    fn get_account(&self, pubkey: Pubkey) -> Option<Account> {
874        self.client
875            .get_account_with_commitment(&pubkey, self.client.commitment())
876            .unwrap()
877            .value
878    }
879}
880
881/// Utility trait for printing transaction results.
882pub trait PrintableTransaction {
883    /// Pretty print the transaction results, tagged with the given name for distinguishability.
884    fn print_named(&self, name: &str);
885
886    /// Pretty print the transaction results.
887    fn print(&self) {
888        self.print_named("");
889    }
890
891    /// Panic and print the transaction if it did not execute successfully
892    fn assert_success(&self);
893}
894
895impl PrintableTransaction for ConfirmedTransactionWithStatusMeta {
896    fn print_named(&self, name: &str) {
897        let tx = self.tx_with_meta.get_transaction();
898        let encoded = self
899            .clone()
900            .encode(UiTransactionEncoding::JsonParsed, None)
901            .expect("Failed to encode");
902        println!("EXECUTE {} (slot {})", name, encoded.slot);
903        println_transaction(&tx, encoded.transaction.meta.as_ref(), "  ", None, None);
904    }
905
906    fn assert_success(&self) {
907        match &self.tx_with_meta.get_status_meta() {
908            Some(meta) if meta.status.is_err() => {
909                self.print();
910                panic!("tx failed!")
911            }
912            _ => (),
913        }
914    }
915}
916
917impl PrintableTransaction for EncodedConfirmedTransactionWithStatusMeta {
918    fn print_named(&self, name: &str) {
919        let tx = self.transaction.transaction.decode().unwrap();
920        println!("EXECUTE {} (slot {})", name, self.slot);
921        println_transaction(&tx, self.transaction.meta.as_ref(), "  ", None, None);
922    }
923
924    fn assert_success(&self) {
925        match &self.transaction.meta {
926            Some(meta) if meta.err.is_some() => {
927                self.print();
928                panic!("tx failed!")
929            }
930            _ => (),
931        }
932    }
933}
934
935pub enum LogLevel {
936    TRACE,
937    DEBUG,
938    INFO,
939    WARN,
940    ERROR,
941}
942
943/// Setup solana logging. This is heavily recommended if you're using a local environment.
944pub fn setup_logging(level: LogLevel) {
945    match level {
946        LogLevel::TRACE => solana_logger::setup_with_default(
947            "trace,solana_runtime::message_processor=trace,solana_metrics::metrics=error",
948        ),
949        LogLevel::DEBUG => solana_logger::setup_with_default(
950            "debug,solana_runtime::message_processor=debug,solana_metrics::metrics=error",
951        ),
952        LogLevel::INFO => solana_logger::setup_with_default(
953            "info,solana_runtime::message_processor=info,solana_metrics::metrics=error",
954        ),
955        LogLevel::WARN => solana_logger::setup_with_default(
956            "warn,solana_runtime::message_processor=warn,solana_metrics::metrics=error",
957        ),
958        LogLevel::ERROR => solana_logger::setup_with_default(
959            "error,solana_runtime::message_processor=error,solana_metrics::metrics=error",
960        ),
961    }
962}
963
964/// Clone the given keypair.
965pub fn clone_keypair(keypair: &Keypair) -> Keypair {
966    Keypair::from_bytes(&keypair.to_bytes()).unwrap()
967}
968
969/// Generate a random keypair.
970pub fn random_keypair() -> Keypair {
971    Keypair::generate(&mut OsRng::default())
972}
973
974/// Return a recognisable Keypair. The public key will start with `Kxxx`, where xxx are the three digits of the number.
975/// `o` is used instead of `0`, as `0` is not part of the base58 charset.
976pub fn keypair(n: u8) -> Keypair {
977    Keypair::from_bytes(&keys::KEYPAIRS[n as usize]).unwrap()
978}
979
980/// Constructs a devnet client using `CommitmentConfig::confirmed()`.
981pub fn devnet_client() -> RpcClient {
982    RpcClient::new_with_commitment(
983        "https://api.devnet.solana.com/".to_string(),
984        CommitmentConfig::confirmed(),
985    )
986}
987
988/// Constructs a testnet client using `CommitmentConfig::confirmed()`.
989pub fn testnet_client() -> RpcClient {
990    RpcClient::new_with_commitment(
991        "https://api.testnet.solana.com/".to_string(),
992        CommitmentConfig::confirmed(),
993    )
994}
995
996/// Constructs a client connecting to localhost:8899 using `CommitmentConfig::confirmed()`.
997pub fn localhost_client() -> RpcClient {
998    RpcClient::new_with_commitment(
999        "http://localhost:8899/".to_string(),
1000        CommitmentConfig::confirmed(),
1001    )
1002}
poc_framework/lib.rs

poc_framework/
lib.rs
