cfx_executor/stack/
stack_info.rs

1use cfx_parameters::internal_contract_addresses::ADMIN_CONTROL_CONTRACT_ADDRESS;
2use cfx_types::{AddressSpaceUtil, AddressWithSpace};
3use cfx_vm_types::Spec;
4use std::collections::HashMap;
5
6#[derive(Debug)]
7pub struct CallStackInfo {
8    call_stack_recipient_addresses: Vec<(AddressWithSpace, bool)>,
9    address_counter: HashMap<AddressWithSpace, u32>,
10    create_address_counter: HashMap<AddressWithSpace, u32>,
11    first_reentrancy_depth: Option<usize>,
12}
13
14impl CallStackInfo {
15    pub fn new() -> Self {
16        CallStackInfo {
17            call_stack_recipient_addresses: Vec::default(),
18            address_counter: HashMap::default(),
19            create_address_counter: HashMap::default(),
20            first_reentrancy_depth: None,
21        }
22    }
23
24    pub fn push(&mut self, address: AddressWithSpace, is_create: bool) {
25        // We should still use the correct behaviour to check if reentrancy
26        // happens.
27        if self.last() != Some(&address) && self.contains_key(&address) {
28            self.first_reentrancy_depth
29                .get_or_insert(self.call_stack_recipient_addresses.len());
30        }
31
32        self.call_stack_recipient_addresses
33            .push((address.clone(), is_create));
34        *self.address_counter.entry(address).or_insert(0) += 1;
35        if is_create {
36            *self.create_address_counter.entry(address).or_insert(0) += 1;
37        }
38    }
39
40    pub fn pop(&mut self) -> Option<(AddressWithSpace, bool)> {
41        let maybe_address = self.call_stack_recipient_addresses.pop();
42        if let Some((address, is_create)) = &maybe_address {
43            Self::decrease_counter(&mut self.address_counter, address);
44            if *is_create {
45                Self::decrease_counter(
46                    &mut self.create_address_counter,
47                    address,
48                );
49            }
50            if self.first_reentrancy_depth
51                == Some(self.call_stack_recipient_addresses.len())
52            {
53                self.first_reentrancy_depth = None
54            }
55        }
56        maybe_address
57    }
58
59    fn decrease_counter(
60        counter_map: &mut HashMap<AddressWithSpace, u32>,
61        address: &AddressWithSpace,
62    ) {
63        let poped_address_cnt = counter_map
64            .get_mut(address)
65            .expect("The lookup table should consistent with call stack");
66        *poped_address_cnt -= 1;
67        if *poped_address_cnt == 0 {
68            counter_map.remove(address);
69        }
70    }
71
72    pub fn last(&self) -> Option<&AddressWithSpace> {
73        self.call_stack_recipient_addresses
74            .last()
75            .map(|(address, _is_create)| address)
76    }
77
78    pub fn contains_key(&self, key: &AddressWithSpace) -> bool {
79        self.address_counter.contains_key(key)
80    }
81
82    pub fn in_reentrancy(&self, spec: &Spec) -> bool {
83        if spec.cip71 {
84            // After CIP-71, anti-reentrancy will closed.
85            false
86        } else {
87            // Consistent with old behaviour
88            // The old (unexpected) behaviour is equivalent to the top element
89            // is lost.
90            self.first_reentrancy_depth.map_or(false, |depth| {
91                (depth as isize)
92                    < self.call_stack_recipient_addresses.len() as isize - 1
93            })
94        }
95    }
96
97    pub fn admin_control_in_creation(&self) -> Option<&AddressWithSpace> {
98        if let [.., second_last, last] =
99            self.call_stack_recipient_addresses.as_slice()
100        {
101            if last.0 == ADMIN_CONTROL_CONTRACT_ADDRESS.with_native_space()
102                && second_last.1
103            {
104                Some(&second_last.0)
105            } else {
106                None
107            }
108        } else {
109            None
110        }
111    }
112
113    pub fn creating_contract(&self, address: &AddressWithSpace) -> bool {
114        self.create_address_counter
115            .get(address)
116            .map_or(false, |x| *x > 0)
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::CallStackInfo;
123    use cfx_types::{Address, AddressSpaceUtil, AddressWithSpace};
124
125    fn get_test_address_raw(n: u8) -> Address { Address::from([n; 20]) }
126
127    fn get_test_address(n: u8) -> AddressWithSpace {
128        get_test_address_raw(n).with_native_space()
129    }
130
131    #[test]
132    fn test_callstack_info() {
133        let mut call_stack = CallStackInfo::new();
134        call_stack.push(get_test_address(1), false);
135        call_stack.push(get_test_address(2), false);
136        assert_eq!(call_stack.pop(), Some((get_test_address(2), false)));
137        assert_eq!(call_stack.contains_key(&get_test_address(2)), false);
138
139        call_stack.push(get_test_address(3), true);
140        call_stack.push(get_test_address(4), false);
141        call_stack.push(get_test_address(3), false);
142        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(3));
143
144        assert_eq!(call_stack.pop(), Some((get_test_address(3), false)));
145        assert_eq!(call_stack.contains_key(&get_test_address(3)), true);
146        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(4));
147
148        assert_eq!(call_stack.pop(), Some((get_test_address(4), false)));
149        assert_eq!(call_stack.contains_key(&get_test_address(4)), false);
150        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(3));
151
152        assert_eq!(call_stack.pop(), Some((get_test_address(3), true)));
153        assert_eq!(call_stack.contains_key(&get_test_address(3)), false);
154        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(1));
155
156        call_stack.push(get_test_address(3), true);
157        call_stack.push(get_test_address(4), false);
158        call_stack.push(get_test_address(3), false);
159        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(3));
160
161        assert_eq!(call_stack.pop(), Some((get_test_address(3), false)));
162        assert_eq!(call_stack.contains_key(&get_test_address(3)), true);
163        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(4));
164
165        assert_eq!(call_stack.pop(), Some((get_test_address(4), false)));
166        assert_eq!(call_stack.contains_key(&get_test_address(4)), false);
167        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(3));
168
169        assert_eq!(call_stack.pop(), Some((get_test_address(3), true)));
170        assert_eq!(call_stack.contains_key(&get_test_address(3)), false);
171        assert_eq!(call_stack.last().unwrap().clone(), get_test_address(1));
172
173        assert_eq!(call_stack.pop(), Some((get_test_address(1), false)));
174        assert_eq!(call_stack.pop(), None);
175        assert_eq!(call_stack.last(), None);
176    }
177}