zcash_client_backend: Add `put_sapling_subtree_roots` to `WalletCommitmentTrees`

Also add the `zcash_client_sqlite` implementation & tests for the new
method.
This commit is contained in:
Kris Nuttycombe 2023-06-19 19:05:35 -06:00
parent d8148f90e7
commit 1e5b23aeba
5 changed files with 270 additions and 18 deletions

View File

@ -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`

View File

@ -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<A, E>,
E: From<ShardTreeError<Self::Error>>;
/// 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<sapling::Node>],
) -> Result<(), ShardTreeError<Self::Error>>;
}
#[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<sapling::Node>],
) -> Result<(), ShardTreeError<Self::Error>> {
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<Self::Error>>(())
})?;
Ok(())
}
}
}

View File

@ -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<H> {
subtree_end_height: BlockHeight,
root_hash: H,
}
impl<H> CommitmentTreeRoot<H> {
/// 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 {

View File

@ -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<P: consensus::Parameters> WalletCommitmentTrees for WalletDb<rusqlite::Conn
let mut shardtree = ShardTree::new(shard_store, PRUNING_DEPTH.try_into().unwrap());
callback(&mut shardtree)?
};
tx.commit()
.map_err(|e| ShardTreeError::Storage(Either::Right(e)))?;
Ok(result)
}
fn put_sapling_subtree_roots(
&mut self,
start_index: u64,
roots: &[CommitmentTreeRoot<sapling::Node>],
) -> Result<(), ShardTreeError<Self::Error>> {
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<SqlTransaction<'conn>, P> {
@ -674,6 +697,19 @@ impl<'conn, P: consensus::Parameters> WalletCommitmentTrees for WalletDb<SqlTran
Ok(result)
}
fn put_sapling_subtree_roots(
&mut self,
start_index: u64,
roots: &[CommitmentTreeRoot<sapling::Node>],
) -> Result<(), ShardTreeError<Self::Error>> {
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.

View File

@ -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<io::Error, rusqlite::Error>;
pub(crate) fn get_shard<H: HashSer>(
conn: &rusqlite::Connection,
table_prefix: &'static str,
shard_root: Address,
shard_root_addr: Address,
) -> Result<Option<LocatedPrunableTree<H>>, 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<u8>>(0),
named_params![":shard_index": shard_root_addr.index()],
|row| Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Option<Vec<u8>>>(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<H>],
) -> Result<(), ShardTreeError<Error>> {
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<u8> = 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<u8> = 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<SqliteShardStore<rusqlite::Connection, String, 3>, 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::<Vec<_>>();
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",
"________________________________"
]
);
}
}