trident_fuzz/accounts_storage/
stake_store.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
use solana_sdk::{
    account::AccountSharedData, clock::Epoch, native_token::LAMPORTS_PER_SOL, pubkey::Pubkey,
    rent::Rent, signature::Keypair, signer::Signer, stake::stake_flags::StakeFlags,
};
use solana_stake_program::stake_state::{
    Authorized, Delegation, Lockup, Meta, Stake, StakeStateV2,
};

use crate::{fuzz_client::FuzzClient, AccountId};

use super::AccountsStorage;

pub struct StakeStore {
    pub pubkey: Pubkey,
}
impl From<Pubkey> for StakeStore {
    fn from(pubkey: Pubkey) -> Self {
        StakeStore { pubkey }
    }
}

impl AccountsStorage<StakeStore> {
    #[allow(clippy::too_many_arguments)]
    pub fn get_or_create_delegated_account(
        &mut self,
        account_id: AccountId,
        client: &mut impl FuzzClient,
        voter_pubkey: Pubkey,
        staker: Pubkey,
        withdrawer: Pubkey,
        stake: u64,
        activation_epoch: Epoch,
        deactivation_epoch: Option<Epoch>,
        lockup: Option<Lockup>,
    ) -> Pubkey {
        let key = self.accounts.entry(account_id).or_insert_with(|| {
            let stake_account = Keypair::new();

            let rent = Rent::default();
            let rent_exempt_lamports = rent.minimum_balance(StakeStateV2::size_of());
            let minimum_delegation = LAMPORTS_PER_SOL; // TODO: a way to get minimum delegation with feature set?
            let minimum_lamports = rent_exempt_lamports.saturating_add(minimum_delegation);

            let stake_state = StakeStateV2::Stake(
                Meta {
                    authorized: Authorized { staker, withdrawer },
                    lockup: lockup.unwrap_or_default(),
                    rent_exempt_reserve: rent_exempt_lamports,
                },
                Stake {
                    delegation: Delegation {
                        stake,
                        activation_epoch,
                        voter_pubkey,
                        deactivation_epoch: if let Some(epoch) = deactivation_epoch {
                            epoch
                        } else {
                            u64::MAX
                        },
                        ..Delegation::default()
                    },
                    ..Stake::default()
                },
                StakeFlags::default(),
            );
            let account = AccountSharedData::new_data_with_space(
                if stake > minimum_lamports {
                    stake
                } else {
                    minimum_lamports
                },
                &stake_state,
                StakeStateV2::size_of(),
                &solana_sdk::stake::program::ID,
            )
            .unwrap();

            client.set_account_custom(&stake_account.pubkey(), &account);

            StakeStore {
                pubkey: stake_account.pubkey(),
            }
        });
        key.pubkey
    }
    pub fn get_or_create_initialized_account(
        &mut self,
        account_id: AccountId,
        client: &mut impl FuzzClient,
        staker: Pubkey,
        withdrawer: Pubkey,
        lockup: Option<Lockup>,
    ) -> Pubkey {
        let key = self.accounts.entry(account_id).or_insert_with(|| {
            let stake_account = Keypair::new();

            let rent = Rent::default();
            let rent_exempt_lamports = rent.minimum_balance(StakeStateV2::size_of());

            let stake_state = StakeStateV2::Initialized(Meta {
                authorized: Authorized { staker, withdrawer },
                lockup: lockup.unwrap_or_default(),
                rent_exempt_reserve: rent_exempt_lamports,
            });
            let account = AccountSharedData::new_data_with_space(
                rent_exempt_lamports,
                &stake_state,
                StakeStateV2::size_of(),
                &solana_sdk::stake::program::ID,
            )
            .unwrap();
            client.set_account_custom(&stake_account.pubkey(), &account);

            StakeStore {
                pubkey: stake_account.pubkey(),
            }
        });
        key.pubkey
    }
    pub fn get(&self, account_id: AccountId) -> Pubkey {
        match self.accounts.get(&account_id) {
            Some(v) => v.pubkey,
            None => Pubkey::new_unique(),
        }
    }
}