diff --git a/zaino-state/src/chain_index.rs b/zaino-state/src/chain_index.rs index bd0217e89..522abc2cd 100644 --- a/zaino-state/src/chain_index.rs +++ b/zaino-state/src/chain_index.rs @@ -268,6 +268,14 @@ pub trait ChainIndex { hash: &types::BlockHash, ) -> impl std::future::Future>, Option>), Self::Error>>; + /// Returns the subtree roots + fn get_subtree_roots( + &self, + pool: ShieldedPool, + start_index: u16, + max_entries: Option, + ) -> impl std::future::Future, Self::Error>>; + /// given a transaction id, returns the transaction, along with /// its consensus branch ID if available #[allow(clippy::type_complexity)] @@ -1471,6 +1479,38 @@ impl ChainIndex for NodeBackedChainIndexSubscriber MempoolInfo { self.mempool.get_mempool_info().await } + + /// Gets the subtree roots of a given pool and the end heights of each root, + /// starting at the provided index, up to an optional maximum number of roots. + async fn get_subtree_roots( + &self, + pool: ShieldedPool, + start_index: u16, + max_entries: Option, + ) -> Result, Self::Error> { + self.source() + .get_subtree_roots(pool, start_index, max_entries) + .await + .map_err(ChainIndexError::backing_validator) + } +} + +/// The available shielded pools +#[non_exhaustive] +pub enum ShieldedPool { + /// Sapling + Sapling, + /// Orchard + Orchard, +} + +impl ShieldedPool { + fn pool_string(&self) -> String { + match self { + ShieldedPool::Sapling => "sapling".to_string(), + ShieldedPool::Orchard => "orchard".to_string(), + } + } } impl NonFinalizedSnapshot for Arc diff --git a/zaino-state/src/chain_index/source.rs b/zaino-state/src/chain_index/source.rs index e2864d003..cdf2c7908 100644 --- a/zaino-state/src/chain_index/source.rs +++ b/zaino-state/src/chain_index/source.rs @@ -2,7 +2,10 @@ use std::{error::Error, str::FromStr as _, sync::Arc}; -use crate::chain_index::types::{BlockHash, TransactionHash}; +use crate::chain_index::{ + types::{BlockHash, TransactionHash}, + ShieldedPool, +}; use async_trait::async_trait; use futures::{future::join, TryFutureExt as _}; use tower::{Service, ServiceExt as _}; @@ -12,7 +15,9 @@ use zaino_fetch::jsonrpsee::{ response::{GetBlockError, GetBlockResponse, GetTransactionResponse, GetTreestateResponse}, }; use zcash_primitives::merkle_tree::read_commitment_tree; -use zebra_chain::{block::TryIntoHeight, serialization::ZcashDeserialize}; +use zebra_chain::{ + block::TryIntoHeight, serialization::ZcashDeserialize, subtree::NoteCommitmentSubtreeIndex, +}; use zebra_state::{HashOrHeight, ReadRequest, ReadResponse, ReadStateService}; macro_rules! expected_read_response { @@ -85,6 +90,15 @@ pub trait BlockchainSource: Clone + Send + Sync + 'static { >, Box, >; + + /// Gets the subtree roots of a given pool and the end heights of each root, + /// starting at the provided index, up to an optional maximum number of roots. + async fn get_subtree_roots( + &self, + pool: ShieldedPool, + start_index: u16, + max_entries: Option, + ) -> BlockchainSourceResult>; } /// An error originating from a blockchain source. @@ -664,6 +678,80 @@ impl BlockchainSource for ValidatorConnector { ValidatorConnector::Fetch(_fetch) => Ok(None), } } + + async fn get_subtree_roots( + &self, + pool: ShieldedPool, + start_index: u16, + max_entries: Option, + ) -> BlockchainSourceResult> { + match self { + ValidatorConnector::State(state) => { + let start_index = NoteCommitmentSubtreeIndex(start_index); + let limit = max_entries.map(NoteCommitmentSubtreeIndex); + let request = match pool { + ShieldedPool::Sapling => ReadRequest::SaplingSubtrees { start_index, limit }, + ShieldedPool::Orchard => ReadRequest::OrchardSubtrees { start_index, limit }, + }; + state + .read_state_service + .clone() + .call(request) + .await + .map(|response| match pool { + ShieldedPool::Sapling => expected_read_response!(response, SaplingSubtrees) + .iter() + .map(|(_index, subtree)| { + (subtree.root.to_bytes(), subtree.end_height.0) + }) + .collect(), + ShieldedPool::Orchard => expected_read_response!(response, OrchardSubtrees) + .iter() + .map(|(_index, subtree)| (subtree.root.to_repr(), subtree.end_height.0)) + .collect(), + }) + .map_err(|e| { + BlockchainSourceError::Unrecoverable(format!( + "could not get subtrees from validator: {e}" + )) + }) + } + + ValidatorConnector::Fetch(json_rp_see_connector) => { + let subtrees = json_rp_see_connector + .get_subtrees_by_index(pool.pool_string(), start_index, max_entries) + .await + .map_err(|e| { + BlockchainSourceError::Unrecoverable(format!( + "could not get subtrees from validator: {e}" + )) + })?; + + Ok(subtrees + .subtrees + .iter() + .map(|subtree| { + Ok::<_, Box>(( + <[u8; 32]>::try_from(hex::decode(&subtree.root)?).map_err( + |_subtree| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "received subtree root not 32 bytes", + ) + }, + )?, + subtree.end_height.0, + )) + }) + .collect::, _>>() + .map_err(|e| { + BlockchainSourceError::Unrecoverable(format!( + "could not get subtrees from validator: {e}" + )) + })?) + } + } + } } /// The location of a transaction returned by @@ -987,5 +1075,14 @@ pub(crate) mod test { > { Ok(None) } + + async fn get_subtree_roots( + &self, + _pool: ShieldedPool, + _start_index: u16, + _max_entries: Option, + ) -> BlockchainSourceResult> { + todo!() + } } } diff --git a/zaino-state/src/chain_index/tests/proptest_blockgen.rs b/zaino-state/src/chain_index/tests/proptest_blockgen.rs index 4abc970c4..b26156e40 100644 --- a/zaino-state/src/chain_index/tests/proptest_blockgen.rs +++ b/zaino-state/src/chain_index/tests/proptest_blockgen.rs @@ -664,6 +664,15 @@ impl BlockchainSource for ProptestMockchain { .unwrap(); Ok(Some(receiver)) } + + async fn get_subtree_roots( + &self, + _pool: crate::chain_index::ShieldedPool, + _start_index: u16, + _max_entries: Option, + ) -> BlockchainSourceResult> { + todo!() + } } type ChainSegment = SummaryDebug>>;