Selaa lähdekoodia

sdk/wasm: Add _local db functions for wasm contract usage.

This adds the contract-side API for DEP-0008.
x 5 kuukautta sitten
vanhempi
sitoutus
49a85c9d12
4 muutettua tiedostoa jossa 335 lisäystä ja 102 poistoa
  1. 20 13
      src/runtime/import/merkle.rs
  2. 238 80
      src/sdk/src/wasm/db.rs
  3. 65 9
      src/sdk/src/wasm/merkle.rs
  4. 12 0
      src/serial/src/lib.rs

+ 20 - 13
src/runtime/import/merkle.rs

@@ -19,8 +19,9 @@
 use std::io::Cursor;
 
 use darkfi_sdk::{
-    crypto::{MerkleNode, MerkleTree},
+    crypto::{pasta_prelude::Field, MerkleNode, MerkleTree},
     hex::AsHex,
+    pasta::pallas,
     wasm,
 };
 use darkfi_serial::{serialize, Decodable, Encodable, WriteExt};
@@ -206,7 +207,7 @@ pub(crate) fn merkle_add_internal(
 
     // Read the current Merkle tree.
     let tree_bytes = if local {
-        let Some(db_cid) = tx_local_db.get(&db_info.contract_id) else {
+        let Some(db_cid) = tx_local_db.get_mut(&db_info.contract_id) else {
             error!(
                 target: "runtime::db::{lt}",
                 "[WASM] [{cid}] {lt}(): Could not find db for {}",
@@ -215,16 +216,22 @@ pub(crate) fn merkle_add_internal(
             return darkfi_sdk::error::INTERNAL_ERROR
         };
 
-        let Some(tree) = db_cid.get(&db_info.tree) else {
-            error!(
-                target: "runtime::db::{lt}",
-                "[WASM] [{cid}] {lt}(): Could not find db tree for {}",
-                db_info.contract_id,
-            );
-            return darkfi_sdk::error::INTERNAL_ERROR
-        };
-
-        tree.get(&tree_key).cloned()
+        // Fetch or initialize this db tree
+        let tree = db_cid.entry(db_info.tree).or_default();
+
+        match tree.get(&tree_key) {
+            Some(v) => Some(v).cloned(),
+            None => {
+                // If our tx-local db does not contain the Merkle tree,
+                // initialize it with a "zero" leaf.
+                let mut merkle_tree = MerkleTree::new(1);
+                merkle_tree.append(MerkleNode::from(pallas::Base::ZERO));
+                let mut merkle_tree_data = vec![];
+                merkle_tree_data.write_u32(0).unwrap();
+                merkle_tree.encode(&mut merkle_tree_data).unwrap();
+                Some(merkle_tree_data)
+            }
+        }
     } else {
         match overlay.get(&db_info.tree, &tree_key) {
             Ok(v) => v.map(|iv| iv.to_vec()),
@@ -338,7 +345,7 @@ pub(crate) fn merkle_add_internal(
         let info_tree = db_cid.get_mut(&db_info.tree).unwrap();
         info_tree.insert(root_key, latest_root_data.clone());
 
-        let roots_tree = db_cid.get_mut(&db_roots.tree).unwrap();
+        let roots_tree = db_cid.entry(db_roots.tree).or_default();
         roots_tree.insert(latest_root_data, value_data);
     } else {
         if let Err(e) = overlay.insert(&db_roots.tree, &latest_root_data, &value_data) {

+ 238 - 80
src/sdk/src/wasm/db.rs

@@ -26,75 +26,170 @@ use crate::{
 
 pub type DbHandle = u32;
 
-/// Create a new database instance for the given contract.
-/// This should be called in the `init_contract()` section to create any databases
-/// that the contract might need or use.
+/// Create a new on-chain database instance for the given contract.
+/// A contract is only able to create a db for itself.
 ///
 /// Returns a `DbHandle` which provides methods for reading and writing.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
 pub fn db_init(contract_id: ContractId, db_name: &str) -> GenericResult<DbHandle> {
-    unsafe {
-        let mut len = 0;
-        let mut buf = vec![];
-        len += contract_id.encode(&mut buf)?;
-        len += db_name.to_string().encode(&mut buf)?;
-
-        let ret = db_init_(buf.as_ptr(), len as u32);
+    let mut len = 0;
+    let mut buf = vec![];
+    len += contract_id.encode(&mut buf)?;
+    len += db_name.to_string().encode(&mut buf)?;
 
-        if ret < 0 {
-            return Err(ContractError::from(ret))
-        }
+    let ret = unsafe { db_init_(buf.as_ptr(), len as u32) };
 
-        Ok(ret as u32)
+    if ret < 0 {
+        return Err(ContractError::from(ret))
     }
+
+    Ok(ret as u32)
 }
 
-/// Everyone can call this. Assumes that the database already went through `db_init()`.
+/// Open an existing on-chain database instance for the given contract.
+/// A contract is able to read any on-chain database.
+///
+/// Returns a `DbHandle` which is used with methods for reading and writing.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+/// * `ContractSection::Update`
 pub fn db_lookup(contract_id: ContractId, db_name: &str) -> GenericResult<DbHandle> {
-    unsafe {
-        let mut len = 0;
-        let mut buf = vec![];
-        len += contract_id.encode(&mut buf)?;
-        len += db_name.to_string().encode(&mut buf)?;
+    db_lookup_internal(contract_id, db_name, false)
+}
 
-        let ret = db_lookup_(buf.as_ptr(), len as u32);
+/// Open a tx-local database instance for the given contract.
+/// A contract is able to read any tx-local database.
+///
+/// If the calling contract is opening its own db, the db will be created
+/// and initialized in-memory.
+///
+/// Returns a `DbHandle` which is used with methods for reading and writing.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+/// * `ContractSection::Update`
+pub fn db_lookup_local(contract_id: ContractId, db_name: &str) -> GenericResult<DbHandle> {
+    db_lookup_internal(contract_id, db_name, true)
+}
 
-        if ret < 0 {
-            return Err(ContractError::from(ret))
+/// Internal function for `db_lookup` which branches to either on-chain or
+/// transaction-local.
+fn db_lookup_internal(
+    contract_id: ContractId,
+    db_name: &str,
+    local: bool,
+) -> GenericResult<DbHandle> {
+    let mut len = 0;
+    let mut buf = vec![];
+    len += contract_id.encode(&mut buf)?;
+    len += db_name.to_string().encode(&mut buf)?;
+
+    let ret = unsafe {
+        if local {
+            db_lookup_local_(buf.as_ptr(), len as u32)
+        } else {
+            db_lookup_(buf.as_ptr(), len as u32)
         }
+    };
 
-        Ok(ret as u32)
+    if ret < 0 {
+        return Err(ContractError::from(ret))
     }
+
+    Ok(ret as u32)
 }
 
-/// Everyone can call this. Will read a key from the key-value store.
+/// Read a key from the on-chain key-value store given a `DbHandle` and `key`.
+///
+/// Returns the `Vec<u8>` value if the key exists, otherwise `None`.
 ///
-/// ```
-/// value = db_get(db_handle, key);
-/// ```
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
 pub fn db_get(db_handle: DbHandle, key: &[u8]) -> GenericResult<Option<Vec<u8>>> {
+    db_get_internal(db_handle, key, false)
+}
+
+/// Read a key from the tx-local key-value store given a `DbHandle` and `key`.
+///
+/// Returns the `Vec<u8>` value if the key exists, otherwise `None`.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+pub fn db_get_local(db_handle: DbHandle, key: &[u8]) -> GenericResult<Option<Vec<u8>>> {
+    db_get_internal(db_handle, key, true)
+}
+
+/// Internal function for `db_get` which branches to either on-chain or
+/// transaction-local.
+fn db_get_internal(db_handle: DbHandle, key: &[u8], local: bool) -> GenericResult<Option<Vec<u8>>> {
     let mut len = 0;
     let mut buf = vec![];
     len += db_handle.encode(&mut buf)?;
-    len += key.to_vec().encode(&mut buf)?;
+    len += key.encode(&mut buf)?;
+
+    let ret = unsafe {
+        if local {
+            db_get_local_(buf.as_ptr(), len as u32)
+        } else {
+            db_get_(buf.as_ptr(), len as u32)
+        }
+    };
 
-    let ret = unsafe { db_get_(buf.as_ptr(), len as u32) };
     wasm::util::parse_ret(ret)
 }
 
-/// Everyone can call this. Checks if a key is contained in the key-value store.
+/// Check if a key is contained in the on-chain key-value store given a
+/// `DbHandle` and `key`.
+///
+/// Returns a boolean value.
 ///
-/// ```
-/// if db_contains_key(db_handle, key) {
-///     println!("true");
-/// }
-/// ```
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
 pub fn db_contains_key(db_handle: DbHandle, key: &[u8]) -> GenericResult<bool> {
+    db_contains_key_internal(db_handle, key, false)
+}
+
+/// Check if a key is contained in the tx-local key-value store given a
+/// `DbHandle` and `key`.
+///
+/// Returns a boolean value.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+pub fn db_contains_key_local(db_handle: DbHandle, key: &[u8]) -> GenericResult<bool> {
+    db_contains_key_internal(db_handle, key, true)
+}
+
+/// Internal function for `db_contains_key` which branches to either on-chain
+/// or transaction-local.
+fn db_contains_key_internal(db_handle: DbHandle, key: &[u8], local: bool) -> GenericResult<bool> {
     let mut len = 0;
     let mut buf = vec![];
     len += db_handle.encode(&mut buf)?;
-    len += key.to_vec().encode(&mut buf)?;
+    len += key.encode(&mut buf)?;
 
-    let ret = unsafe { db_contains_key_(buf.as_ptr(), len as u32) };
+    let ret = unsafe {
+        if local {
+            db_contains_key_local_(buf.as_ptr(), len as u32)
+        } else {
+            db_contains_key_(buf.as_ptr(), len as u32)
+        }
+    };
 
     if ret < 0 {
         return Err(ContractError::from(ret))
@@ -107,77 +202,140 @@ pub fn db_contains_key(db_handle: DbHandle, key: &[u8]) -> GenericResult<bool> {
     }
 }
 
-/// Only update() can call this. Set a value within the transaction.
+/// Set a key and value in the on-chain database for the given `DbHandle`.
 ///
-/// ```
-/// db_set(tx_handle, key, value);
-/// ```
+/// Returns `Ok` on success.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
 pub fn db_set(db_handle: DbHandle, key: &[u8], value: &[u8]) -> GenericResult<()> {
-    // Check entry for tx_handle is not None
-    unsafe {
-        let mut len = 0;
-        let mut buf = vec![];
-        len += db_handle.encode(&mut buf)?;
-        len += key.to_vec().encode(&mut buf)?;
-        len += value.to_vec().encode(&mut buf)?;
-
-        let ret = db_set_(buf.as_ptr(), len as u32);
-
-        if ret != wasm::entrypoint::SUCCESS {
-            return Err(ContractError::from(ret))
+    db_set_internal(db_handle, key, value, false)
+}
+
+/// Set a key and value in the tx-local database for the given `DbHandle`.
+///
+/// Returns `Ok` on success.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
+pub fn db_set_local(db_handle: DbHandle, key: &[u8], value: &[u8]) -> GenericResult<()> {
+    db_set_internal(db_handle, key, value, true)
+}
+
+/// Internal function for `db_set` which branches to either on-chain or
+/// transaction-local.
+fn db_set_internal(
+    db_handle: DbHandle,
+    key: &[u8],
+    value: &[u8],
+    local: bool,
+) -> GenericResult<()> {
+    let mut len = 0;
+    let mut buf = vec![];
+    len += db_handle.encode(&mut buf)?;
+    len += key.encode(&mut buf)?;
+    len += value.encode(&mut buf)?;
+
+    let ret = unsafe {
+        if local {
+            db_set_local_(buf.as_ptr(), len as u32)
+        } else {
+            db_set_(buf.as_ptr(), len as u32)
         }
+    };
 
-        Ok(())
+    if ret != wasm::entrypoint::SUCCESS {
+        return Err(ContractError::from(ret))
     }
+
+    Ok(())
 }
 
-/// Only update() can call this. Removes a key from the db.
+/// Remove a key from the on-chain database given a `DbHandle` and `key`.
 ///
-/// ```
-///     db_del(tx_handle, key);
-/// ```
+/// Returns `Ok` on success.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
 pub fn db_del(db_handle: DbHandle, key: &[u8]) -> GenericResult<()> {
-    // Check entry for tx_handle is not None
-    unsafe {
-        let mut len = 0;
-        let mut buf = vec![];
-        len += db_handle.encode(&mut buf)?;
-        len += key.to_vec().encode(&mut buf)?;
+    db_del_internal(db_handle, key, false)
+}
 
-        let ret = db_del_(buf.as_ptr(), len as u32);
+/// Remove a key from the tx-local database given a `DbHandle` and `key`.
+///
+/// Returns `Ok` on success.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
+pub fn db_del_local(db_handle: DbHandle, key: &[u8]) -> GenericResult<()> {
+    db_del_internal(db_handle, key, true)
+}
 
-        if ret != wasm::entrypoint::SUCCESS {
-            return Err(ContractError::from(ret))
+/// Internal function for `db_del` which branches to either on-chain or
+/// transaction-local.
+fn db_del_internal(db_handle: DbHandle, key: &[u8], local: bool) -> GenericResult<()> {
+    let mut len = 0;
+    let mut buf = vec![];
+    len += db_handle.encode(&mut buf)?;
+    len += key.encode(&mut buf)?;
+
+    let ret = unsafe {
+        if local {
+            db_del_local_(buf.as_ptr(), len as u32)
+        } else {
+            db_del_(buf.as_ptr(), len as u32)
         }
+    };
 
-        Ok(())
+    if ret != wasm::entrypoint::SUCCESS {
+        return Err(ContractError::from(ret))
     }
+
+    Ok(())
 }
 
-/// Only deploy() can call this.
+/// Given a zkas circuit, create a VerifyingKey and insert them both
+/// into the on-chain db.
+///
+/// Returns `Ok` on success, otherwise returns an error code.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
 pub fn zkas_db_set(bincode: &[u8]) -> GenericResult<()> {
-    unsafe {
-        let mut len = 0;
-        let mut buf = vec![];
-        len += bincode.to_vec().encode(&mut buf)?;
-
-        let ret = zkas_db_set_(buf.as_ptr(), len as u32);
+    let mut len = 0;
+    let mut buf = vec![];
+    len += bincode.encode(&mut buf)?;
 
-        if ret != wasm::entrypoint::SUCCESS {
-            return Err(ContractError::from(ret))
-        }
+    let ret = unsafe { zkas_db_set_(buf.as_ptr(), len as u32) };
 
-        Ok(())
+    if ret != wasm::entrypoint::SUCCESS {
+        return Err(ContractError::from(ret))
     }
+
+    Ok(())
 }
 
 extern "C" {
     fn db_init_(ptr: *const u8, len: u32) -> i64;
+
     fn db_lookup_(ptr: *const u8, len: u32) -> i64;
+    fn db_lookup_local_(ptr: *const u8, len: u32) -> i64;
+
     fn db_get_(ptr: *const u8, len: u32) -> i64;
+    fn db_get_local_(ptr: *const u8, len: u32) -> i64;
+
     fn db_contains_key_(ptr: *const u8, len: u32) -> i64;
+    fn db_contains_key_local_(ptr: *const u8, len: u32) -> i64;
+
     fn db_set_(ptr: *const u8, len: u32) -> i64;
+    fn db_set_local_(ptr: *const u8, len: u32) -> i64;
+
     fn db_del_(ptr: *const u8, len: u32) -> i64;
+    fn db_del_local_(ptr: *const u8, len: u32) -> i64;
 
     fn zkas_db_set_(ptr: *const u8, len: u32) -> i64;
 }

+ 65 - 9
src/sdk/src/wasm/merkle.rs

@@ -25,7 +25,7 @@ use crate::{
     wasm::db::DbHandle,
 };
 
-/// Add given elements into a Merkle tree. Used for inclusion proofs.
+/// Add given elements into an on-chain Merkle tree. Used for inclusion proofs.
 ///
 /// * `db_info` is a handle for a database where the Merkle tree is stored.
 /// * `db_roots` is a handle for a database where all the new Merkle roots are stored.
@@ -53,19 +53,74 @@ pub fn merkle_add(
     root_key: &[u8],
     tree_key: &[u8],
     elements: &[MerkleNode],
+) -> GenericResult<()> {
+    merkle_add_internal(db_info, db_roots, root_key, tree_key, elements, false)
+}
+
+/// Add given elements into a tx-local Merkle tree. Used for inclusion proofs.
+///
+/// * `db_info` is a handle for a database where the Merkle tree is stored.
+/// * `db_roots` is a handle for a database where all the new Merkle roots are stored.
+/// * `root_key` is the serialized key pointing to the latest Merkle root in `db_info`
+/// * `tree_key` is the serialized key pointing to the Merkle tree in `db_info`.
+/// * `elements` are the items we want to add to the Merkle tree.
+///
+/// There are 2 databases:
+///
+/// * `db_info` stores general metadata or info.
+/// * `db_roots` stores a log of all the merkle roots.
+///
+/// Inside `db_info` we store:
+///
+/// * The \[latest root hash:32\] under `root_key`.
+/// * The incremental merkle tree under `tree_key`.
+///
+/// Inside `db_roots` we store:
+///
+/// * All \[merkle root:32\]s as keys. The value is the current \[tx_hash:32\]\[call_idx:1\].
+///   If no new values are added, then the root key is updated to the current (tx_hash, call_idx).
+pub fn merkle_add_local(
+    db_info: DbHandle,
+    db_roots: DbHandle,
+    root_key: &[u8],
+    tree_key: &[u8],
+    elements: &[MerkleNode],
+) -> GenericResult<()> {
+    merkle_add_internal(db_info, db_roots, root_key, tree_key, elements, true)
+}
+
+/// Internal function for `merkle_add` which branches to either on-chain or
+/// transaction-local.
+pub fn merkle_add_internal(
+    db_info: DbHandle,
+    db_roots: DbHandle,
+    root_key: &[u8],
+    tree_key: &[u8],
+    elements: &[MerkleNode],
+    local: bool,
 ) -> GenericResult<()> {
     let mut buf = vec![];
     let mut len = 0;
     len += db_info.encode(&mut buf)?;
     len += db_roots.encode(&mut buf)?;
-    len += root_key.to_vec().encode(&mut buf)?;
-    len += tree_key.to_vec().encode(&mut buf)?;
-    len += elements.to_vec().encode(&mut buf)?;
+    len += root_key.encode(&mut buf)?;
+    len += tree_key.encode(&mut buf)?;
+    len += elements.encode(&mut buf)?;
 
-    match unsafe { merkle_add_(buf.as_ptr(), len as u32) } {
+    let ret = unsafe {
+        if local {
+            merkle_add_local_(buf.as_ptr(), len as u32)
+        } else {
+            merkle_add_(buf.as_ptr(), len as u32)
+        }
+    };
+
+    if ret < 0 {
+        return Err(ContractError::from(ret))
+    }
+
+    match ret {
         0 => Ok(()),
-        -1 => Err(ContractError::CallerAccessDenied),
-        -2 => Err(ContractError::DbSetFailed),
         _ => unreachable!(),
     }
 }
@@ -103,8 +158,8 @@ pub fn sparse_merkle_insert_batch(
     len += db_info.encode(&mut buf)?;
     len += db_smt.encode(&mut buf)?;
     len += db_roots.encode(&mut buf)?;
-    len += root_key.to_vec().encode(&mut buf)?;
-    len += elements.to_vec().encode(&mut buf)?;
+    len += root_key.encode(&mut buf)?;
+    len += elements.encode(&mut buf)?;
 
     match unsafe { sparse_merkle_insert_batch_(buf.as_ptr(), len as u32) } {
         0 => Ok(()),
@@ -116,5 +171,6 @@ pub fn sparse_merkle_insert_batch(
 
 extern "C" {
     fn merkle_add_(ptr: *const u8, len: u32) -> i64;
+    fn merkle_add_local_(ptr: *const u8, len: u32) -> i64;
     fn sparse_merkle_insert_batch_(ptr: *const u8, len: u32) -> i64;
 }

+ 12 - 0
src/serial/src/lib.rs

@@ -510,6 +510,18 @@ impl<T: Encodable> Encodable for Vec<T> {
     }
 }
 
+impl<T: Encodable> Encodable for &[T] {
+    #[inline]
+    fn encode<S: Write>(&self, s: &mut S) -> Result<usize, Error> {
+        let mut len = 0;
+        len += VarInt(self.len() as u64).encode(s)?;
+        for val in self.iter() {
+            len += val.encode(s)?;
+        }
+        Ok(len)
+    }
+}
+
 impl<T: Decodable> Decodable for Vec<T> {
     #[inline]
     fn decode<D: Read>(d: &mut D) -> Result<Self, Error> {