cita_vm/state/
account.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use cita_trie::{PatriciaTrie, Trie, DB};
5use ethereum_types::{BigEndianHash, H256, U256};
6use hashbrown::HashMap;
7
8use crate::common;
9use crate::common::hash;
10use crate::state::err::Error;
11
12/// Single and pure account in the database. Usually, store it according to
13/// the following structure:
14/// Key: address -> Value: rlp::encode(account).
15#[derive(Debug)]
16pub struct Account {
17    pub balance: U256,
18    pub nonce: U256,
19    pub storage_root: H256,
20    pub code_hash: H256,
21    pub abi_hash: H256,
22}
23
24/// Free to use rlp::encode().
25impl rlp::Encodable for Account {
26    fn rlp_append(&self, s: &mut rlp::RlpStream) {
27        s.begin_list(5)
28            .append(&self.nonce)
29            .append(&self.balance)
30            .append(&self.storage_root)
31            .append(&self.code_hash)
32            .append(&self.abi_hash);
33    }
34}
35
36/// Free to use rlp::decode().
37impl rlp::Decodable for Account {
38    fn decode(data: &rlp::Rlp) -> Result<Self, rlp::DecoderError> {
39        Ok(Account {
40            nonce: data.val_at(0)?,
41            balance: data.val_at(1)?,
42            storage_root: data.val_at(2)?,
43            code_hash: data.val_at(3)?,
44            abi_hash: data.val_at(4)?,
45        })
46    }
47}
48
49#[derive(PartialEq, Eq, Clone, Copy, Debug)]
50pub enum CodeState {
51    Clean,
52    Dirty,
53}
54
55#[derive(Debug, Clone)]
56pub struct StateObject {
57    pub balance: U256,
58    pub nonce: U256,
59    pub storage_root: H256,
60    pub code_hash: H256,
61    pub code: Vec<u8>,
62    pub code_size: usize,
63    pub code_state: CodeState,
64    pub abi_hash: H256,
65    pub abi: Vec<u8>,
66    pub abi_size: usize,
67    pub abi_state: CodeState,
68    pub storage_changes: HashMap<H256, H256>,
69}
70
71impl From<Account> for StateObject {
72    fn from(account: Account) -> Self {
73        StateObject {
74            balance: account.balance,
75            nonce: account.nonce,
76            storage_root: account.storage_root,
77            code_hash: account.code_hash,
78            code: vec![],
79            code_size: 0,
80            code_state: CodeState::Clean,
81            abi_hash: account.abi_hash,
82            abi: vec![],
83            abi_size: 0,
84            abi_state: CodeState::Clean,
85            storage_changes: HashMap::new(),
86        }
87    }
88}
89
90//const CODE_PREFIX: &str = "C:";
91//const ABI_PREFIX: &str = "ABI:";
92
93impl StateObject {
94    /// Create a new account.
95    /// Note: If contract account generated, make sure you use `init_code` on
96    /// this before `commit`ing.
97    pub fn new(balance: U256, nonce: U256) -> StateObject {
98        StateObject {
99            balance,
100            nonce,
101            storage_root: common::hash::RLP_NULL,
102            code_hash: common::hash::NIL_DATA,
103            code: vec![],
104            code_size: 0,
105            code_state: CodeState::Clean,
106            abi_hash: common::hash::NIL_DATA,
107            abi: vec![],
108            abi_size: 0,
109            abi_state: CodeState::Clean,
110            storage_changes: HashMap::new(),
111        }
112    }
113
114    /// Create a new account from rlp bytes.
115    /// Note: make sure you use `read_code` after this.
116    pub fn from_rlp(data: &[u8]) -> Result<StateObject, Error> {
117        let account: Account = rlp::decode(data)?;
118        Ok(account.into())
119    }
120
121    /// Get the account from state object.
122    pub fn account(&self) -> Account {
123        Account {
124            balance: self.balance,
125            nonce: self.nonce,
126            storage_root: self.storage_root,
127            code_hash: self.code_hash,
128            abi_hash: self.abi_hash,
129        }
130    }
131
132    /// Get the rlp data.
133    pub fn rlp(&self) -> Vec<u8> {
134        rlp::encode(&self.account()).to_vec()
135    }
136
137    /// Function is_empty returns whether the given account is empty. Empty
138    /// is defined according to EIP161 (balance = nonce = code = 0).
139    pub fn is_empty(&self) -> bool {
140        self.balance.is_zero()
141            && self.nonce.is_zero()
142            && self.code_hash == common::hash::NIL_DATA
143            && self.storage_root == common::hash::RLP_NULL
144    }
145
146    /// Init the code by given data.
147    pub fn init_code(&mut self, code: Vec<u8>) {
148        self.code = code;
149        self.code_size = self.code.len();
150        self.code_hash = {
151            if self.code_size > 0 {
152                H256::from_slice(common::hash::summary(&self.code).as_slice())
153            } else {
154                common::hash::NIL_DATA
155            }
156        };
157        self.code_state = CodeState::Dirty;
158    }
159
160    /// Init the abi by given data.
161    pub fn init_abi(&mut self, abi: Vec<u8>) {
162        self.abi = abi;
163        self.abi_size = self.abi.len();
164        self.abi_hash = {
165            if self.abi_size > 0 {
166                H256::from_slice(common::hash::summary(&self.abi).as_slice())
167            } else {
168                common::hash::NIL_DATA
169            }
170        };
171        self.abi_state = CodeState::Dirty;
172    }
173
174    /// Read the code from database by it's codehash.
175    pub fn read_code<B: DB>(&mut self, db: Arc<B>) -> Result<(), Error> {
176        if self.code_hash == common::hash::NIL_DATA {
177            return Ok(());
178        }
179        let c = db
180            .get(self.code_hash.as_bytes())
181            .map_err(|e| Error::DB(format!("{}", e)))?
182            .unwrap_or_default();
183        self.code = c;
184        self.code_size = self.code.len();
185        self.code_state = CodeState::Clean;
186        Ok(())
187    }
188
189    /// Read the abi from database by it's abihash.
190    pub fn read_abi<B: DB>(&mut self, db: Arc<B>) -> Result<(), Error> {
191        if self.abi_hash == common::hash::NIL_DATA {
192            return Ok(());
193        }
194        let c = db
195            .get(self.abi_hash.as_bytes())
196            .map_err(|e| Error::DB(format!("{}", e)))?
197            .unwrap_or_default();
198        self.abi = c;
199        self.abi_size = self.abi.len();
200        self.abi_state = CodeState::Clean;
201        Ok(())
202    }
203
204    /// Add nonce by 1.
205    pub fn inc_nonce(&mut self) {
206        self.nonce += U256::from(1u8);
207    }
208
209    /// Add balance.
210    /// Note: overflowing is not allowed.
211    pub fn add_balance(&mut self, x: U256) {
212        let (a, b) = self.balance.overflowing_add(x);
213        // overflow is not allowed at state_object.
214        assert!(!b);
215        self.balance = a;
216    }
217
218    /// Sub balance.
219    /// Note: overflowing is not allowed.
220    pub fn sub_balance(&mut self, x: U256) {
221        self.balance = self.balance.saturating_sub(x);
222    }
223
224    /// Set (key, value) in storage cache.
225    pub fn set_storage(&mut self, key: H256, value: H256) {
226        self.storage_changes.insert(key, value);
227    }
228
229    /// Get value by key from database.
230    pub fn get_storage_at_backend<B: DB>(&self, db: Arc<B>, key: &H256) -> Result<Option<H256>, Error> {
231        if self.storage_root == common::hash::RLP_NULL {
232            return Ok(None);
233        }
234        let trie = PatriciaTrie::from(db, Arc::new(hash::get_hasher()), &self.storage_root.0)?;
235        if let Some(b) = trie.get(key.as_bytes())? {
236            let u256_k: U256 = rlp::decode(&b)?;
237            let h256_k = H256::from_uint(&u256_k);
238            return Ok(Some(h256_k));
239        }
240        Ok(None)
241    }
242
243    /// Get value by key from storage cache.
244    pub fn get_storage_at_changes(&self, key: &H256) -> Option<H256> {
245        self.storage_changes.get(key).copied()
246    }
247
248    /// Get the storage changes
249    pub fn get_storage_changes(&self) -> BTreeMap<String, String> {
250        let mut result = BTreeMap::new();
251        for (k, v) in self.storage_changes.iter() {
252            let key = String::from("0x") + &hex::encode(*k);
253            let value = String::from("0x") + &hex::encode(*v);
254            result.insert(key.clone(), value.clone());
255        }
256        result
257    }
258
259    /// Get value by key.
260    pub fn get_storage<B: DB>(&self, db: Arc<B>, key: &H256) -> Result<Option<H256>, Error> {
261        if let Some(value) = self.get_storage_at_changes(key) {
262            return Ok(Some(value));
263        }
264        if let Some(value) = self.get_storage_at_backend(db, key)? {
265            return Ok(Some(value));
266        }
267        Ok(None)
268    }
269
270    /// Get storage proof
271    pub fn get_storage_proof<B: DB>(&self, db: Arc<B>, key: &H256) -> Result<Vec<Vec<u8>>, Error> {
272        let trie = PatriciaTrie::from(db, Arc::new(hash::get_hasher()), &self.storage_root.0)
273            .map_err(|e| Error::DB(format!("StateObject::get_storage_proof: {}", e)))?;
274        let proof = trie.get_proof(&key.0)?;
275        Ok(proof)
276    }
277
278    /// Flush data in storage cache to database.
279    pub fn commit_storage<B: DB>(&mut self, db: Arc<B>) -> Result<(), Error> {
280        let mut trie = if self.storage_root == common::hash::RLP_NULL {
281            PatriciaTrie::new(db, Arc::new(hash::get_hasher()))
282        } else {
283            PatriciaTrie::from(db, Arc::new(hash::get_hasher()), &self.storage_root.0)?
284        };
285
286        for (k, v) in self.storage_changes.drain() {
287            if v.is_zero() {
288                trie.remove(k.as_bytes())?;
289            } else {
290                trie.insert(k.0.to_vec(), rlp::encode(&U256::from_big_endian(v.as_bytes())).to_vec())?;
291            }
292        }
293
294        self.storage_root = H256::from_slice(trie.root()?.as_slice());
295        Ok(())
296    }
297
298    /// Flush code to database if necessary.
299    pub fn commit_code<B: DB>(&mut self, db: Arc<B>) -> Result<(), Error> {
300        match (self.code_state == CodeState::Dirty, self.code.is_empty()) {
301            (true, true) => {
302                self.code_size = 0;
303                self.code_state = CodeState::Clean;
304            }
305            (true, false) => {
306                db.insert(self.code_hash.0.to_vec(), self.code.clone())
307                    .map_err(|e| Error::DB(format!("{}", e)))?;
308                self.code_size = self.code.len();
309                self.code_state = CodeState::Clean;
310            }
311            (false, _) => {}
312        }
313        Ok(())
314    }
315
316    /// Flush abi to database if necessary.
317    pub fn commit_abi<B: DB>(&mut self, db: Arc<B>) -> Result<(), Error> {
318        match (self.abi_state == CodeState::Dirty, self.abi.is_empty()) {
319            (true, true) => {
320                self.abi_size = 0;
321                self.abi_state = CodeState::Clean;
322            }
323            (true, false) => {
324                db.insert(self.abi_hash.0.to_vec(), self.abi.clone())
325                    .map_err(|e| Error::DB(format!("{}", e)))?;
326                self.abi_size = self.abi.len();
327                self.abi_state = CodeState::Clean;
328            }
329            (false, _) => {}
330        }
331        Ok(())
332    }
333
334    /// Clone without storage changes.
335    pub fn clone_clean(&self) -> StateObject {
336        StateObject {
337            balance: self.balance,
338            nonce: self.nonce,
339            storage_root: self.storage_root,
340            code: self.code.clone(),
341            code_hash: self.code_hash,
342            code_size: self.code_size,
343            code_state: self.code_state,
344            abi: self.abi.clone(),
345            abi_hash: self.abi_hash,
346            abi_size: self.abi_size,
347            abi_state: self.abi_state,
348            storage_changes: HashMap::new(),
349        }
350    }
351
352    /// Clone with storage changes.
353    pub fn clone_dirty(&self) -> StateObject {
354        let mut state_object = self.clone_clean();
355        state_object.storage_changes = self.storage_changes.clone();
356        state_object
357    }
358
359    /// Merge with others.
360    pub fn merge(&mut self, other: StateObject) {
361        self.balance = other.balance;
362        self.nonce = other.nonce;
363        self.storage_root = other.storage_root;
364        self.code_hash = other.code_hash;
365        self.code_state = other.code_state;
366        self.code = other.code;
367        self.code_size = other.code_size;
368        self.abi_hash = other.abi_hash;
369        self.abi_state = other.abi_state;
370        self.abi = other.abi;
371        self.abi_size = other.abi_size;
372        self.storage_changes = other.storage_changes;
373    }
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379    use std::str::FromStr;
380
381    #[test]
382    fn state_object_new() {
383        let o = StateObject::new(69u8.into(), 0u8.into());
384        assert_eq!(o.balance, 69u8.into());
385        assert_eq!(o.nonce, 0u8.into());
386        assert_eq!(o.code_hash, common::hash::NIL_DATA);
387        assert_eq!(o.abi_hash, common::hash::NIL_DATA);
388        assert_eq!(o.storage_root, common::hash::RLP_NULL);
389        assert_eq!(hex::encode(rlp::encode(&o.account())), "f8658045a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470");
390    }
391
392    #[test]
393    fn state_object_rlp() {
394        let a = StateObject::new(69u8.into(), 0u8.into());
395        let b = StateObject::from_rlp(&rlp::encode(&a.account())[..]).unwrap();
396        assert_eq!(a.balance, b.balance);
397        assert_eq!(a.nonce, b.nonce);
398        assert_eq!(a.code_hash, b.code_hash);
399        assert_eq!(a.storage_root, b.storage_root);
400    }
401
402    #[test]
403    fn state_object_code() {
404        let mut a = StateObject::new(69u8.into(), 0.into());
405        let db = Arc::new(cita_trie::MemoryDB::new(false));
406        a.init_code(vec![0x55, 0x44, 0xffu8]);
407        assert_eq!(a.code_state, CodeState::Dirty);
408        assert_eq!(a.code_size, 3);
409        a.commit_code(Arc::clone(&db)).unwrap();
410        assert_eq!(a.code_state, CodeState::Clean);
411        assert_eq!(
412            a.code_hash,
413            H256::from_str("af231e631776a517ca23125370d542873eca1fb4d613ed9b5d5335a46ae5b7eb").unwrap()
414        );
415
416        let k = a.code_hash.0.to_vec();
417        assert_eq!(db.get(&k).unwrap().unwrap(), vec![0x55, 0x44, 0xffu8]);
418        a.init_code(vec![0x55]);
419        assert_eq!(a.code_state, CodeState::Dirty);
420        assert_eq!(a.code_size, 1);
421        a.commit_code(Arc::clone(&db)).unwrap();
422        assert_eq!(
423            a.code_hash,
424            H256::from_str("37bf2238b11b68cdc8382cece82651b59d3c3988873b6e0f33d79694aa45f1be").unwrap()
425        );
426
427        let k = a.code_hash.0.to_vec();
428        assert_eq!(db.get(&k).unwrap().unwrap(), vec![0x55]);
429    }
430
431    #[test]
432    fn state_object_storage_1() {
433        let mut a = StateObject::new(69u8.into(), 0.into());
434        let db = Arc::new(cita_trie::MemoryDB::new(false));
435        a.set_storage(H256::zero(), H256::from_low_u64_be(0x1234));
436        a.commit_storage(Arc::clone(&db)).unwrap();
437        assert_eq!(
438            a.storage_root,
439            H256::from_str("71623f5ec821de33ad5aa81f8c82f0916c6f60de0a536f8c466d440c56715bd5").unwrap()
440        );
441    }
442
443    #[test]
444    fn state_object_storage_2() {
445        let mut a = StateObject::new(69u8.into(), 0.into());
446        let db = Arc::new(cita_trie::MemoryDB::new(false));
447        a.set_storage(H256::zero(), H256::from_low_u64_be(0x1234));
448        a.commit_storage(Arc::clone(&db)).unwrap();
449        assert_eq!(
450            a.storage_root,
451            //"c57e1afb758b07f8d2c8f13a3b6e44fa5ff94ab266facc5a4fd3f062426e50b2".into()
452            H256::from_str("71623f5ec821de33ad5aa81f8c82f0916c6f60de0a536f8c466d440c56715bd5").unwrap()
453        );
454        a.set_storage(H256::from_low_u64_be(1), H256::from_low_u64_be(0x1234));
455        a.commit_storage(Arc::clone(&db)).unwrap();
456        assert_eq!(
457            a.storage_root,
458            //"4e49574efd650366d071855e0a3975123ea9d64cc945e8f5de8c8c517e1b4ca5".into()
459            H256::from_str("a3db671bd0653a641fb031dccb869982da390eade9e6f993802ed09c4f6b7b2a").unwrap()
460        );
461        a.set_storage(H256::from_low_u64_be(1), H256::zero());
462        a.commit_storage(Arc::clone(&db)).unwrap();
463        assert_eq!(
464            a.storage_root,
465            //"c57e1afb758b07f8d2c8f13a3b6e44fa5ff94ab266facc5a4fd3f062426e50b2".into()
466            H256::from_str("71623f5ec821de33ad5aa81f8c82f0916c6f60de0a536f8c466d440c56715bd5").unwrap()
467        );
468    }
469
470    #[test]
471    fn state_object_storage_3() {
472        let mut a = StateObject::new(69u8.into(), 0.into());
473        let db = Arc::new(cita_trie::MemoryDB::new(false));
474        let a_rlp = {
475            a.set_storage(H256::zero(), H256::from_low_u64_be(0x1234));
476            a.commit_storage(Arc::clone(&db)).unwrap();
477            a.init_code(vec![]);
478            a.commit_code(Arc::clone(&db)).unwrap();
479            rlp::encode(&a.account())
480        };
481        a = StateObject::from_rlp(&a_rlp[..]).unwrap();
482        assert_eq!(
483            a.storage_root,
484            //"c57e1afb758b07f8d2c8f13a3b6e44fa5ff94ab266facc5a4fd3f062426e50b2".into()
485            H256::from_str("71623f5ec821de33ad5aa81f8c82f0916c6f60de0a536f8c466d440c56715bd5").unwrap()
486        );
487        assert_eq!(
488            a.get_storage(Arc::clone(&db), &H256::zero()).unwrap().unwrap(),
489            H256::from_low_u64_be(0x1234)
490        );
491        assert_eq!(a.get_storage(Arc::clone(&db), &H256::from_low_u64_be(1)).unwrap(), None);
492    }
493
494    #[test]
495    fn state_object_note_code() {
496        let mut a = StateObject::new(69u8.into(), 0.into());
497        let db = Arc::new(cita_trie::MemoryDB::new(false));
498        let a_rlp = {
499            a.init_code(vec![0x55, 0x44, 0xffu8]);
500            a.commit_code(Arc::clone(&db)).unwrap();
501            a.rlp()
502        };
503        a = StateObject::from_rlp(&a_rlp[..]).unwrap();
504        a.read_code(Arc::clone(&db)).unwrap();
505        assert_eq!(a.code, vec![0x55, 0x44, 0xffu8]);
506    }
507}
cita_vm/state/account.rs

cita_vm/state/
account.rs
