From 1e5b23aeba4e1da4f814fb75191c6f99cb724c3e Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Mon, 19 Jun 2023 19:05:35 -0600 Subject: [PATCH] zcash_client_backend: Add `put_sapling_subtree_roots` to `WalletCommitmentTrees` Also add the `zcash_client_sqlite` implementation & tests for the new method. --- zcash_client_backend/CHANGELOG.md | 1 + zcash_client_backend/src/data_api.rs | 32 ++- zcash_client_backend/src/data_api/chain.rs | 29 +++ zcash_client_sqlite/src/lib.rs | 42 +++- .../src/wallet/commitment_tree.rs | 184 ++++++++++++++++-- 5 files changed, 270 insertions(+), 18 deletions(-) diff --git a/zcash_client_backend/CHANGELOG.md b/zcash_client_backend/CHANGELOG.md index 1cb28eaef..999d91860 100644 --- a/zcash_client_backend/CHANGELOG.md +++ b/zcash_client_backend/CHANGELOG.md @@ -18,6 +18,7 @@ and this library adheres to Rust's notion of - `NullifierQuery` for use with `WalletRead::get_sapling_nullifiers` - `BlockMetadata` - `ScannedBlock` + - `chain::CommitmentTreeRoot` - `wallet::input_sellection::Proposal::{min_target_height, min_anchor_height}`: - `zcash_client_backend::wallet::WalletSaplingOutput::note_commitment_tree_position` - `zcash_client_backend::scanning::ScanError` diff --git a/zcash_client_backend/src/data_api.rs b/zcash_client_backend/src/data_api.rs index bd281e282..7cc945f18 100644 --- a/zcash_client_backend/src/data_api.rs +++ b/zcash_client_backend/src/data_api.rs @@ -28,6 +28,8 @@ use crate::{ wallet::{ReceivedSaplingNote, WalletTransparentOutput, WalletTx}, }; +use self::chain::CommitmentTreeRoot; + pub mod chain; pub mod error; pub mod wallet; @@ -545,10 +547,18 @@ pub trait WalletCommitmentTrees { >, ) -> Result, E: From>; + + /// Adds a sequence of note commitment tree subtree roots to the data store. + fn put_sapling_subtree_roots( + &mut self, + start_index: u64, + roots: &[CommitmentTreeRoot], + ) -> Result<(), ShardTreeError>; } #[cfg(feature = "test-dependencies")] pub mod testing { + use incrementalmerkletree::Address; use secrecy::{ExposeSecret, SecretVec}; use shardtree::{MemoryShardStore, ShardTree, ShardTreeError}; use std::{collections::HashMap, convert::Infallible, ops::Range}; @@ -573,8 +583,9 @@ pub mod testing { }; use super::{ - BlockMetadata, DecryptedTransaction, NullifierQuery, ScannedBlock, SentTransaction, - WalletCommitmentTrees, WalletRead, WalletWrite, SAPLING_SHARD_HEIGHT, + chain::CommitmentTreeRoot, BlockMetadata, DecryptedTransaction, NullifierQuery, + ScannedBlock, SentTransaction, WalletCommitmentTrees, WalletRead, WalletWrite, + SAPLING_SHARD_HEIGHT, }; pub struct MockWalletDb { @@ -805,5 +816,22 @@ pub mod testing { { callback(&mut self.sapling_tree) } + + fn put_sapling_subtree_roots( + &mut self, + start_index: u64, + roots: &[CommitmentTreeRoot], + ) -> Result<(), ShardTreeError> { + self.with_sapling_tree_mut(|t| { + for (root, i) in roots.iter().zip(0u64..) { + let root_addr = + Address::from_parts(SAPLING_SHARD_HEIGHT.into(), start_index + i); + t.insert(root_addr, *root.root_hash())?; + } + Ok::<_, ShardTreeError>(()) + })?; + + Ok(()) + } } } diff --git a/zcash_client_backend/src/data_api/chain.rs b/zcash_client_backend/src/data_api/chain.rs index fab86eaf6..16d530ac7 100644 --- a/zcash_client_backend/src/data_api/chain.rs +++ b/zcash_client_backend/src/data_api/chain.rs @@ -66,6 +66,35 @@ use crate::{ pub mod error; use error::Error; +/// A struct containing metadata about a subtree root of the note commitment tree. +/// +/// This stores the block height at which the leaf that completed the subtree was +/// added, and the root hash of the complete subtree. +pub struct CommitmentTreeRoot { + subtree_end_height: BlockHeight, + root_hash: H, +} + +impl CommitmentTreeRoot { + /// Construct a new `CommitmentTreeRoot` from its constituent parts. + pub fn from_parts(subtree_end_height: BlockHeight, root_hash: H) -> Self { + Self { + subtree_end_height, + root_hash, + } + } + + /// Returns the block height at which the leaf that completed the subtree was added. + pub fn subtree_end_height(&self) -> BlockHeight { + self.subtree_end_height + } + + /// Returns the root of the complete subtree. + pub fn root_hash(&self) -> &H { + &self.root_hash + } +} + /// This trait provides sequential access to raw blockchain data via a callback-oriented /// API. pub trait BlockSource { diff --git a/zcash_client_sqlite/src/lib.rs b/zcash_client_sqlite/src/lib.rs index b7351edf8..2f65053c4 100644 --- a/zcash_client_sqlite/src/lib.rs +++ b/zcash_client_sqlite/src/lib.rs @@ -36,6 +36,7 @@ use either::Either; use rusqlite::{self, Connection}; use secrecy::{ExposeSecret, SecretVec}; use std::{borrow::Borrow, collections::HashMap, convert::AsRef, fmt, io, ops::Range, path::Path}; +use wallet::commitment_tree::put_shard_roots; use incrementalmerkletree::Position; use shardtree::{ShardTree, ShardTreeError}; @@ -55,9 +56,10 @@ use zcash_primitives::{ use zcash_client_backend::{ address::{AddressMetadata, UnifiedAddress}, data_api::{ - self, chain::BlockSource, BlockMetadata, DecryptedTransaction, NullifierQuery, PoolType, - Recipient, ScannedBlock, SentTransaction, WalletCommitmentTrees, WalletRead, WalletWrite, - SAPLING_SHARD_HEIGHT, + self, + chain::{BlockSource, CommitmentTreeRoot}, + BlockMetadata, DecryptedTransaction, NullifierQuery, PoolType, Recipient, ScannedBlock, + SentTransaction, WalletCommitmentTrees, WalletRead, WalletWrite, SAPLING_SHARD_HEIGHT, }, keys::{UnifiedFullViewingKey, UnifiedSpendingKey}, proto::compact_formats::CompactBlock, @@ -643,10 +645,31 @@ impl WalletCommitmentTrees for WalletDb], + ) -> Result<(), ShardTreeError> { + let tx = self + .conn + .transaction() + .map_err(|e| ShardTreeError::Storage(Either::Right(e)))?; + put_shard_roots::<_, { sapling::NOTE_COMMITMENT_TREE_DEPTH }, SAPLING_SHARD_HEIGHT>( + &tx, + SAPLING_TABLES_PREFIX, + start_index, + roots, + )?; + tx.commit() + .map_err(|e| ShardTreeError::Storage(Either::Right(e)))?; + Ok(()) + } } impl<'conn, P: consensus::Parameters> WalletCommitmentTrees for WalletDb, P> { @@ -674,6 +697,19 @@ impl<'conn, P: consensus::Parameters> WalletCommitmentTrees for WalletDb], + ) -> Result<(), ShardTreeError> { + put_shard_roots::<_, { sapling::NOTE_COMMITMENT_TREE_DEPTH }, SAPLING_SHARD_HEIGHT>( + self.conn.0, + SAPLING_TABLES_PREFIX, + start_index, + roots, + ) + } } /// A handle for the SQLite block source. diff --git a/zcash_client_sqlite/src/wallet/commitment_tree.rs b/zcash_client_sqlite/src/wallet/commitment_tree.rs index 78548836a..12cf24333 100644 --- a/zcash_client_sqlite/src/wallet/commitment_tree.rs +++ b/zcash_client_sqlite/src/wallet/commitment_tree.rs @@ -4,10 +4,15 @@ use std::{ collections::BTreeSet, io::{self, Cursor}, marker::PhantomData, + rc::Rc, }; +use zcash_client_backend::data_api::chain::CommitmentTreeRoot; -use incrementalmerkletree::{Address, Level, Position}; -use shardtree::{Checkpoint, LocatedPrunableTree, PrunableTree, ShardStore, TreeState}; +use incrementalmerkletree::{Address, Hashable, Level, Position, Retention}; +use shardtree::{ + Checkpoint, LocatedPrunableTree, LocatedTree, PrunableTree, RetentionFlags, ShardStore, + ShardTreeError, TreeState, +}; use zcash_primitives::{consensus::BlockHeight, merkle_tree::HashSer}; @@ -257,23 +262,29 @@ type Error = Either; pub(crate) fn get_shard( conn: &rusqlite::Connection, table_prefix: &'static str, - shard_root: Address, + shard_root_addr: Address, ) -> Result>, Error> { conn.query_row( &format!( - "SELECT shard_data + "SELECT shard_data, root_hash FROM {}_tree_shards WHERE shard_index = :shard_index", table_prefix ), - named_params![":shard_index": shard_root.index()], - |row| row.get::<_, Vec>(0), + named_params![":shard_index": shard_root_addr.index()], + |row| Ok((row.get::<_, Vec>(0)?, row.get::<_, Option>>(1)?)), ) .optional() .map_err(Either::Right)? - .map(|shard_data| { + .map(|(shard_data, root_hash)| { let shard_tree = read_shard(&mut Cursor::new(shard_data)).map_err(Either::Left)?; - Ok(LocatedPrunableTree::from_parts(shard_root, shard_tree)) + let located_tree = LocatedPrunableTree::from_parts(shard_root_addr, shard_tree); + if let Some(root_hash_data) = root_hash { + let root_hash = H::read(Cursor::new(root_hash_data)).map_err(Either::Left)?; + Ok(located_tree.reannotate_root(Some(Rc::new(root_hash)))) + } else { + Ok(located_tree) + } }) .transpose() } @@ -746,18 +757,102 @@ pub(crate) fn truncate_checkpoints( Ok(()) } +pub(crate) fn put_shard_roots< + H: Hashable + HashSer + Clone + Eq, + const DEPTH: u8, + const SHARD_HEIGHT: u8, +>( + conn: &rusqlite::Transaction<'_>, + table_prefix: &'static str, + start_index: u64, + roots: &[CommitmentTreeRoot], +) -> Result<(), ShardTreeError> { + if roots.is_empty() { + // nothing to do + return Ok(()); + } + + // We treat the cap as a DEPTH-SHARD_HEIGHT tree so that we can make a batch insertion of + // root data using `Position::from(start_index)` as the starting position and treating the + // roots as level-0 leaves. + let cap = LocatedTree::from_parts( + Address::from_parts((DEPTH - SHARD_HEIGHT).into(), 0), + get_cap(conn, table_prefix).map_err(ShardTreeError::Storage)?, + ); + + let cap_result = cap + .batch_insert( + Position::from(start_index), + roots.iter().map(|r| { + ( + r.root_hash().clone(), + Retention::Checkpoint { + id: (), + is_marked: false, + }, + ) + }), + ) + .map_err(ShardTreeError::Insert)? + .expect("slice of inserted roots was verified to be nonempty"); + + put_cap(conn, table_prefix, cap_result.subtree.take_root()).map_err(ShardTreeError::Storage)?; + + for (root, i) in roots.iter().zip(0u64..) { + // We want to avoid deserializing the subtree just to annotate its root node, so we simply + // cache the downloaded root alongside of any already-persisted subtree. We will update the + // subtree data itself by reannotating the root node of the tree, handling conflicts, at + // the time that we deserialize the tree. + let mut stmt = conn + .prepare_cached(&format!( + "INSERT INTO {}_tree_shards (shard_index, subtree_end_height, root_hash, shard_data) + VALUES (:shard_index, :subtree_end_height, :root_hash, :shard_data) + ON CONFLICT (shard_index) DO UPDATE + SET subtree_end_height = :subtree_end_height, root_hash = :root_hash", + table_prefix + )) + .map_err(|e| ShardTreeError::Storage(Either::Right(e)))?; + + // The `shard_data` value will only be used in the case that no tree already exists. + let mut shard_data: Vec = vec![]; + let tree = PrunableTree::leaf((root.root_hash().clone(), RetentionFlags::EPHEMERAL)); + write_shard(&mut shard_data, &tree) + .map_err(|e| ShardTreeError::Storage(Either::Left(e)))?; + + let mut root_hash_data: Vec = vec![]; + root.root_hash() + .write(&mut root_hash_data) + .map_err(|e| ShardTreeError::Storage(Either::Left(e)))?; + + stmt.execute(named_params![ + ":shard_index": start_index + i, + ":subtree_end_height": u32::from(root.subtree_end_height()), + ":root_hash": root_hash_data, + ":shard_data": shard_data, + ]) + .map_err(|e| ShardTreeError::Storage(Either::Right(e)))?; + } + + Ok(()) +} + #[cfg(test)] mod tests { use tempfile::NamedTempFile; - use incrementalmerkletree::testing::{ - check_append, check_checkpoint_rewind, check_remove_mark, check_rewind_remove_mark, - check_root_hashes, check_witness_consistency, check_witnesses, + use incrementalmerkletree::{ + testing::{ + check_append, check_checkpoint_rewind, check_remove_mark, check_rewind_remove_mark, + check_root_hashes, check_witness_consistency, check_witnesses, + }, + Position, Retention, }; use shardtree::ShardTree; + use zcash_client_backend::data_api::chain::CommitmentTreeRoot; + use zcash_primitives::consensus::BlockHeight; use super::SqliteShardStore; - use crate::{tests, wallet::init::init_wallet_db, WalletDb}; + use crate::{tests, wallet::init::init_wallet_db, WalletDb, SAPLING_TABLES_PREFIX}; fn new_tree(m: usize) -> ShardTree, 4, 3> { let data_file = NamedTempFile::new().unwrap(); @@ -766,7 +861,8 @@ mod tests { init_wallet_db(&mut db_data, None).unwrap(); let store = - SqliteShardStore::<_, String, 3>::from_connection(db_data.conn, "sapling").unwrap(); + SqliteShardStore::<_, String, 3>::from_connection(db_data.conn, SAPLING_TABLES_PREFIX) + .unwrap(); ShardTree::new(store, m) } @@ -804,4 +900,66 @@ mod tests { fn rewind_remove_mark() { check_rewind_remove_mark(new_tree); } + + #[test] + fn put_shard_roots() { + let data_file = NamedTempFile::new().unwrap(); + let mut db_data = WalletDb::for_path(data_file.path(), tests::network()).unwrap(); + data_file.keep().unwrap(); + + init_wallet_db(&mut db_data, None).unwrap(); + let tx = db_data.conn.transaction().unwrap(); + let store = + SqliteShardStore::<_, String, 3>::from_connection(&tx, SAPLING_TABLES_PREFIX).unwrap(); + + // introduce some roots + let roots = (0u32..4) + .into_iter() + .map(|idx| { + CommitmentTreeRoot::from_parts( + BlockHeight::from((idx + 1) * 3), + if idx == 3 { + "abcdefgh".to_string() + } else { + idx.to_string() + }, + ) + }) + .collect::>(); + super::put_shard_roots::<_, 6, 3>(store.conn, SAPLING_TABLES_PREFIX, 0, &roots).unwrap(); + + // simulate discovery of a note + let mut tree = ShardTree::<_, 6, 3>::new(store, 10); + tree.batch_insert( + Position::from(24), + ('a'..='h').into_iter().map(|c| { + ( + c.to_string(), + match c { + 'c' => Retention::Marked, + 'h' => Retention::Checkpoint { + id: BlockHeight::from(3), + is_marked: false, + }, + _ => Retention::Ephemeral, + }, + ) + }), + ) + .unwrap(); + + // construct a witness for the note + let witness = tree.witness(Position::from(26), 0).unwrap(); + assert_eq!( + witness.path_elems(), + &[ + "d", + "ab", + "efgh", + "2", + "01", + "________________________________" + ] + ); + } }