darkfi/runtime/import/
smt.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2025 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 *
10 * 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::io::Cursor;
20
21use darkfi_sdk::{
22    crypto::{
23        pasta_prelude::*,
24        smt::{PoseidonFp, SparseMerkleTree, StorageAdapter, EMPTY_NODES_FP, SMT_FP_DEPTH},
25    },
26    error::{ContractError, ContractResult},
27    wasm,
28};
29use darkfi_serial::{deserialize, serialize, Decodable, Encodable};
30use halo2_proofs::pasta::pallas;
31use log::{debug, error};
32use num_bigint::BigUint;
33use wasmer::{FunctionEnvMut, WasmPtr};
34
35use super::acl::acl_allow;
36use crate::runtime::vm_runtime::{ContractSection, Env};
37
38/// An SMT adapter for sled overlay storage. Compatible with the WasmDb SMT adapter
39pub struct SledStorage<'a> {
40    overlay: &'a mut sled_overlay::SledDbOverlay,
41    tree_key: &'a [u8],
42}
43
44impl StorageAdapter for SledStorage<'_> {
45    type Value = pallas::Base;
46
47    fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
48        if let Err(e) = self.overlay.insert(self.tree_key, &key.to_bytes_le(), &value.to_repr()) {
49            error!(
50                target: "runtime::smt::SledStorage::put",
51                "[WASM] SledStorage::put(): inserting key {:?}, value {:?} into DB tree: {:?}: {}",
52                key, value, self.tree_key, e,
53            );
54            return Err(ContractError::SmtPutFailed)
55        }
56
57        Ok(())
58    }
59
60    fn get(&self, key: &BigUint) -> Option<pallas::Base> {
61        let value = match self.overlay.get(self.tree_key, &key.to_bytes_le()) {
62            Ok(v) => v,
63            Err(e) => {
64                error!(
65                    target: "runtime::smt::SledStorage::get",
66                    "[WASM] SledStorage::get(): Fetching key {:?} from DB tree: {:?}: {}",
67                    key, self.tree_key, e,
68                );
69                return None
70            }
71        };
72
73        let value = value?;
74        let mut repr = [0; 32];
75        repr.copy_from_slice(&value);
76
77        pallas::Base::from_repr(repr).into()
78    }
79
80    fn del(&mut self, key: &BigUint) -> ContractResult {
81        if let Err(e) = self.overlay.remove(self.tree_key, &key.to_bytes_le()) {
82            error!(
83                target: "runtime::smt::SledStorage::del",
84                "[WASM] SledStorage::del(): Removing key {:?} from DB tree: {:?}: {}",
85                key, self.tree_key, e,
86            );
87            return Err(ContractError::SmtDelFailed)
88        }
89
90        Ok(())
91    }
92}
93
94/// Adds data to sparse merkle tree. The tree, database connection, and new data to add is
95/// read from `ptr` at offset specified by `len`.
96/// Returns `0` on success; otherwise, returns an error-code corresponding to a
97/// [`ContractError`] (defined in the SDK).
98/// See also the method `merkle_add` in `sdk/src/merkle.rs`.
99///
100/// Permissions: update
101pub(crate) fn sparse_merkle_insert_batch(
102    mut ctx: FunctionEnvMut<Env>,
103    ptr: WasmPtr<u8>,
104    len: u32,
105) -> i64 {
106    let (env, mut store) = ctx.data_and_store_mut();
107    let cid = env.contract_id;
108
109    // Enforce function ACL
110    if let Err(e) = acl_allow(env, &[ContractSection::Update]) {
111        error!(
112            target: "runtime::smt::sparse_merkle_insert_batch",
113            "[WASM] [{}] sparse_merkle_insert_batch(): Called in unauthorized section: {}", cid, e,
114        );
115        return darkfi_sdk::error::CALLER_ACCESS_DENIED
116    }
117
118    // Subtract used gas.
119    // This makes calling the function which returns early have some (small) cost.
120    env.subtract_gas(&mut store, 1);
121
122    let memory_view = env.memory_view(&store);
123    let Ok(mem_slice) = ptr.slice(&memory_view, len) else {
124        error!(
125            target: "runtime::smt::sparse_merkle_insert_batch",
126            "[WASM] [{}] sparse_merkle_insert_batch(): Failed to make slice from ptr", cid,
127        );
128        return darkfi_sdk::error::INTERNAL_ERROR
129    };
130
131    let mut buf = vec![0_u8; len as usize];
132    if let Err(e) = mem_slice.read_slice(&mut buf) {
133        error!(
134            target: "runtime::smt::sparse_merkle_insert_batch",
135            "[WASM] [{}] sparse_merkle_insert_batch(): Failed to read from memory slice: {}", cid, e,
136        );
137        return darkfi_sdk::error::INTERNAL_ERROR
138    };
139
140    // The buffer should deserialize into:
141    // - db_smt
142    // - db_roots
143    // - nullifiers (as Vec<pallas::Base>)
144    let mut buf_reader = Cursor::new(buf);
145    let db_info_index: u32 = match Decodable::decode(&mut buf_reader) {
146        Ok(v) => v,
147        Err(e) => {
148            error!(
149                target: "runtime::smt::sparse_merkle_insert_batch",
150                "[WASM] [{}] sparse_merkle_insert_batch(): Failed to decode db_info DbHandle: {}", cid, e,
151            );
152            return darkfi_sdk::error::INTERNAL_ERROR
153        }
154    };
155    let db_info_index = db_info_index as usize;
156
157    let db_smt_index: u32 = match Decodable::decode(&mut buf_reader) {
158        Ok(v) => v,
159        Err(e) => {
160            error!(
161            target: "runtime::smt::sparse_merkle_insert_batch",
162                "[WASM] [{}] sparse_merkle_insert_batch(): Failed to decode db_smt DbHandle: {}", cid, e,
163            );
164            return darkfi_sdk::error::INTERNAL_ERROR
165        }
166    };
167    let db_smt_index = db_smt_index as usize;
168
169    let db_roots_index: u32 = match Decodable::decode(&mut buf_reader) {
170        Ok(v) => v,
171        Err(e) => {
172            error!(
173                target: "runtime::smt::sparse_merkle_insert_batch",
174                "[WASM] [{}] sparse_merkle_insert_batch(): Failed to decode db_roots DbHandle: {}", cid, e,
175            );
176            return darkfi_sdk::error::INTERNAL_ERROR
177        }
178    };
179    let db_roots_index = db_roots_index as usize;
180
181    let db_handles = env.db_handles.borrow();
182    let n_dbs = db_handles.len();
183
184    if n_dbs <= db_info_index || n_dbs <= db_smt_index || n_dbs <= db_roots_index {
185        error!(
186            target: "runtime::smt::sparse_merkle_insert_batch",
187            "[WASM] [{}] sparse_merkle_insert_batch(): Requested DbHandle that is out of bounds", cid,
188        );
189        return darkfi_sdk::error::INTERNAL_ERROR
190    }
191    let db_info = &db_handles[db_info_index];
192    let db_smt = &db_handles[db_smt_index];
193    let db_roots = &db_handles[db_roots_index];
194
195    // Make sure that the contract owns the dbs it wants to write to
196    if db_info.contract_id != env.contract_id ||
197        db_smt.contract_id != env.contract_id ||
198        db_roots.contract_id != env.contract_id
199    {
200        error!(
201            target: "runtime::smt::sparse_merkle_insert_batch",
202            "[WASM] [{}] sparse_merkle_insert_batch(): Unauthorized to write to DbHandle", cid,
203        );
204        return darkfi_sdk::error::CALLER_ACCESS_DENIED
205    }
206
207    // This `key` represents the sled key in info where the latest root is
208    let root_key: Vec<u8> = match Decodable::decode(&mut buf_reader) {
209        Ok(v) => v,
210        Err(e) => {
211            error!(
212                target: "runtime::smt::sparse_merkle_insert_batch",
213                "[WASM] [{}] sparse_merkle_insert_batch(): Failed to decode key vec: {}", cid, e,
214            );
215            return darkfi_sdk::error::INTERNAL_ERROR
216        }
217    };
218
219    // This `nullifier` represents the leaf we're adding to the Merkle tree
220    let nullifiers: Vec<pallas::Base> = match Decodable::decode(&mut buf_reader) {
221        Ok(v) => v,
222        Err(e) => {
223            error!(
224                target: "runtime::smt::sparse_merkle_insert_batch",
225                "[WASM] [{}] sparse_merkle_insert_batch(): Failed to decode pallas::Base: {}", cid, e,
226            );
227            return darkfi_sdk::error::INTERNAL_ERROR
228        }
229    };
230
231    // Make sure we've read the entire buffer
232    if buf_reader.position() != (len as u64) {
233        error!(
234            target: "runtime::smt::sparse_merkle_insert_batch",
235            "[WASM] [{}] sparse_merkle_insert_batch(): Mismatch between given length, and cursor length", cid,
236        );
237        return darkfi_sdk::error::INTERNAL_ERROR
238    }
239
240    // Generate the SledStorage SMT
241    let hasher = PoseidonFp::new();
242    let lock = env.blockchain.lock().unwrap();
243    let mut overlay = lock.overlay.lock().unwrap();
244    let smt_store = SledStorage { overlay: &mut overlay, tree_key: &db_smt.tree };
245    let mut smt = SparseMerkleTree::<
246        SMT_FP_DEPTH,
247        { SMT_FP_DEPTH + 1 },
248        pallas::Base,
249        PoseidonFp,
250        SledStorage,
251    >::new(smt_store, hasher, &EMPTY_NODES_FP);
252
253    // Count the nullifiers for gas calculation
254    let inserted_nullifiers = nullifiers.len() * 32;
255
256    // Insert the new nullifiers
257    let leaves: Vec<_> = nullifiers.iter().map(|x| (*x, *x)).collect();
258    if let Err(e) = smt.insert_batch(leaves) {
259        error!(
260            target: "runtime::smt::sparse_merkle_insert_batch",
261            "[WASM] [{}] sparse_merkle_insert_batch(): SMT failed to insert batch: {}", cid, e,
262        );
263        return darkfi_sdk::error::INTERNAL_ERROR
264    };
265
266    // Grab the current SMT root to add in our set of roots.
267    // Since each update to the tree is atomic, we only need to add the last root.
268    let latest_root = smt.root();
269
270    // Validate latest root data, to ensure their integrity
271    let latest_root_data = serialize(&latest_root);
272    if latest_root_data.len() != 32 {
273        error!(
274            target: "runtime::smt::sparse_merkle_insert_batch",
275            "[WASM] [{}] sparse_merkle_insert_batch(): Latest root data length missmatch: {}", cid, latest_root_data.len(),
276        );
277        return darkfi_sdk::error::INTERNAL_ERROR
278    }
279
280    // Validate the new value data, to ensure their integrity
281    let mut new_value_data = Vec::with_capacity(32 + 1);
282    if let Err(e) = env.tx_hash.inner().encode(&mut new_value_data) {
283        error!(
284            target: "runtime::smt::sparse_merkle_insert_batch",
285            "[WASM] [{}] sparse_merkle_insert_batch(): Failed to serialize transaction hash: {}", cid, e,
286        );
287        return darkfi_sdk::error::INTERNAL_ERROR
288    };
289    if let Err(e) = env.call_idx.encode(&mut new_value_data) {
290        error!(
291            target: "runtime::smt::sparse_merkle_insert_batch",
292            "[WASM] [{}] sparse_merkle_insert_batch(): Failed to serialize call index: {}", cid, e,
293        );
294        return darkfi_sdk::error::INTERNAL_ERROR
295    };
296    if new_value_data.len() != 32 + 1 {
297        error!(
298            target: "runtime::smt::sparse_merkle_insert_batch",
299            "[WASM] [{}] sparse_merkle_insert_batch(): New value data length missmatch: {}", cid, new_value_data.len(),
300        );
301        return darkfi_sdk::error::INTERNAL_ERROR
302    }
303
304    // Retrieve snapshot root data set
305    let root_value_data_set = match overlay.get(&db_roots.tree, &latest_root_data) {
306        Ok(data) => data,
307        Err(e) => {
308            error!(
309                target: "runtime::smt::sparse_merkle_insert_batch",
310                "[WASM] [{}] sparse_merkle_insert_batch(): SMT failed to retrieve current root snapshot: {}", cid, e,
311            );
312            return darkfi_sdk::error::INTERNAL_ERROR
313        }
314    };
315
316    // If the record exists, append the new value data,
317    // otherwise create a new set with it.
318    let root_value_data_set = match root_value_data_set {
319        Some(value_data_set) => {
320            let mut value_data_set: Vec<Vec<u8>> = match deserialize(&value_data_set) {
321                Ok(set) => set,
322                Err(e) => {
323                    error!(
324                        target: "runtime::smt::sparse_merkle_insert_batch",
325                        "[WASM] [{}] sparse_merkle_insert_batch(): Failed to deserialize current root snapshot: {}", cid, e,
326                    );
327                    return darkfi_sdk::error::INTERNAL_ERROR
328                }
329            };
330
331            if !value_data_set.contains(&new_value_data) {
332                value_data_set.push(new_value_data);
333            }
334
335            value_data_set
336        }
337        None => vec![new_value_data],
338    };
339
340    // Write the latest root snapshot
341    debug!(
342        target: "runtime::smt::sparse_merkle_insert_batch",
343        "[WASM] [{}] sparse_merkle_insert_batch(): Appending SMT root to db: {:?}", cid, latest_root,
344    );
345    if overlay.insert(&db_roots.tree, &latest_root_data, &serialize(&root_value_data_set)).is_err()
346    {
347        error!(
348            target: "runtime::smt::sparse_merkle_insert_batch",
349            "[WASM] [{}] sparse_merkle_insert_batch(): Couldn't insert to db_roots tree", cid,
350        );
351        return darkfi_sdk::error::INTERNAL_ERROR
352    }
353
354    // Update the pointer to the latest known root
355    debug!(
356        target: "runtime::smt::sparse_merkle_insert_batch",
357        "[WASM] [{}] sparse_merkle_insert_batch(): Replacing latest SMT root pointer", cid,
358    );
359    if overlay.insert(&db_info.tree, &root_key, &latest_root_data).is_err() {
360        error!(
361            target: "runtime::smt::sparse_merkle_insert_batch",
362            "[WASM] [{}] sparse_merkle_insert_batch(): Couldn't insert latest root to db_info tree", cid,
363        );
364        return darkfi_sdk::error::INTERNAL_ERROR
365    }
366
367    // Subtract used gas.
368    // Here we count:
369    // * The number of nullifiers we inserted into the DB
370    drop(overlay);
371    drop(lock);
372    drop(db_handles);
373    env.subtract_gas(&mut store, inserted_nullifiers as u64);
374
375    wasm::entrypoint::SUCCESS
376}
darkfi/runtime/import/smt.rs

darkfi/runtime/import/
smt.rs
