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

* Fix the combining of Sapling subtrees

This commit fixes a bug due to which the function `sapling_subtrees`
used to always include all subtrees from the non-finalized state without
respecting the given limit.

* Fix the combining of Orchard subtrees

This commit fixes a bug due to which the function `orchard_subtrees`
used to always include all subtrees from the non-finalized state without
respecting the given limit.

* Add additional checks when retrieving subtrees

* Allow raw subtree insertions into `Chain` in tests

* Derive `Default` for `Chain`

* Derive `Default` for note commitment tree nodes

* Use `pub(super)` for `DiskWriteBatch`

Set the visibility of `DiskWriteBatch` to `pub(super)`.

* Add tests for retrieving subtrees

* Use `default()` for subtree roots in tests

This change is unrelated to the PR.

* Refactor docs for inserting subtrees

This change is unrelated to the PR.

zebra-chain/src/amount.rs

@@ -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")]
@@ -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 {

zebra-chain/src/orchard/tree.rs

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

zebra-chain/src/sapling/tree.rs

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

zebra-chain/src/value_balance.rs

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

zebra-rpc/src/methods/tests/snapshot.rs

@@ -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(
+    let subtree_root = sapling::tree::Node::default();
-        "The array [0u8; 32] should be convertible to a Sapling note commitment tree node.",
-    );
 
     for i in 0..2u16 {
         let subtree = NoteCommitmentSubtreeData::new(Height(i.into()), subtree_root);
@@ -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(
+    let subtree_root = orchard::tree::Node::default();
-        "The array [0u8; 32] should be convertible to a Sapling note commitment tree node.",
-    );
 
     for i in 0..2u16 {
         let subtree = NoteCommitmentSubtreeData::new(Height(i.into()), subtree_root);

zebra-state/src/service/finalized_state.rs

@@ -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.
 ///

zebra-state/src/service/finalized_state/zebra_db/shielded.rs

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

zebra-state/src/service/non_finalized_state/chain.rs

@@ -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
     //
@@ -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
     //
@@ -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());
+    }
+}

zebra-state/src/service/read/tests/vectors.rs

@@ -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,
 };
 
@@ -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]
@@ -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
@@ -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()
+}

zebra-state/src/service/read/tree.rs

@@ -74,6 +74,19 @@ where
 {
     let mut db_list = db.sapling_subtree_list_by_index_for_rpc(start_index, limit);
 
+    if let Some(limit) = limit {
+        let subtrees_num = u16::try_from(db_list.len())
+            .expect("There can't be more than `u16::MAX` Sapling subtrees.");
+
+        // Return the subtrees if the amount of them reached the given limit.
+        if subtrees_num == limit.0 {
+            return db_list;
+        }
+
+        // If not, make sure the amount is below the limit.
+        debug_assert!(subtrees_num < limit.0);
+    }
+
     // If there's no chain, then we have the complete list.
     let Some(chain) = chain else {
         return db_list;
@@ -93,6 +106,17 @@ where
         // If there's no matching index, just update the list of trees.
         let Some(db_subtree) = db_list.get(&fork_index) else {
             db_list.insert(fork_index, fork_subtree);
+
+            // Stop adding new subtrees once their amount reaches the given limit.
+            if let Some(limit) = limit {
+                let subtrees_num = u16::try_from(db_list.len())
+                    .expect("There can't be more than `u16::MAX` Sapling subtrees.");
+
+                if subtrees_num == limit.0 {
+                    break;
+                }
+            }
+
             continue;
         };
 
@@ -104,6 +128,12 @@ where
         // Otherwise, the subtree is already in the list, so we don't need to add it.
     }
 
+    // Make sure the amount of retrieved subtrees does not exceed the given limit.
+    #[cfg(debug_assertions)]
+    if let Some(limit) = limit {
+        assert!(db_list.len() <= limit.0.into());
+    }
+
     // Check that we got the start subtree from the non-finalized or finalized state.
     // (The non-finalized state doesn't do this check.)
     if db_list.get(&start_index).is_some() {
@@ -160,6 +190,19 @@ where
 {
     let mut db_list = db.orchard_subtree_list_by_index_for_rpc(start_index, limit);
 
+    if let Some(limit) = limit {
+        let subtrees_num = u16::try_from(db_list.len())
+            .expect("There can't be more than `u16::MAX` Orchard subtrees.");
+
+        // Return the subtrees if the amount of them reached the given limit.
+        if subtrees_num == limit.0 {
+            return db_list;
+        }
+
+        // If not, make sure the amount is below the limit.
+        debug_assert!(subtrees_num < limit.0);
+    }
+
     // If there's no chain, then we have the complete list.
     let Some(chain) = chain else {
         return db_list;
@@ -179,6 +222,17 @@ where
         // If there's no matching index, just update the list of trees.
         let Some(db_subtree) = db_list.get(&fork_index) else {
             db_list.insert(fork_index, fork_subtree);
+
+            // Stop adding new subtrees once their amount reaches the given limit.
+            if let Some(limit) = limit {
+                let subtrees_num = u16::try_from(db_list.len())
+                    .expect("There can't be more than `u16::MAX` Orchard subtrees.");
+
+                if subtrees_num == limit.0 {
+                    break;
+                }
+            }
+
             continue;
         };
 
@@ -190,6 +244,12 @@ where
         // Otherwise, the subtree is already in the list, so we don't need to add it.
     }
 
+    // Make sure the amount of retrieved subtrees does not exceed the given limit.
+    #[cfg(debug_assertions)]
+    if let Some(limit) = limit {
+        assert!(db_list.len() <= limit.0.into());
+    }
+
     // Check that we got the start subtree from the non-finalized or finalized state.
     // (The non-finalized state doesn't do this check.)
     if db_list.get(&start_index).is_some() {