cita_state/
account.rs

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

cita_state/
account.rs
