diff --git a/packages/rs-drive-proof-verifier/src/proof.rs b/packages/rs-drive-proof-verifier/src/proof.rs index 0a83e9dec8..64b2ea20c9 100644 --- a/packages/rs-drive-proof-verifier/src/proof.rs +++ b/packages/rs-drive-proof-verifier/src/proof.rs @@ -1030,7 +1030,6 @@ impl FromProof &proof.grovedb_proof, start_block_height, limit, - false, platform_version, ) .map_drive_error(proof, mtd)?; diff --git a/packages/rs-drive/src/drive/saved_block_transactions/fetch_compacted_address_balances/v0/mod.rs b/packages/rs-drive/src/drive/saved_block_transactions/fetch_compacted_address_balances/v0/mod.rs index 18fd9db946..efa7838341 100644 --- a/packages/rs-drive/src/drive/saved_block_transactions/fetch_compacted_address_balances/v0/mod.rs +++ b/packages/rs-drive/src/drive/saved_block_transactions/fetch_compacted_address_balances/v0/mod.rs @@ -16,7 +16,10 @@ pub type CompactedAddressBalanceChanges = impl Drive { /// Version 0 implementation of fetching compacted address balance changes. /// - /// Retrieves all compacted address balance change records from `start_block_height` onwards. + /// Retrieves all compacted address balance change records where `end_block >= start_block_height`. + /// This includes ranges that contain `start_block_height` (e.g., range 400-600 when querying + /// from block 505) as well as ranges that start after `start_block_height`. + /// /// Returns a vector of (start_block, end_block, address_balance_map) tuples. pub(super) fn fetch_compacted_address_balance_changes_v0( &self, @@ -27,34 +30,118 @@ impl Drive { ) -> Result { let path = Self::saved_compacted_block_transactions_address_balances_path_vec(); - // Create a range query starting from the specified height - // Keys are 16 bytes: (start_block, end_block), both big-endian - // We query from (start_block_height, 0) onwards - let mut start_key = Vec::with_capacity(16); - start_key.extend_from_slice(&start_block_height.to_be_bytes()); - start_key.extend_from_slice(&0u64.to_be_bytes()); + // Keys are 16 bytes: (start_block, end_block), both big-endian. + // We want ranges where end_block >= start_block_height, which includes: + // - Ranges that contain start_block_height (e.g., 400-600 contains 505) + // - Ranges that start at or after start_block_height + // + // Strategy: + // 1. First query: descending from (start_block_height, u64::MAX) with limit 1 + // to find any range where start_block <= start_block_height that might contain it + // 2. Second query: ascending from (start_block_height, 0) to get ranges + // that start at or after start_block_height - let mut query = Query::new(); - query.insert_range_from(start_key..); + let config = bincode::config::standard() + .with_big_endian() + .with_no_limit(); - let path_query = PathQuery::new(path, SizedQuery::new(query, limit, None)); + let mut compacted_changes = Vec::new(); + let limit_usize = limit.map(|l| l as usize); - let (results, _) = self.grove_get_path_query( - &path_query, + // Query 1: Find if there's a range containing start_block_height + // Query descending from (start_block_height, u64::MAX) with limit 1 + let mut desc_end_key = Vec::with_capacity(16); + desc_end_key.extend_from_slice(&start_block_height.to_be_bytes()); + desc_end_key.extend_from_slice(&u64::MAX.to_be_bytes()); + + let mut desc_query = Query::new_with_direction(false); // descending + desc_query.insert_range_to_inclusive(..=desc_end_key); + + let desc_path_query = + PathQuery::new(path.clone(), SizedQuery::new(desc_query, Some(1), None)); + + let (desc_results, _) = self.grove_get_path_query( + &desc_path_query, transaction, QueryResultType::QueryKeyElementPairResultType, &mut vec![], &platform_version.drive, )?; - let config = bincode::config::standard() - .with_big_endian() - .with_no_limit(); + // Check if we found a range that contains start_block_height + if let Some((key, element)) = desc_results.to_key_elements().into_iter().next() { + if key.len() != 16 { + return Err(Error::Protocol(Box::new( + ProtocolError::CorruptedSerialization( + "invalid compacted block key length, expected 16 bytes".to_string(), + ), + ))); + } - let mut compacted_changes = Vec::new(); + let start_block = u64::from_be_bytes(key[0..8].try_into().unwrap()); + let end_block = u64::from_be_bytes(key[8..16].try_into().unwrap()); + + // Only include if end_block >= start_block_height (range contains our block) + if end_block >= start_block_height { + let Element::Item(serialized_data, _) = element else { + return Err(Error::Protocol(Box::new( + ProtocolError::CorruptedSerialization( + "expected item element for compacted address balances".to_string(), + ), + ))); + }; + + let (address_balances, _): (BTreeMap, usize) = + bincode::decode_from_slice(&serialized_data, config).map_err(|e| { + Error::Protocol(Box::new(ProtocolError::CorruptedSerialization(format!( + "cannot decode compacted address balances: {}", + e + )))) + })?; + + compacted_changes.push((start_block, end_block, address_balances)); + } + } + + // Check if we've already hit the limit + if let Some(l) = limit_usize { + if compacted_changes.len() >= l { + return Ok(compacted_changes); + } + } + + // Query 2: Get ranges that start at or after start_block_height (ascending) + // Always use (start_block_height, 0) for consistent proof verification + // The result may overlap with descending query if descending found a range + // starting exactly at start_block_height - we dedupe below + let mut asc_start_key = Vec::with_capacity(16); + asc_start_key.extend_from_slice(&start_block_height.to_be_bytes()); + asc_start_key.extend_from_slice(&0u64.to_be_bytes()); + + let mut asc_query = Query::new(); + asc_query.insert_range_from(asc_start_key..); + + let asc_path_query = PathQuery::new(path, SizedQuery::new(asc_query, limit, None)); + + let (asc_results, _) = self.grove_get_path_query( + &asc_path_query, + transaction, + QueryResultType::QueryKeyElementPairResultType, + &mut vec![], + &platform_version.drive, + )?; + + // Track the start_block from descending query to avoid duplicates + let desc_start_block = compacted_changes.first().map(|(start, _, _)| *start); + + for (key, element) in asc_results.to_key_elements() { + // Check if we've reached the limit + if let Some(l) = limit_usize { + if compacted_changes.len() >= l { + break; + } + } - for (key, element) in results.to_key_elements() { - // Parse start_block and end_block from key (16 bytes total, both big-endian) if key.len() != 16 { return Err(Error::Protocol(Box::new( ProtocolError::CorruptedSerialization( @@ -66,7 +153,11 @@ impl Drive { let start_block = u64::from_be_bytes(key[0..8].try_into().unwrap()); let end_block = u64::from_be_bytes(key[8..16].try_into().unwrap()); - // Get the serialized data from the Item element + // Skip if this is the same range we got from descending query + if Some(start_block) == desc_start_block { + continue; + } + let Element::Item(serialized_data, _) = element else { return Err(Error::Protocol(Box::new( ProtocolError::CorruptedSerialization( @@ -75,7 +166,6 @@ impl Drive { ))); }; - // Deserialize the address balance map let (address_balances, _): (BTreeMap, usize) = bincode::decode_from_slice(&serialized_data, config).map_err(|e| { Error::Protocol(Box::new(ProtocolError::CorruptedSerialization(format!( @@ -91,6 +181,12 @@ impl Drive { } /// Version 0 implementation for proving compacted address balance changes. + /// + /// Uses a two-step approach: + /// 1. First query (non-proving): descending to find any range containing start_block_height + /// 2. Second query (proving): ascending from the found start_block or start_block_height + /// + /// This ensures the proof covers all relevant ranges efficiently. pub(super) fn prove_compacted_address_balance_changes_v0( &self, start_block_height: u64, @@ -100,10 +196,55 @@ impl Drive { ) -> Result, Error> { let path = Self::saved_compacted_block_transactions_address_balances_path_vec(); - // Create a range query starting from the specified height - let mut start_key = Vec::with_capacity(16); - start_key.extend_from_slice(&start_block_height.to_be_bytes()); - start_key.extend_from_slice(&0u64.to_be_bytes()); + // Step 1: Non-proving descending query to find any range containing start_block_height + let mut desc_end_key = Vec::with_capacity(16); + desc_end_key.extend_from_slice(&start_block_height.to_be_bytes()); + desc_end_key.extend_from_slice(&u64::MAX.to_be_bytes()); + + let mut desc_query = Query::new_with_direction(false); // descending + desc_query.insert_range_to_inclusive(..=desc_end_key); + + let desc_path_query = + PathQuery::new(path.clone(), SizedQuery::new(desc_query, Some(1), None)); + + let (desc_results, _) = self.grove_get_path_query( + &desc_path_query, + transaction, + QueryResultType::QueryKeyElementPairResultType, + &mut vec![], + &platform_version.drive, + )?; + + // Determine the actual start key for the proved query + // If we found a containing range, use its exact key + // Otherwise use (start_block_height, start_block_height) since end_block >= start_block always + let start_key = if let Some((key, _)) = desc_results.to_key_elements().into_iter().next() { + if key.len() == 16 { + let end_block = u64::from_be_bytes(key[8..16].try_into().unwrap()); + // If this range contains start_block_height, use its exact key + if end_block >= start_block_height { + key + } else { + // No containing range, use (start_block_height, start_block_height) + let mut key = Vec::with_capacity(16); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key + } + } else { + let mut key = Vec::with_capacity(16); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key + } + } else { + let mut key = Vec::with_capacity(16); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key + }; + + // Step 2: Proved ascending query from start_key let mut query = Query::new(); query.insert_range_from(start_key..); diff --git a/packages/rs-drive/src/drive/saved_block_transactions/store_address_balances/v0/mod.rs b/packages/rs-drive/src/drive/saved_block_transactions/store_address_balances/v0/mod.rs index 2073854959..e4221e6e53 100644 --- a/packages/rs-drive/src/drive/saved_block_transactions/store_address_balances/v0/mod.rs +++ b/packages/rs-drive/src/drive/saved_block_transactions/store_address_balances/v0/mod.rs @@ -1090,4 +1090,132 @@ mod tests { panic!("expected Item element for block ranges"); } } + + #[test] + fn should_fetch_compacted_range_when_start_height_falls_within_range() { + let drive = setup_drive_with_initial_state_structure(None); + // Set compaction threshold to 4 blocks so we can trigger compaction easily + let platform_version = create_test_platform_version_for_compaction(4, 1000); + + // Store 4 blocks (100-103) to trigger compaction into range (100, 103) + for block_height in 100u64..=103 { + let mut balances = BTreeMap::new(); + balances.insert(ADDR_1, CreditOperation::AddToCredits(1000)); + + drive + .store_address_balances_for_block_v0( + &balances, + block_height, + 1700000000000, + None, + &platform_version, + ) + .expect("should store block"); + } + + // Verify compaction happened - fetch from block 0 should return the compacted range + let all_results = drive + .fetch_compacted_address_balance_changes(0, None, None, &platform_version) + .expect("should fetch compacted changes"); + + assert_eq!(all_results.len(), 1, "should have 1 compacted entry"); + let (start, end, _) = &all_results[0]; + assert_eq!(*start, 100, "compacted range should start at 100"); + assert_eq!(*end, 103, "compacted range should end at 103"); + + // Now query with start_height=101, which falls WITHIN the compacted range (100-103) + // We should still get this range since block 101 is within it + let results_from_101 = drive + .fetch_compacted_address_balance_changes(101, None, None, &platform_version) + .expect("should fetch compacted changes from 101"); + + assert_eq!( + results_from_101.len(), + 1, + "should include compacted range (100, 103) when querying from block 101 \ + since 101 falls within that range" + ); + } + + #[test] + fn should_prove_and_verify_compacted_address_balances() { + use crate::drive::Drive; + + let drive = setup_drive_with_initial_state_structure(None); + // Set compaction threshold to 4 blocks + let platform_version = create_test_platform_version_for_compaction(4, 1000); + + // Store 4 blocks (100-103) to trigger compaction into range (100, 103) + for block_height in 100u64..=103 { + let mut balances = BTreeMap::new(); + balances.insert(ADDR_1, CreditOperation::AddToCredits(1000)); + + drive + .store_address_balances_for_block_v0( + &balances, + block_height, + 1700000000000, + None, + &platform_version, + ) + .expect("should store block"); + } + + // Generate proof with limit (required for verify_query_with_absence_proof) + let limit = Some(100u16); + let proof = drive + .prove_compacted_address_balance_changes(0, limit, None, &platform_version) + .expect("should generate proof"); + + // Verify proof + let (root_hash, verified_changes) = + Drive::verify_compacted_address_balance_changes(&proof, 0, limit, &platform_version) + .expect("should verify proof"); + + assert!(!root_hash.is_empty(), "root hash should not be empty"); + assert_eq!(verified_changes.len(), 1, "should have 1 compacted entry"); + assert_eq!(verified_changes[0].0, 100, "start block should be 100"); + assert_eq!(verified_changes[0].1, 103, "end block should be 103"); + } + + #[test] + fn should_prove_and_verify_with_containing_range() { + use crate::drive::Drive; + + let drive = setup_drive_with_initial_state_structure(None); + // Set compaction threshold to 4 blocks + let platform_version = create_test_platform_version_for_compaction(4, 1000); + + // Store 4 blocks (100-103) to trigger compaction + for block_height in 100u64..=103 { + let mut balances = BTreeMap::new(); + balances.insert(ADDR_1, CreditOperation::AddToCredits(1000)); + + drive + .store_address_balances_for_block_v0( + &balances, + block_height, + 1700000000000, + None, + &platform_version, + ) + .expect("should store block"); + } + + // Generate proof starting from block 101 (which falls within range 100-103) + let limit = Some(100u16); + let proof = drive + .prove_compacted_address_balance_changes(101, limit, None, &platform_version) + .expect("should generate proof"); + + // Verify proof - should include range (100, 103) since 101 is within it + let (root_hash, verified_changes) = + Drive::verify_compacted_address_balance_changes(&proof, 101, limit, &platform_version) + .expect("should verify proof"); + + assert!(!root_hash.is_empty(), "root hash should not be empty"); + assert_eq!(verified_changes.len(), 1, "should have 1 compacted entry"); + assert_eq!(verified_changes[0].0, 100, "start block should be 100"); + assert_eq!(verified_changes[0].1, 103, "end block should be 103"); + } } diff --git a/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/mod.rs b/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/mod.rs index e7a940434a..fe7de6913d 100644 --- a/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/mod.rs +++ b/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/mod.rs @@ -24,7 +24,6 @@ impl Drive { /// - `proof`: A byte slice containing the cryptographic proof for the compacted address balance changes. /// - `start_block_height`: The block height to start verifying from. /// - `limit`: Optional maximum number of compacted entries to verify. - /// - `verify_subset_of_proof`: A boolean flag indicating whether to verify only a subset of the proof. /// - `platform_version`: A reference to the platform version. /// /// # Returns @@ -36,7 +35,6 @@ impl Drive { proof: &[u8], start_block_height: u64, limit: Option, - verify_subset_of_proof: bool, platform_version: &PlatformVersion, ) -> Result<(RootHash, VerifiedCompactedAddressBalanceChanges), Error> { match platform_version @@ -50,7 +48,6 @@ impl Drive { proof, start_block_height, limit, - verify_subset_of_proof, platform_version, ), version => Err(Error::Drive(DriveError::UnknownVersionMismatch { diff --git a/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/v0/mod.rs b/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/v0/mod.rs index 3cf1e05b5b..504bc49608 100644 --- a/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/v0/mod.rs +++ b/packages/rs-drive/src/verify/address_funds/verify_compacted_address_balance_changes/v0/mod.rs @@ -15,11 +15,17 @@ use std::collections::BTreeMap; use super::VerifiedCompactedAddressBalanceChanges; impl Drive { + /// Verifies compacted address balance changes proof. + /// + /// Verification is done in two steps: + /// 1. First verify as SUBSET to examine what's in the proof and determine + /// what actual_start_block was used when generating the proof. + /// 2. Then verify the main ascending query (not as subset) using the exact + /// same query that was used for proving. pub(super) fn verify_compacted_address_balance_changes_v0( proof: &[u8], start_block_height: u64, limit: Option, - verify_subset_of_proof: bool, platform_version: &PlatformVersion, ) -> Result<(RootHash, VerifiedCompactedAddressBalanceChanges), Error> { let path = vec![ @@ -27,35 +33,74 @@ impl Drive { vec![COMPACTED_ADDRESS_BALANCES_KEY_U8], ]; - // Create a range query starting from the specified height - // Keys are 16 bytes: (start_block, end_block), both big-endian - // We query from (start_block_height, 0) onwards - let mut start_key = Vec::with_capacity(16); - start_key.extend_from_slice(&start_block_height.to_be_bytes()); - start_key.extend_from_slice(&0u64.to_be_bytes()); + let config = bincode::config::standard() + .with_big_endian() + .with_no_limit(); + + // Step 1: Use subset query with insert_all() to examine everything in the proof + // This allows us to find entries that might start before start_block_height + // but contain it (e.g., range 100-200 when querying from block 150) + let mut subset_query = Query::new(); + subset_query.insert_all(); + + let subset_path_query = + PathQuery::new(path.clone(), SizedQuery::new(subset_query, None, None)); + + // Use verify_subset_query to look at what's in the proof + let (root_hash, subset_results) = GroveDb::verify_subset_query( + proof, + &subset_path_query, + &platform_version.drive.grove_version, + )?; + + // Collect results into a BTreeMap to ensure proper ordering by key + let results_map: BTreeMap, Element> = subset_results + .into_iter() + .filter_map(|(_, key, maybe_element)| maybe_element.map(|element| (key, element))) + .collect(); + + // Get the first entry and check if it's a containing range + // Only use the key from the proof if it contains start_block_height + // (start <= start_block_height <= end) + // Otherwise fall back to (start_block_height, start_block_height) + let start_key = results_map.first_key_value().and_then(|(key, _)| { + if key.len() != 16 { + return None; + } + let start_block = u64::from_be_bytes(key[0..8].try_into().unwrap()); + let end_block = u64::from_be_bytes(key[8..16].try_into().unwrap()); + + // Only return the key if it's a containing range + if start_block <= start_block_height && start_block_height <= end_block { + Some(key.clone()) + } else { + None + } + }); + + // Step 2: Verify the proof using the start_key discovered from the proof + // The smallest key in the proof is what the prove function used as its starting point + // If no entries exist, fall back to (start_block_height, start_block_height) + let start_key = start_key.unwrap_or_else(|| { + let mut key = Vec::with_capacity(16); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key.extend_from_slice(&start_block_height.to_be_bytes()); + key + }); let mut query = Query::new(); query.insert_range_from(start_key..); let path_query = PathQuery::new(path, SizedQuery::new(query, limit, None)); - - let (root_hash, proved_key_values) = if verify_subset_of_proof { - GroveDb::verify_subset_query_with_absence_proof( - proof, - &path_query, - &platform_version.drive.grove_version, - )? - } else { - GroveDb::verify_query_with_absence_proof( - proof, - &path_query, - &platform_version.drive.grove_version, - )? - }; - - let config = bincode::config::standard() - .with_big_endian() - .with_no_limit(); + let (verified_root_hash, proved_key_values) = + GroveDb::verify_query(proof, &path_query, &platform_version.drive.grove_version)?; + + // Both verifications must have the same root hash + if root_hash != verified_root_hash { + return Err(Error::Proof(ProofError::CorruptedProof( + "root hash mismatch between subset and main query verification".to_string(), + ))); + } let mut compacted_changes = Vec::new();