anvil_zksync_core/node/
impersonate.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
use std::collections::HashSet;
use std::sync::{Arc, RwLock};
use zksync_types::Address;

/// Manages impersonated accounts across the system.
///
/// Clones always agree on the set of impersonated accounts and updating one affects all other
/// instances.
#[derive(Clone, Debug, Default)]
pub struct ImpersonationManager {
    state: Arc<RwLock<ImpersonationState>>,
}

impl ImpersonationManager {
    /// Sets the auto impersonation flag, when `true` it makes all accounts impersonated by default.
    /// Setting to `false` disabled this behavior.
    pub fn set_auto_impersonation(&self, enabled: bool) {
        tracing::trace!(enabled, "auto impersonation status set");
        self.state
            .write()
            .expect("ImpersonationManager lock is poisoned")
            .auto = enabled
    }

    /// Starts impersonation for the provided account.
    ///
    /// Returns `true` if the account was not impersonated before.
    pub fn impersonate(&self, addr: Address) -> bool {
        tracing::trace!(?addr, "start impersonation");
        let mut state = self
            .state
            .write()
            .expect("ImpersonationManager lock is poisoned");
        state.accounts.insert(addr)
    }

    /// Stops impersonation for the provided account.
    ///
    /// Returns `true` if the account was impersonated before.
    pub fn stop_impersonating(&self, addr: &Address) -> bool {
        tracing::trace!(?addr, "stop impersonation");
        self.state
            .write()
            .expect("ImpersonationManager lock is poisoned")
            .accounts
            .remove(addr)
    }

    /// Returns whether the provided account is currently impersonated.
    pub fn is_impersonating(&self, addr: &Address) -> bool {
        let state = self
            .state
            .read()
            .expect("ImpersonationManager lock is poisoned");
        state.is_impersonating(addr)
    }

    /// Returns internal state representation.
    pub fn state(&self) -> ImpersonationState {
        self.state
            .read()
            .expect("ImpersonationManager lock is poisoned")
            .clone()
    }

    /// Overrides current internal state with the provided value.
    pub fn set_state(&self, state: ImpersonationState) {
        *self
            .state
            .write()
            .expect("ImpersonationManager lock is poisoned") = state;
    }

    /// Inspects the entire account set on a user-provided function without dropping the lock.
    pub fn inspect<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&ImpersonationState) -> R,
    {
        let guard = self
            .state
            .read()
            .expect("ImpersonationManager lock is poisoned");
        f(&guard)
    }
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct ImpersonationState {
    /// If `true` then all accounts are impersonated regardless of `accounts` contents
    pub auto: bool,
    /// Accounts that are currently impersonated
    pub accounts: HashSet<Address>,
}

impl ImpersonationState {
    pub fn is_impersonating(&self, addr: &Address) -> bool {
        self.auto || self.accounts.contains(addr)
    }
}