darkfi/blockchain/
contract_store.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2026 Dyne.org foundation
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Affero General Public License as
7 * published by the Free Software Foundation, either version 3 of the
8 * License, or (at your option) any later version.
9 *
10r* This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 * GNU Affero General Public License for more details.
14 *
15 * You should have received a copy of the GNU Affero General Public License
16 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
17 */
18
19use std::{collections::BTreeMap, io::Cursor};
20
21use darkfi_sdk::{
22    crypto::contract_id::{
23        ContractId, NATIVE_CONTRACT_IDS_BYTES, NATIVE_CONTRACT_ZKAS_DB_NAMES,
24        SMART_CONTRACT_MONOTREE_DB_NAME, SMART_CONTRACT_ZKAS_DB_NAME,
25    },
26    monotree::{MemoryDb, Monotree, SledOverlayDb, SledTreeDb, EMPTY_HASH},
27};
28use darkfi_serial::{deserialize, serialize};
29use sled_overlay::{sled, SledDbOverlayStateDiff};
30use tracing::{debug, error};
31
32use crate::{
33    zk::{empty_witnesses, VerifyingKey, ZkCircuit},
34    zkas::ZkBinary,
35    Error, Result,
36};
37
38use super::{parse_record, SledDbOverlayPtr};
39
40pub const SLED_CONTRACTS_TREE: &[u8] = b"_contracts";
41pub const SLED_CONTRACTS_TREES_TREE: &[u8] = b"_contracts_trees";
42pub const SLED_BINCODE_TREE: &[u8] = b"_wasm_bincode";
43
44/// The `ContractStore` is a structure representing all `sled` trees related
45/// to storing the blockchain's contracts information.
46#[derive(Clone)]
47pub struct ContractStore {
48    /// The `sled` tree storing the wasm bincode for deployed contracts.
49    /// The layout looks like this:
50    /// ```plaintext
51    ///  tree: "_wasm_bincode"
52    ///   key: ContractId
53    /// value: Vec<u8>
54    pub wasm: sled::Tree,
55    /// The `sled` tree storing the pointers to contracts' databases.
56    /// See the rustdoc for the impl functions for more info.
57    /// The layout looks like this:
58    /// ```plaintext
59    ///  tree: "_contracts"
60    ///   key: ContractId
61    /// value: Vec<blake3(ContractId || tree_name)>
62    /// ```
63    /// These values get mutated with `init()` and `remove()`.
64    pub state: sled::Tree,
65    /// The `sled` tree storing the inverse pointers to contracts'
66    /// databases. See the rustdoc for the impl functions for more
67    /// info.
68    /// The layout looks like this:
69    /// ```plaintext
70    ///  tree: "_contracts_trees"
71    ///   key: blake3(ContractId || tree_name)
72    /// value: ContractId
73    /// ```
74    /// These values get mutated with `init()` and `remove()`.
75    pub state_trees: sled::Tree,
76}
77
78impl ContractStore {
79    /// Opens a new or existing `ContractStore` on the given sled database.
80    pub fn new(db: &sled::Db) -> Result<Self> {
81        let wasm = db.open_tree(SLED_BINCODE_TREE)?;
82        let state = db.open_tree(SLED_CONTRACTS_TREE)?;
83        let state_trees = db.open_tree(SLED_CONTRACTS_TREES_TREE)?;
84        Ok(Self { wasm, state, state_trees })
85    }
86
87    /// Fetches the bincode for a given ContractId from the store's wasm tree.
88    /// Returns an error if the bincode is not found.
89    pub fn get(&self, contract_id: ContractId) -> Result<Vec<u8>> {
90        if let Some(bincode) = self.wasm.get(serialize(&contract_id))? {
91            return Ok(bincode.to_vec())
92        }
93
94        Err(Error::WasmBincodeNotFound)
95    }
96
97    /// Do a lookup of an existing contract state. In order to succeed, the
98    /// state must have been previously initialized with `init()`. If the
99    /// state has been found, a handle to it will be returned. Otherwise, we
100    /// return an error.
101    pub fn lookup(
102        &self,
103        db: &sled::Db,
104        contract_id: &ContractId,
105        tree_name: &str,
106    ) -> Result<sled::Tree> {
107        debug!(target: "blockchain::contractstore::lookup", "Looking up state tree for {contract_id}:{tree_name}");
108
109        // A guard to make sure we went through init()
110        let contract_id_bytes = serialize(contract_id);
111        if !self.state.contains_key(&contract_id_bytes)? {
112            return Err(Error::ContractNotFound(contract_id.to_string()))
113        }
114
115        let state_pointers = self.state.get(&contract_id_bytes)?.unwrap();
116        let state_pointers: Vec<[u8; 32]> = deserialize(&state_pointers)?;
117
118        // We assume the tree has been created already, so it should be listed
119        // in this array. If not, that's an error.
120        let ptr = contract_id.hash_state_id(tree_name);
121        if !state_pointers.contains(&ptr) {
122            return Err(Error::ContractStateNotFound)
123        }
124
125        // We open the tree and return its handle
126        let tree = db.open_tree(ptr)?;
127        Ok(tree)
128    }
129
130    /// Attempt to remove an existing contract state. In order to succeed, the
131    /// state must have been previously initialized with `init()`. If the state
132    /// has been found, its contents in the tree will be cleared, and the pointer
133    /// will be removed from the main `ContractStateStore`. If anything is not
134    /// found as initialized, an error is returned.
135    /// NOTE: this function is not used right now, we keep it for future proofing,
136    ///       and its obviously untested.
137    pub fn remove(&self, db: &sled::Db, contract_id: &ContractId, tree_name: &str) -> Result<()> {
138        debug!(target: "blockchain::contractstore::remove", "Removing state tree for {contract_id}:{tree_name}");
139
140        // A guard to make sure we went through init()
141        let contract_id_bytes = serialize(contract_id);
142        if !self.state.contains_key(&contract_id_bytes)? {
143            return Err(Error::ContractNotFound(contract_id.to_string()))
144        }
145
146        let state_pointers = self.state.get(&contract_id_bytes)?.unwrap();
147        let mut state_pointers: Vec<[u8; 32]> = deserialize(&state_pointers)?;
148
149        // We assume the tree has been created already, so it should be listed
150        // in this array. If not, that's an error.
151        let ptr = contract_id.hash_state_id(tree_name);
152        if !state_pointers.contains(&ptr) {
153            return Err(Error::ContractStateNotFound)
154        }
155        if !self.state_trees.contains_key(ptr)? {
156            return Err(Error::ContractStateNotFound)
157        }
158
159        // Remove the deleted tree from the state pointer set.
160        state_pointers.retain(|x| *x != ptr);
161        self.state.insert(contract_id_bytes, serialize(&state_pointers))?;
162        self.state_trees.remove(ptr)?;
163
164        // Drop the deleted tree from the database
165        db.drop_tree(ptr)?;
166
167        Ok(())
168    }
169
170    /// Abstraction function for fetching a `ZkBinary` and its respective `VerifyingKey`
171    /// from a contract's zkas sled tree.
172    pub fn get_zkas(
173        &self,
174        db: &sled::Db,
175        contract_id: &ContractId,
176        zkas_ns: &str,
177    ) -> Result<(ZkBinary, VerifyingKey)> {
178        debug!(target: "blockchain::contractstore::get_zkas", "Looking up \"{contract_id}:{zkas_ns}\" zkas circuit & vk");
179
180        let zkas_tree = self.lookup(db, contract_id, SMART_CONTRACT_ZKAS_DB_NAME)?;
181
182        let Some(zkas_bytes) = zkas_tree.get(serialize(&zkas_ns))? else {
183            return Err(Error::ZkasBincodeNotFound)
184        };
185
186        // If anything in this function panics, that means corrupted data managed
187        // to get into this sled tree. This should not be possible.
188        let (zkbin, vkbin): (Vec<u8>, Vec<u8>) = deserialize(&zkas_bytes).unwrap();
189
190        // The first vec is the compiled zkas binary
191        let zkbin = ZkBinary::decode(&zkbin, false).unwrap();
192
193        // Construct the circuit to be able to read the VerifyingKey
194        let circuit = ZkCircuit::new(empty_witnesses(&zkbin).unwrap(), &zkbin);
195
196        // The second one is the serialized VerifyingKey for it
197        let mut vk_buf = Cursor::new(vkbin);
198        let vk = VerifyingKey::read::<Cursor<Vec<u8>>, ZkCircuit>(&mut vk_buf, circuit).unwrap();
199
200        Ok((zkbin, vk))
201    }
202
203    /// Retrieve all wasm bincodes from the store's wasm tree in the form
204    /// of a tuple (`contract_id`, `bincode`).
205    /// Be careful as this will try to load everything in memory.
206    pub fn get_all_wasm(&self) -> Result<Vec<(ContractId, Vec<u8>)>> {
207        let mut bincodes = vec![];
208
209        for bincode in self.wasm.iter() {
210            let bincode = bincode.unwrap();
211            let contract_id = deserialize(&bincode.0)?;
212            bincodes.push((contract_id, bincode.1.to_vec()));
213        }
214
215        Ok(bincodes)
216    }
217
218    /// Retrieve all contract states from the store's state tree in the
219    /// form of a tuple (`contract_id`, `state_hashes`).
220    /// Be careful as this will try to load everything in memory.
221    pub fn get_all_states(&self) -> Result<Vec<(ContractId, Vec<blake3::Hash>)>> {
222        let mut contracts = vec![];
223
224        for contract in self.state.iter() {
225            contracts.push(parse_record(contract.unwrap())?);
226        }
227
228        Ok(contracts)
229    }
230
231    /// Retrieve provided key value bytes from a contract's zkas sled tree.
232    pub fn get_state_tree_value(
233        &self,
234        db: &sled::Db,
235        contract_id: &ContractId,
236        tree_name: &str,
237        key: &[u8],
238    ) -> Result<Vec<u8>> {
239        debug!(target: "blockchain::contractstore::get_state_tree_value", "Looking up state tree value for {contract_id}:{tree_name}");
240
241        // Grab the state tree
242        let state_tree = self.lookup(db, contract_id, tree_name)?;
243
244        // Grab the key value
245        match state_tree.get(key)? {
246            Some(value) => Ok(value.to_vec()),
247            None => Err(Error::DatabaseError(format!(
248                "State tree {contract_id}:{tree_name} doesn't contain key: {key:?}"
249            ))),
250        }
251    }
252
253    /// Retrieve all records from a contract's zkas sled tree, as a `BTreeMap`.
254    /// Be careful as this will try to load everything in memory.
255    pub fn get_state_tree_records(
256        &self,
257        db: &sled::Db,
258        contract_id: &ContractId,
259        tree_name: &str,
260    ) -> Result<BTreeMap<Vec<u8>, Vec<u8>>> {
261        debug!(target: "blockchain::contractstore::get_state_tree_records", "Looking up state tree records for {contract_id}:{tree_name}");
262
263        // Grab the state tree
264        let state_tree = self.lookup(db, contract_id, tree_name)?;
265
266        // Retrieve its records
267        let mut ret = BTreeMap::new();
268        for record in state_tree.iter() {
269            let (key, value) = record.unwrap();
270            ret.insert(key.to_vec(), value.to_vec());
271        }
272
273        Ok(ret)
274    }
275
276    /// Generate a Monotree(SMT) containing all contracts states
277    /// roots, along with the wasm bincodes monotree root.
278    ///
279    /// Note: native contracts zkas tree and wasm bincodes are excluded.
280    pub fn get_state_monotree(&self, db: &sled::Db) -> Result<Monotree<MemoryDb>> {
281        // Initialize the monotree
282        debug!(target: "blockchain::contractstore::get_state_monotree", "Initializing global monotree...");
283        let mut root = None;
284        let monotree_db = MemoryDb::new();
285        let mut tree = Monotree::new(monotree_db);
286
287        // Iterate over current contracts states records
288        for state_record in self.state.iter() {
289            // Grab its monotree pointer
290            let (contract_id, state_pointers): (ContractId, Vec<[u8; 32]>) =
291                parse_record(state_record?)?;
292            let state_monotree_ptr = contract_id.hash_state_id(SMART_CONTRACT_MONOTREE_DB_NAME);
293
294            // Check it exists
295            if !state_pointers.contains(&state_monotree_ptr) {
296                return Err(Error::ContractStateNotFound)
297            }
298            if !self.state_trees.contains_key(state_monotree_ptr)? {
299                return Err(Error::ContractStateNotFound)
300            }
301
302            // Grab its monotree
303            let state_tree = db.open_tree(state_monotree_ptr)?;
304            let state_monotree_db = SledTreeDb::new(&state_tree);
305            let state_monotree = Monotree::new(state_monotree_db);
306
307            // Insert its root to the global monotree
308            let state_monotree_root = match state_monotree.get_headroot()? {
309                Some(hash) => hash,
310                None => *EMPTY_HASH,
311            };
312            debug!(target: "blockchain::contractstore::get_state_monotree", "Contract {contract_id} root: {}", blake3::Hash::from(state_monotree_root));
313            root = tree.insert(root.as_ref(), &contract_id.to_bytes(), &state_monotree_root)?;
314            debug!(target: "blockchain::contractstore::get_state_monotree", "New global root: {}", blake3::Hash::from(root.unwrap()));
315        }
316
317        // Iterate over current contracts wasm bincodes to compute its monotree root
318        debug!(target: "blockchain::contractstore::get_state_monotree", "Initializing wasm bincodes monotree...");
319        let mut wasm_monotree_root = None;
320        let wasm_monotree_db = MemoryDb::new();
321        let mut wasm_monotree = Monotree::new(wasm_monotree_db);
322        for record in self.wasm.iter() {
323            let (key, value) = record?;
324
325            // Skip native ones
326            if NATIVE_CONTRACT_IDS_BYTES.contains(&deserialize(&key)?) {
327                continue
328            }
329
330            // Insert record
331            let key = blake3::hash(&key);
332            let value = blake3::hash(&value);
333            debug!(target: "blockchain::contractstore::get_state_monotree", "Inserting key {key} with value: {value}");
334            wasm_monotree_root = wasm_monotree.insert(
335                wasm_monotree_root.as_ref(),
336                key.as_bytes(),
337                value.as_bytes(),
338            )?;
339        }
340
341        // Insert wasm bincodes root to the global monotree
342        let wasm_monotree_root = match wasm_monotree_root {
343            Some(hash) => hash,
344            None => *EMPTY_HASH,
345        };
346        debug!(target: "blockchain::contractstore::get_state_monotree", "New root: {}", blake3::Hash::from(wasm_monotree_root));
347        root = tree.insert(
348            root.as_ref(),
349            blake3::hash(SLED_BINCODE_TREE).as_bytes(),
350            &wasm_monotree_root,
351        )?;
352        debug!(target: "blockchain::contractstore::get_state_monotree", "New global root: {}", blake3::Hash::from(root.unwrap()));
353        tree.set_headroot(root.as_ref());
354
355        Ok(tree)
356    }
357}
358
359/// Overlay structure over a [`ContractStore`] instance.
360pub struct ContractStoreOverlay(SledDbOverlayPtr);
361
362impl ContractStoreOverlay {
363    pub fn new(overlay: &SledDbOverlayPtr) -> Result<Self> {
364        overlay.lock().unwrap().open_tree(SLED_BINCODE_TREE, true)?;
365        overlay.lock().unwrap().open_tree(SLED_CONTRACTS_TREE, true)?;
366        overlay.lock().unwrap().open_tree(SLED_CONTRACTS_TREES_TREE, true)?;
367        Ok(Self(overlay.clone()))
368    }
369
370    /// Fetches the bincode for a given ContractId from the overlay's wasm tree.
371    /// Returns an error if the bincode is not found.
372    pub fn get(&self, contract_id: ContractId) -> Result<Vec<u8>> {
373        if let Some(bincode) =
374            self.0.lock().unwrap().get(SLED_BINCODE_TREE, &serialize(&contract_id))?
375        {
376            return Ok(bincode.to_vec())
377        }
378
379        Err(Error::WasmBincodeNotFound)
380    }
381
382    /// Inserts or replaces the bincode for a given ContractId into the overlay's
383    /// wasm tree.
384    pub fn insert(&self, contract_id: ContractId, bincode: &[u8]) -> Result<()> {
385        if let Err(e) =
386            self.0.lock().unwrap().insert(SLED_BINCODE_TREE, &serialize(&contract_id), bincode)
387        {
388            error!(target: "blockchain::contractstoreoverlay::insert", "Failed to insert bincode to Wasm tree: {e}");
389            return Err(e.into())
390        }
391
392        Ok(())
393    }
394
395    /// Try to initialize a new contract state. Contracts can create a number
396    /// of trees, separated by `tree_name`, which they can then use from the
397    /// smart contract API. `init()` will look into the main `ContractStateStoreOverlay`
398    /// tree to check if the smart contract was already deployed, and if so
399    /// it will fetch a vector of these states that were initialized. If the
400    /// state was already found, this function will return an error, because
401    /// in this case the handle should be fetched using `lookup()`.
402    /// If the tree was not initialized previously, it will be appended to
403    /// the main `ContractStateStoreOverlay` tree and a handle to it will be
404    /// returned.
405    pub fn init(&self, contract_id: &ContractId, tree_name: &str) -> Result<[u8; 32]> {
406        debug!(target: "blockchain::contractstoreoverlay::init", "Initializing state overlay tree for {contract_id}:{tree_name}");
407        let mut lock = self.0.lock().unwrap();
408
409        // See if there are existing state trees.
410        // If not, just start with an empty vector.
411        let contract_id_bytes = serialize(contract_id);
412        let mut state_pointers: Vec<[u8; 32]> =
413            if lock.contains_key(SLED_CONTRACTS_TREE, &contract_id_bytes)? {
414                let bytes = lock.get(SLED_CONTRACTS_TREE, &contract_id_bytes)?.unwrap();
415                deserialize(&bytes)?
416            } else {
417                vec![]
418            };
419
420        // If the db was never initialized, it should not be in here.
421        let ptr = contract_id.hash_state_id(tree_name);
422        if state_pointers.contains(&ptr) {
423            return Err(Error::ContractAlreadyInitialized)
424        }
425
426        // Now we add it so it's marked as initialized and create its tree.
427        state_pointers.push(ptr);
428        lock.insert(SLED_CONTRACTS_TREE, &contract_id_bytes, &serialize(&state_pointers))?;
429        lock.insert(SLED_CONTRACTS_TREES_TREE, &ptr, &contract_id_bytes)?;
430        lock.open_tree(&ptr, false)?;
431
432        Ok(ptr)
433    }
434
435    /// Do a lookup of an existing contract state. In order to succeed, the
436    /// state must have been previously initialized with `init()`. If the
437    /// state has been found, a handle to it will be returned. Otherwise, we
438    /// return an error.
439    pub fn lookup(&self, contract_id: &ContractId, tree_name: &str) -> Result<[u8; 32]> {
440        debug!(target: "blockchain::contractstoreoverlay::lookup", "Looking up state tree for {contract_id}:{tree_name}");
441        let mut lock = self.0.lock().unwrap();
442
443        // A guard to make sure we went through init()
444        let contract_id_bytes = serialize(contract_id);
445        if !lock.contains_key(SLED_CONTRACTS_TREE, &contract_id_bytes)? {
446            return Err(Error::ContractNotFound(contract_id.to_string()))
447        }
448
449        let state_pointers = lock.get(SLED_CONTRACTS_TREE, &contract_id_bytes)?.unwrap();
450        let state_pointers: Vec<[u8; 32]> = deserialize(&state_pointers)?;
451
452        // We assume the tree has been created already, so it should be listed
453        // in this array. If not, that's an error.
454        let ptr = contract_id.hash_state_id(tree_name);
455        if !state_pointers.contains(&ptr) {
456            return Err(Error::ContractStateNotFound)
457        }
458        if !lock.contains_key(SLED_CONTRACTS_TREES_TREE, &ptr)? {
459            return Err(Error::ContractStateNotFound)
460        }
461
462        // We open the tree and return its handle
463        lock.open_tree(&ptr, false)?;
464        Ok(ptr)
465    }
466
467    /// Abstraction function for fetching a `ZkBinary` and its respective `VerifyingKey`
468    /// from a contract's zkas sled tree.
469    pub fn get_zkas(
470        &self,
471        contract_id: &ContractId,
472        zkas_ns: &str,
473    ) -> Result<(ZkBinary, VerifyingKey)> {
474        debug!(target: "blockchain::contractstoreoverlay::get_zkas", "Looking up \"{contract_id}:{zkas_ns}\" zkas circuit & vk");
475
476        let zkas_tree = self.lookup(contract_id, SMART_CONTRACT_ZKAS_DB_NAME)?;
477
478        let Some(zkas_bytes) = self.0.lock().unwrap().get(&zkas_tree, &serialize(&zkas_ns))? else {
479            return Err(Error::ZkasBincodeNotFound)
480        };
481
482        // If anything in this function panics, that means corrupted data managed
483        // to get into this sled tree. This should not be possible.
484        let (zkbin, vkbin): (Vec<u8>, Vec<u8>) = deserialize(&zkas_bytes).unwrap();
485
486        // The first vec is the compiled zkas binary
487        let zkbin = ZkBinary::decode(&zkbin, false).unwrap();
488
489        // Construct the circuit to be able to read the VerifyingKey
490        let circuit = ZkCircuit::new(empty_witnesses(&zkbin).unwrap(), &zkbin);
491
492        // The second one is the serialized VerifyingKey for it
493        let mut vk_buf = Cursor::new(vkbin);
494        let vk = VerifyingKey::read::<Cursor<Vec<u8>>, ZkCircuit>(&mut vk_buf, circuit).unwrap();
495
496        Ok((zkbin, vk))
497    }
498
499    /// Generate a Monotree(SMT) containing all contracts states
500    /// roots, along with the wasm bincodes monotree roots.
501    /// Be carefull as this will open all states monotrees in the
502    /// overlay, and all contract state trees if their monotrees
503    /// need rebuild.
504    ///
505    /// Note: native contracts zkas tree and wasm bincodes are
506    /// excluded.
507    pub fn get_state_monotree(&self) -> Result<Monotree<MemoryDb>> {
508        let mut lock = self.0.lock().unwrap();
509
510        // Grab all states pointers
511        debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Retrieving state pointers...");
512        let mut states_monotrees_pointers = vec![];
513        for state_record in lock.iter(SLED_CONTRACTS_TREE)? {
514            // Grab its monotree pointer
515            let (contract_id, mut state_pointers): (ContractId, Vec<[u8; 32]>) =
516                parse_record(state_record?)?;
517            let state_monotree_ptr = contract_id.hash_state_id(SMART_CONTRACT_MONOTREE_DB_NAME);
518
519            // Check it exists
520            if !state_pointers.contains(&state_monotree_ptr) {
521                return Err(Error::ContractStateNotFound)
522            }
523            if !lock.contains_key(SLED_CONTRACTS_TREES_TREE, &state_monotree_ptr)? {
524                return Err(Error::ContractStateNotFound)
525            }
526
527            // Skip native zkas trees
528            if NATIVE_CONTRACT_IDS_BYTES.contains(&contract_id.to_bytes()) {
529                state_pointers.retain(|ptr| !NATIVE_CONTRACT_ZKAS_DB_NAMES.contains(ptr));
530            }
531
532            states_monotrees_pointers.push((contract_id, state_pointers, state_monotree_ptr));
533        }
534
535        // Initialize the monotree
536        debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Initializing global monotree...");
537        let mut root = None;
538        let monotree_db = MemoryDb::new();
539        let mut tree = Monotree::new(monotree_db);
540
541        // Iterate over contract states monotrees pointers
542        for (contract_id, state_pointers, state_monotree_ptr) in states_monotrees_pointers {
543            // Iterate over contract state pointers to find their
544            // inserted keys. If any of them has dropped keys, we must
545            // rebuild the contract state monotree.
546            debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Updating monotree for contract: {contract_id}");
547            let mut rebuild = false;
548            let mut inserts = vec![];
549            'outer: for state_ptr in &state_pointers {
550                // Skip the actual monotree state pointer
551                if state_ptr == &state_monotree_ptr {
552                    continue
553                }
554
555                // Look for it in the overlay
556                for (state_key, state_cache) in &lock.state.caches {
557                    if state_key != state_ptr {
558                        continue
559                    }
560
561                    // Check if it has dropped keys
562                    if !state_cache.state.removed.is_empty() {
563                        rebuild = true;
564                        break 'outer
565                    }
566
567                    // Grab the new/updated keys
568                    for (key, value) in &state_cache.state.cache {
569                        let key = blake3::hash(key);
570                        let value = blake3::hash(value);
571                        inserts.push((key, value))
572                    }
573                    break
574                }
575            }
576
577            // Check if we need to rebuild it
578            if rebuild {
579                // Iterate over all contract states to grab the monotree keys
580                debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Rebuilding monotree...");
581                inserts = vec![];
582                for state_ptr in state_pointers {
583                    // Open the contract state
584                    lock.open_tree(&state_ptr, false)?;
585
586                    // If the pointer is the monotree one, clear it
587                    if state_ptr == state_monotree_ptr {
588                        lock.clear(&state_ptr)?;
589                        continue
590                    }
591
592                    // Grab all its keys
593                    for record in lock.iter(&state_ptr)? {
594                        let (key, value) = record?;
595                        let key = blake3::hash(&key);
596                        let value = blake3::hash(&value);
597                        inserts.push((key, value))
598                    }
599                }
600            }
601
602            // Grab its monotree
603            let state_monotree_db = SledOverlayDb::new(&mut lock, &state_monotree_ptr)?;
604            let mut state_monotree = Monotree::new(state_monotree_db);
605            let mut state_monotree_root =
606                if rebuild { None } else { state_monotree.get_headroot()? };
607            let state_monotree_root_str = match state_monotree_root {
608                Some(hash) => blake3::Hash::from(hash),
609                None => blake3::Hash::from(*EMPTY_HASH),
610            };
611            debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Current root: {state_monotree_root_str}");
612
613            // Update or insert new records
614            for (key, value) in &inserts {
615                debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Inserting key {key} with value: {value}");
616                state_monotree_root = state_monotree.insert(
617                    state_monotree_root.as_ref(),
618                    key.as_bytes(),
619                    value.as_bytes(),
620                )?;
621            }
622
623            // Set root
624            state_monotree.set_headroot(state_monotree_root.as_ref());
625
626            // Insert its root to the global monotree
627            let state_monotree_root = match state_monotree_root {
628                Some(hash) => hash,
629                None => *EMPTY_HASH,
630            };
631            debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "New root: {}", blake3::Hash::from(state_monotree_root));
632            root = tree.insert(root.as_ref(), &contract_id.to_bytes(), &state_monotree_root)?;
633        }
634
635        // Iterate over current contracts wasm bincodes to compute its monotree root
636        debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Initializing wasm bincodes monotree...");
637        let mut wasm_monotree_root = None;
638        let wasm_monotree_db = MemoryDb::new();
639        let mut wasm_monotree = Monotree::new(wasm_monotree_db);
640        for record in lock.iter(SLED_BINCODE_TREE)? {
641            let (key, value) = record?;
642
643            // Skip native ones
644            if NATIVE_CONTRACT_IDS_BYTES.contains(&deserialize(&key)?) {
645                continue
646            }
647
648            // Insert record
649            let key = blake3::hash(&key);
650            let value = blake3::hash(&value);
651            debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "Inserting key {key} with value: {value}");
652            wasm_monotree_root = wasm_monotree.insert(
653                wasm_monotree_root.as_ref(),
654                key.as_bytes(),
655                value.as_bytes(),
656            )?;
657        }
658
659        // Insert wasm bincodes root to the global monotree
660        let wasm_monotree_root = match wasm_monotree_root {
661            Some(hash) => hash,
662            None => *EMPTY_HASH,
663        };
664        debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "New root: {}", blake3::Hash::from(wasm_monotree_root));
665        root = tree.insert(
666            root.as_ref(),
667            blake3::hash(SLED_BINCODE_TREE).as_bytes(),
668            &wasm_monotree_root,
669        )?;
670        debug!(target: "blockchain::contractstoreoverlay::get_state_monotree", "New global root: {}", blake3::Hash::from(root.unwrap()));
671        tree.set_headroot(root.as_ref());
672
673        Ok(tree)
674    }
675
676    /// Retrieve all updated contracts states and wasm bincodes from
677    /// provided overlay diff, update their monotrees in the overlay
678    /// and their records in the provided Monotree(SMT).
679    ///
680    /// Note: native contracts zkas tree and wasm bincodes are
681    /// excluded.
682    pub fn update_state_monotree(
683        &self,
684        diff: &SledDbOverlayStateDiff,
685        tree: &mut Monotree<MemoryDb>,
686    ) -> Result<()> {
687        // If a contract was dropped, we must rebuild the monotree from
688        // scratch.
689        if let Some((state_cache, _)) = diff.caches.get(SLED_CONTRACTS_TREE) {
690            if !state_cache.removed.is_empty() {
691                debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Rebuilding global monotree...");
692                *tree = self.get_state_monotree()?;
693                return Ok(());
694            }
695        }
696
697        // Grab lock over the overlay
698        let mut lock = self.0.lock().unwrap();
699        debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Retrieving contracts updates...");
700
701        // If a contract tree was dropped, we must rebuild its monotree
702        // from scratch.
703        let mut contracts_updates = BTreeMap::new();
704        if let Some((state_cache, _)) = diff.caches.get(SLED_CONTRACTS_TREES_TREE) {
705            // Mark all the contracts of dropped trees for rebuild
706            for contract_id_bytes in state_cache.removed.values() {
707                contracts_updates.insert(contract_id_bytes.clone(), (true, vec![]));
708            }
709        }
710
711        // Iterate over diff caches to find all contracts updates
712        for (state_key, state_cache) in &diff.caches {
713            // Check if that cache is not a contract state one.
714            // Overlay protected trees are all the native/non-contract
715            // ones.
716            if lock.state.protected_tree_names.contains(state_key) {
717                continue
718            }
719
720            // Grab the actual state key
721            let state_key = deserialize(state_key)?;
722
723            // Skip native zkas tree
724            if NATIVE_CONTRACT_ZKAS_DB_NAMES.contains(&state_key) {
725                continue
726            }
727
728            // Grab its contract id
729            let Some(contract_id_bytes) = lock.get(SLED_CONTRACTS_TREES_TREE, &state_key)? else {
730                return Err(Error::ContractStateNotFound)
731            };
732            let contract_id: ContractId = deserialize(&contract_id_bytes)?;
733
734            // Skip the actual monotree state cache
735            let state_monotree_ptr = contract_id.hash_state_id(SMART_CONTRACT_MONOTREE_DB_NAME);
736            if state_monotree_ptr == state_key {
737                continue
738            }
739
740            // Grab its record from the map
741            let (rebuild, mut inserts) = match contracts_updates.get(&contract_id_bytes) {
742                Some(r) => r.clone(),
743                None => (false, vec![]),
744            };
745
746            // Check if the contract monotree is already marked for
747            // rebuild.
748            if rebuild {
749                continue
750            }
751
752            // If records have been dropped, mark the contract monotree
753            // for rebuild.
754            if !state_cache.0.removed.is_empty() {
755                contracts_updates.insert(contract_id_bytes, (true, vec![]));
756                continue
757            }
758
759            // Grab the new/updated keys
760            for (key, (_, value)) in &state_cache.0.cache {
761                let key = blake3::hash(key);
762                let value = blake3::hash(value);
763                inserts.push((key, value))
764            }
765            contracts_updates.insert(contract_id_bytes, (rebuild, inserts));
766        }
767
768        // Grab current root
769        let mut root = tree.get_headroot()?;
770        let root_str = match root {
771            Some(hash) => blake3::Hash::from(hash),
772            None => blake3::Hash::from(*EMPTY_HASH),
773        };
774        debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Updating global monotree with root: {root_str}");
775
776        // Iterate over contracts updates
777        for (contract_id_bytes, (rebuild, mut inserts)) in contracts_updates {
778            let contract_id: ContractId = deserialize(&contract_id_bytes)?;
779            let state_monotree_ptr = contract_id.hash_state_id(SMART_CONTRACT_MONOTREE_DB_NAME);
780            debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Updating monotree for contract: {contract_id}");
781
782            // Check if we need to rebuild it
783            if rebuild {
784                // Grab its state pointers
785                let state_pointers = lock.get(SLED_CONTRACTS_TREE, &contract_id_bytes)?.unwrap();
786                let mut state_pointers: Vec<[u8; 32]> = deserialize(&state_pointers)?;
787
788                // Skip native zkas trees
789                if NATIVE_CONTRACT_IDS_BYTES.contains(&contract_id.to_bytes()) {
790                    state_pointers.retain(|ptr| !NATIVE_CONTRACT_ZKAS_DB_NAMES.contains(ptr));
791                }
792
793                // Iterate over all contract states to grab the monotree keys
794                debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Rebuilding monotree...");
795                for state_ptr in state_pointers {
796                    // Open the contract state
797                    lock.open_tree(&state_ptr, false)?;
798
799                    // If the pointer is the monotree one, clear it
800                    if state_ptr == state_monotree_ptr {
801                        lock.clear(&state_ptr)?;
802                        continue
803                    }
804
805                    // Grab all its keys
806                    for record in lock.iter(&state_ptr)? {
807                        let (key, value) = record?;
808                        let key = blake3::hash(&key);
809                        let value = blake3::hash(&value);
810                        inserts.push((key, value))
811                    }
812                }
813            }
814
815            // Grab its monotree
816            let state_monotree_db = SledOverlayDb::new(&mut lock, &state_monotree_ptr)?;
817            let mut state_monotree = Monotree::new(state_monotree_db);
818            let mut state_monotree_root =
819                if rebuild { None } else { state_monotree.get_headroot()? };
820            let state_monotree_root_str = match state_monotree_root {
821                Some(hash) => blake3::Hash::from(hash),
822                None => blake3::Hash::from(*EMPTY_HASH),
823            };
824            debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Current root: {state_monotree_root_str}");
825
826            // Update or insert new records
827            for (key, value) in &inserts {
828                debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Inserting key {key} with value: {value}");
829                state_monotree_root = state_monotree.insert(
830                    state_monotree_root.as_ref(),
831                    key.as_bytes(),
832                    value.as_bytes(),
833                )?;
834            }
835
836            // Set root
837            state_monotree.set_headroot(state_monotree_root.as_ref());
838
839            // Insert its root to the global monotree
840            let state_monotree_root = match state_monotree_root {
841                Some(hash) => hash,
842                None => *EMPTY_HASH,
843            };
844            debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "New root: {}", blake3::Hash::from(state_monotree_root));
845            root = tree.insert(root.as_ref(), &contract_id.to_bytes(), &state_monotree_root)?;
846        }
847
848        // Check if wasm bincodes cache exists
849        let Some((state_cache, _)) = diff.caches.get(SLED_CONTRACTS_TREES_TREE) else {
850            debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "New global root: {}", blake3::Hash::from(root.unwrap()));
851            tree.set_headroot(root.as_ref());
852            return Ok(())
853        };
854
855        // Check if wasm bincodes cache is updated
856        if state_cache.cache.is_empty() && state_cache.removed.is_empty() {
857            debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "New global root: {}", blake3::Hash::from(root.unwrap()));
858            tree.set_headroot(root.as_ref());
859            return Ok(())
860        }
861
862        // Iterate over current contracts wasm bincodes to compute its monotree root
863        debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Updating wasm bincodes monotree...");
864        let mut wasm_monotree_root = None;
865        let wasm_monotree_db = MemoryDb::new();
866        let mut wasm_monotree = Monotree::new(wasm_monotree_db);
867        for record in lock.iter(SLED_BINCODE_TREE)? {
868            let (key, value) = record?;
869
870            // Skip native ones
871            if NATIVE_CONTRACT_IDS_BYTES.contains(&deserialize(&key)?) {
872                continue
873            }
874
875            // Insert record
876            let key = blake3::hash(&key);
877            let value = blake3::hash(&value);
878            debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "Inserting key {key} with value: {value}");
879            wasm_monotree_root = wasm_monotree.insert(
880                wasm_monotree_root.as_ref(),
881                key.as_bytes(),
882                value.as_bytes(),
883            )?;
884        }
885
886        // Insert wasm bincodes root to the global monotree
887        let wasm_monotree_root = match wasm_monotree_root {
888            Some(hash) => hash,
889            None => *EMPTY_HASH,
890        };
891        debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "New root: {}", blake3::Hash::from(wasm_monotree_root));
892        root = tree.insert(
893            root.as_ref(),
894            blake3::hash(SLED_BINCODE_TREE).as_bytes(),
895            &wasm_monotree_root,
896        )?;
897        debug!(target: "blockchain::contractstoreoverlay::update_state_monotree", "New global root: {}", blake3::Hash::from(root.unwrap()));
898        tree.set_headroot(root.as_ref());
899
900        Ok(())
901    }
902}
darkfi/blockchain/contract_store.rs

darkfi/blockchain/
contract_store.rs
