anvil_zksync_core/node/
impersonate.rs

1use std::collections::HashSet;
2use std::sync::{Arc, RwLock};
3use zksync_types::Address;
4
5/// Manages impersonated accounts across the system.
6///
7/// Clones always agree on the set of impersonated accounts and updating one affects all other
8/// instances.
9#[derive(Clone, Debug, Default)]
10pub struct ImpersonationManager {
11    state: Arc<RwLock<ImpersonationState>>,
12}
13
14impl ImpersonationManager {
15    /// Sets the auto impersonation flag, when `true` it makes all accounts impersonated by default.
16    /// Setting to `false` disabled this behavior.
17    pub fn set_auto_impersonation(&self, enabled: bool) {
18        tracing::trace!(enabled, "auto impersonation status set");
19        self.state
20            .write()
21            .expect("ImpersonationManager lock is poisoned")
22            .auto = enabled
23    }
24
25    /// Starts impersonation for the provided account.
26    ///
27    /// Returns `true` if the account was not impersonated before.
28    pub fn impersonate(&self, addr: Address) -> bool {
29        tracing::trace!(?addr, "start impersonation");
30        let mut state = self
31            .state
32            .write()
33            .expect("ImpersonationManager lock is poisoned");
34        state.accounts.insert(addr)
35    }
36
37    /// Stops impersonation for the provided account.
38    ///
39    /// Returns `true` if the account was impersonated before.
40    pub fn stop_impersonating(&self, addr: &Address) -> bool {
41        tracing::trace!(?addr, "stop impersonation");
42        self.state
43            .write()
44            .expect("ImpersonationManager lock is poisoned")
45            .accounts
46            .remove(addr)
47    }
48
49    /// Returns whether the provided account is currently impersonated.
50    pub fn is_impersonating(&self, addr: &Address) -> bool {
51        let state = self
52            .state
53            .read()
54            .expect("ImpersonationManager lock is poisoned");
55        state.is_impersonating(addr)
56    }
57
58    /// Returns internal state representation.
59    pub fn state(&self) -> ImpersonationState {
60        self.state
61            .read()
62            .expect("ImpersonationManager lock is poisoned")
63            .clone()
64    }
65
66    /// Overrides current internal state with the provided value.
67    pub fn set_state(&self, state: ImpersonationState) {
68        *self
69            .state
70            .write()
71            .expect("ImpersonationManager lock is poisoned") = state;
72    }
73
74    /// Inspects the entire account set on a user-provided function without dropping the lock.
75    pub fn inspect<F, R>(&self, f: F) -> R
76    where
77        F: FnOnce(&ImpersonationState) -> R,
78    {
79        let guard = self
80            .state
81            .read()
82            .expect("ImpersonationManager lock is poisoned");
83        f(&guard)
84    }
85}
86
87#[derive(Clone, Debug, Default, PartialEq)]
88pub struct ImpersonationState {
89    /// If `true` then all accounts are impersonated regardless of `accounts` contents
90    pub auto: bool,
91    /// Accounts that are currently impersonated
92    pub accounts: HashSet<Address>,
93}
94
95impl ImpersonationState {
96    pub fn is_impersonating(&self, addr: &Address) -> bool {
97        self.auto || self.accounts.contains(addr)
98    }
99}
anvil_zksync_core/node/impersonate.rs

anvil_zksync_core/node/
impersonate.rs
