Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(state): Limit the retrieval of note commitment subtrees #7734

Merged
merged 10 commits into from
Oct 12, 2023
4 changes: 2 additions & 2 deletions zebra-chain/src/amount.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ pub type Result<T, E = Error> = std::result::Result<T, E>;
// TODO:
// - remove the default NegativeAllowed bound, to make consensus rule reviews easier
// - put a Constraint bound on the type generic, not just some implementations
#[derive(Clone, Copy, Serialize, Deserialize)]
#[derive(Clone, Copy, Serialize, Deserialize, Default)]
#[serde(try_from = "i64")]
#[serde(into = "i64")]
#[serde(bound = "C: Constraint + Clone")]
Expand Down Expand Up @@ -509,7 +509,7 @@ impl Constraint for NegativeAllowed {
/// 0..=MAX_MONEY,
/// );
/// ```
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Default)]
pub struct NonNegative;

impl Constraint for NonNegative {
Expand Down
2 changes: 1 addition & 1 deletion zebra-chain/src/orchard/tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ impl ZcashDeserialize for Root {
}

/// A node of the Orchard Incremental Note Commitment Tree.
#[derive(Copy, Clone, Eq, PartialEq)]
#[derive(Copy, Clone, Eq, PartialEq, Default)]
pub struct Node(pallas::Base);

impl Node {
Expand Down
2 changes: 1 addition & 1 deletion zebra-chain/src/sapling/tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ impl ZcashDeserialize for Root {
///
/// Note that it's handled as a byte buffer and not a point coordinate (jubjub::Fq)
/// because that's how the spec handles the MerkleCRH^Sapling function inputs and outputs.
#[derive(Copy, Clone, Eq, PartialEq)]
#[derive(Copy, Clone, Eq, PartialEq, Default)]
pub struct Node([u8; 32]);

impl AsRef<[u8; 32]> for Node {
Expand Down
2 changes: 1 addition & 1 deletion zebra-chain/src/value_balance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ mod tests;
use ValueBalanceError::*;

/// An amount spread between different Zcash pools.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)]
pub struct ValueBalance<C> {
transparent: Amount<C>,
sprout: Amount<C>,
Expand Down
8 changes: 2 additions & 6 deletions zebra-rpc/src/methods/tests/snapshot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -347,9 +347,7 @@ async fn test_mocked_rpc_response_data_for_network(network: Network) {

// Mock the data for the response.
let mut subtrees = BTreeMap::new();
let subtree_root = [0u8; 32].as_slice().try_into().expect(
"The array [0u8; 32] should be convertible to a Sapling note commitment tree node.",
);
let subtree_root = sapling::tree::Node::default();

for i in 0..2u16 {
let subtree = NoteCommitmentSubtreeData::new(Height(i.into()), subtree_root);
Expand Down Expand Up @@ -377,9 +375,7 @@ async fn test_mocked_rpc_response_data_for_network(network: Network) {

// Mock the data for the response.
let mut subtrees = BTreeMap::new();
let subtree_root = [0u8; 32].try_into().expect(
"The array [0u8; 32] should be convertible to a Sapling note commitment tree node.",
);
let subtree_root = orchard::tree::Node::default();

for i in 0..2u16 {
let subtree = NoteCommitmentSubtreeData::new(Height(i.into()), subtree_root);
Expand Down
2 changes: 1 addition & 1 deletion zebra-state/src/service/finalized_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ pub use disk_format::{OutputIndex, OutputLocation, TransactionLocation, MAX_ON_D

pub(super) use zebra_db::ZebraDb;

use disk_db::DiskWriteBatch;
pub(super) use disk_db::DiskWriteBatch;

/// The finalized part of the chain state, stored in the db.
///
Expand Down
16 changes: 14 additions & 2 deletions zebra-state/src/service/finalized_state/zebra_db/shielded.rs
Original file line number Diff line number Diff line change
Expand Up @@ -295,6 +295,12 @@ impl ZebraDb {
.collect();
}

// Make sure the amount of retrieved subtrees does not exceed the given limit.
#[cfg(debug_assertions)]
if let Some(limit) = limit {
assert!(list.len() <= limit.0.into());
}

// Check that we got the start subtree.
if list.get(&start_index).is_some() {
list
Expand Down Expand Up @@ -462,6 +468,12 @@ impl ZebraDb {
.collect();
}

// Make sure the amount of retrieved subtrees does not exceed the given limit.
#[cfg(debug_assertions)]
if let Some(limit) = limit {
assert!(list.len() <= limit.0.into());
}

// Check that we got the start subtree.
if list.get(&start_index).is_some() {
list
Expand Down Expand Up @@ -645,7 +657,7 @@ impl DiskWriteBatch {

// Sapling tree methods

/// Inserts the Sapling note commitment subtree.
/// Inserts the Sapling note commitment subtree into the batch.
pub fn insert_sapling_subtree(
&mut self,
zebra_db: &ZebraDb,
Expand Down Expand Up @@ -698,7 +710,7 @@ impl DiskWriteBatch {

// Orchard tree methods

/// Inserts the Orchard note commitment subtree.
/// Inserts the Orchard note commitment subtree into the batch.
pub fn insert_orchard_subtree(
&mut self,
zebra_db: &ZebraDb,
Expand Down
27 changes: 25 additions & 2 deletions zebra-state/src/service/non_finalized_state/chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ pub mod index;

/// A single non-finalized partial chain, from the child of the finalized tip,
/// to a non-finalized chain tip.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, Default)]
pub struct Chain {
// Config
//
Expand Down Expand Up @@ -67,7 +67,7 @@ pub struct Chain {
}

/// The internal state of [`Chain`].
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct ChainInner {
// Blocks, heights, hashes, and transaction locations
//
Expand Down Expand Up @@ -2224,3 +2224,26 @@ impl PartialEq for Chain {
}

impl Eq for Chain {}

#[cfg(test)]
impl Chain {
/// Inserts the supplied Sapling note commitment subtree into the chain.
pub(crate) fn insert_sapling_subtree(
&mut self,
subtree: NoteCommitmentSubtree<sapling::tree::Node>,
) {
self.inner
.sapling_subtrees
.insert(subtree.index, subtree.into_data());
}

/// Inserts the supplied Orchard note commitment subtree into the chain.
pub(crate) fn insert_orchard_subtree(
&mut self,
subtree: NoteCommitmentSubtree<orchard::tree::Node>,
) {
self.inner
.orchard_subtrees
.insert(subtree.index, subtree.into_data());
}
}
156 changes: 155 additions & 1 deletion zebra-state/src/service/read/tests/vectors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,11 @@ use std::sync::Arc;

use zebra_chain::{
block::{Block, Height},
orchard,
parameters::Network::*,
sapling,
serialization::ZcashDeserializeInto,
subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
transaction,
};

Expand All @@ -14,7 +17,16 @@ use zebra_test::{
transcript::{ExpectedTranscriptError, Transcript},
};

use crate::{init_test_services, populated_state, response::MinedTx, ReadRequest, ReadResponse};
use crate::{
init_test_services, populated_state,
response::MinedTx,
service::{
finalized_state::{DiskWriteBatch, ZebraDb},
non_finalized_state::Chain,
read::{orchard_subtrees, sapling_subtrees},
},
Config, ReadRequest, ReadResponse,
};

/// Test that ReadStateService responds correctly when empty.
#[tokio::test]
Expand Down Expand Up @@ -88,6 +100,136 @@ async fn populated_read_state_responds_correctly() -> Result<()> {
Ok(())
}

/// Tests if Zebra combines the Sapling note commitment subtrees from the finalized and
/// non-finalized states correctly.
#[tokio::test]
async fn test_sapling_subtrees() -> Result<()> {
let dummy_subtree_root = sapling::tree::Node::default();

// Prepare the finalized state.
let db_subtree = NoteCommitmentSubtree::new(0, Height(1), dummy_subtree_root);
let db = ZebraDb::new(&Config::ephemeral(), Mainnet, true);
let mut db_batch = DiskWriteBatch::new();
db_batch.insert_sapling_subtree(&db, &db_subtree);
db.write(db_batch)
.expect("Writing a batch with a Sapling subtree should succeed.");

// Prepare the non-fianlized state.
let chain_subtree = NoteCommitmentSubtree::new(1, Height(3), dummy_subtree_root);
let mut chain = Chain::default();
chain.insert_sapling_subtree(chain_subtree);
let chain = Some(Arc::new(chain));

// At this point, we have one Sapling subtree in the finalized state and one Sapling subtree in
// the non-finalized state.

// Retrieve only the first subtree and check its properties.
let subtrees = sapling_subtrees(chain.clone(), &db, 0.into(), Some(1.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree));

// Retrieve both subtrees using a limit and check their properties.
let subtrees = sapling_subtrees(chain.clone(), &db, 0.into(), Some(2.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 2);
assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree));
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve both subtrees without using a limit and check their properties.
let subtrees = sapling_subtrees(chain.clone(), &db, 0.into(), None);
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 2);
assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree));
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve only the second subtree and check its properties.
let subtrees = sapling_subtrees(chain.clone(), &db, 1.into(), Some(1.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve only the second subtree, using a limit that would allow for more trees if they were
// present, and check its properties.
let subtrees = sapling_subtrees(chain.clone(), &db, 1.into(), Some(2.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve only the second subtree, without using any limit, and check its properties.
let subtrees = sapling_subtrees(chain, &db, 1.into(), None);
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

Ok(())
}

/// Tests if Zebra combines the Orchard note commitment subtrees from the finalized and
/// non-finalized states correctly.
#[tokio::test]
async fn test_orchard_subtrees() -> Result<()> {
let dummy_subtree_root = orchard::tree::Node::default();

// Prepare the finalized state.
let db_subtree = NoteCommitmentSubtree::new(0, Height(1), dummy_subtree_root);
let db = ZebraDb::new(&Config::ephemeral(), Mainnet, true);
let mut db_batch = DiskWriteBatch::new();
db_batch.insert_orchard_subtree(&db, &db_subtree);
db.write(db_batch)
.expect("Writing a batch with an Orchard subtree should succeed.");

// Prepare the non-fianlized state.
let chain_subtree = NoteCommitmentSubtree::new(1, Height(3), dummy_subtree_root);
let mut chain = Chain::default();
chain.insert_orchard_subtree(chain_subtree);
let chain = Some(Arc::new(chain));

// At this point, we have one Orchard subtree in the finalized state and one Orchard subtree in
// the non-finalized state.

// Retrieve only the first subtree and check its properties.
let subtrees = orchard_subtrees(chain.clone(), &db, 0.into(), Some(1.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree));

// Retrieve both subtrees using a limit and check their properties.
let subtrees = orchard_subtrees(chain.clone(), &db, 0.into(), Some(2.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 2);
assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree));
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve both subtrees without using a limit and check their properties.
let subtrees = orchard_subtrees(chain.clone(), &db, 0.into(), None);
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 2);
assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree));
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve only the second subtree and check its properties.
let subtrees = orchard_subtrees(chain.clone(), &db, 1.into(), Some(1.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve only the second subtree, using a limit that would allow for more trees if they were
// present, and check its properties.
let subtrees = orchard_subtrees(chain.clone(), &db, 1.into(), Some(2.into()));
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

// Retrieve only the second subtree, without using any limit, and check its properties.
let subtrees = orchard_subtrees(chain, &db, 1.into(), None);
let mut subtrees = subtrees.iter();
assert_eq!(subtrees.len(), 1);
assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree));

Ok(())
}

/// Returns test cases for the empty state and missing blocks.
fn empty_state_test_cases() -> Vec<(ReadRequest, Result<ReadResponse, ExpectedTranscriptError>)> {
let block: Arc<Block> = zebra_test::vectors::BLOCK_MAINNET_419200_BYTES
Expand All @@ -109,3 +251,15 @@ fn empty_state_test_cases() -> Vec<(ReadRequest, Result<ReadResponse, ExpectedTr
),
]
}

/// Returns `true` if `index` and `subtree_data` match the contents of `subtree`. Otherwise, returns
/// `false`.
fn subtrees_eq<N>(
(index, subtree_data): (&NoteCommitmentSubtreeIndex, &NoteCommitmentSubtreeData<N>),
subtree: &NoteCommitmentSubtree<N>,
) -> bool
where
N: PartialEq + Copy,
{
index == &subtree.index && subtree_data == &subtree.into_data()
}
Loading
Loading