trident_fuzz/
snapshot.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
#![allow(dead_code)] // The Snapshot is constructed in the FuzzTestExecutor macro and is generated automatically

use solana_sdk::account::{AccountSharedData, ReadableAccount};
use solana_sdk::clock::Epoch;
use solana_sdk::instruction::AccountMeta;
use solana_sdk::pubkey::Pubkey;

use crate::fuzz_client::FuzzClient;

use crate::error::*;

pub struct SnapshotAccount {
    address: Pubkey,
    account: AccountSharedData,
}

impl SnapshotAccount {
    pub fn get_account(&self) -> &AccountSharedData {
        &self.account
    }
    pub fn pubkey(&self) -> Pubkey {
        self.address
    }
    pub fn data(&self) -> &[u8] {
        self.account.data()
    }
    pub fn lamports(&self) -> u64 {
        self.account.lamports()
    }
    pub fn owner(&self) -> &Pubkey {
        self.account.owner()
    }
    pub fn executable(&self) -> bool {
        self.account.executable()
    }
    pub fn rent_epoch(&self) -> Epoch {
        self.account.rent_epoch()
    }
}
pub struct Snapshot {
    before: Vec<SnapshotAccount>,
    after: Vec<SnapshotAccount>,
    metas: Vec<AccountMeta>,
}

impl Snapshot {
    pub fn new(metas: &[AccountMeta]) -> Snapshot {
        Self {
            before: Default::default(),
            after: Default::default(),
            metas: metas.to_vec(),
        }
    }
    pub fn capture_before(
        &mut self,
        client: &mut impl FuzzClient,
    ) -> Result<(), FuzzClientErrorWithOrigin> {
        self.before = self
            .capture(client)
            .map_err(|e| e.with_context(Context::Pre))?;
        Ok(())
    }

    pub fn capture_after(
        &mut self,
        client: &mut impl FuzzClient,
    ) -> Result<(), FuzzClientErrorWithOrigin> {
        self.after = self
            .capture(client)
            .map_err(|e| e.with_context(Context::Post))?;
        Ok(())
    }

    fn capture(
        &mut self,
        client: &mut impl FuzzClient,
    ) -> Result<Vec<SnapshotAccount>, FuzzClientErrorWithOrigin> {
        let snapshot_accounts = self
            .metas
            .iter()
            .map(|meta| {
                let account = client.get_account(&meta.pubkey);
                SnapshotAccount {
                    address: meta.pubkey,
                    account,
                }
            })
            .collect();

        Ok(snapshot_accounts)
    }

    pub fn get_before(&self) -> &[SnapshotAccount] {
        &self.before
    }
    pub fn get_after(&self) -> &[SnapshotAccount] {
        &self.after
    }

    pub fn get_snapshot(&self) -> (&[SnapshotAccount], &[SnapshotAccount]) {
        (self.get_before(), self.get_after())
    }
}