From 966f29f40102a96ed5126598f584ac11aeb6164a Mon Sep 17 00:00:00 2001 From: QuantumExplorer Date: Thu, 20 Jun 2024 20:27:51 +0300 Subject: [PATCH] fix: sum tree verification with specialized balances (#1899) --- .../src/balances/total_credits_balance/mod.rs | 20 +++++++++++++++---- .../tests/strategy_tests/main.rs | 6 +++--- .../calculate_total_credits_balance/v0/mod.rs | 10 ++++++++++ .../document/document_create_transition.rs | 8 ++++++++ 4 files changed, 37 insertions(+), 7 deletions(-) diff --git a/packages/rs-dpp/src/balances/total_credits_balance/mod.rs b/packages/rs-dpp/src/balances/total_credits_balance/mod.rs index 5ba240e7d8..58d131a678 100644 --- a/packages/rs-dpp/src/balances/total_credits_balance/mod.rs +++ b/packages/rs-dpp/src/balances/total_credits_balance/mod.rs @@ -11,6 +11,8 @@ pub struct TotalCreditsBalance { pub total_in_pools: SignedCredits, /// all the credits in identity balances pub total_identity_balances: SignedCredits, + /// all the credits in specialized balances + pub total_specialized_balances: SignedCredits, } impl TotalCreditsBalance { @@ -21,6 +23,7 @@ impl TotalCreditsBalance { total_credits_in_platform, total_in_pools, total_identity_balances, + total_specialized_balances, } = *self; if total_in_pools < 0 { @@ -35,6 +38,12 @@ impl TotalCreditsBalance { )); } + if total_specialized_balances < 0 { + return Err(ProtocolError::CriticalCorruptedCreditsCodeExecution( + "Credits of specialized balances are less than 0".to_string(), + )); + } + if total_credits_in_platform > MAX_CREDITS { return Err(ProtocolError::CriticalCorruptedCreditsCodeExecution( "Total credits in platform more than max credits size".to_string(), @@ -43,6 +52,7 @@ impl TotalCreditsBalance { let total_from_trees = (total_in_pools) .checked_add(total_identity_balances) + .and_then(|partial_sum| partial_sum.checked_add(total_specialized_balances)) .ok_or(ProtocolError::CriticalCorruptedCreditsCodeExecution( "Overflow of total credits".to_string(), ))?; @@ -55,14 +65,16 @@ impl TotalCreditsBalance { let TotalCreditsBalance { total_in_pools, total_identity_balances, + total_specialized_balances, .. } = *self; - let total_in_trees = total_in_pools.checked_add(total_identity_balances).ok_or( - ProtocolError::CriticalCorruptedCreditsCodeExecution( + let total_in_trees = total_in_pools + .checked_add(total_identity_balances) + .and_then(|partial_sum| partial_sum.checked_add(total_specialized_balances)) + .ok_or(ProtocolError::CriticalCorruptedCreditsCodeExecution( "Overflow of total credits".to_string(), - ), - )?; + ))?; Ok(total_in_trees.to_unsigned()) } diff --git a/packages/rs-drive-abci/tests/strategy_tests/main.rs b/packages/rs-drive-abci/tests/strategy_tests/main.rs index e2ac288985..b7976a10ca 100644 --- a/packages/rs-drive-abci/tests/strategy_tests/main.rs +++ b/packages/rs-drive-abci/tests/strategy_tests/main.rs @@ -1089,7 +1089,7 @@ mod tests { .unwrap() .unwrap() ), - "838bf7225bdd30f5ac86d876f879b5e6721d594a9072560b27b9e33614bf5bf3".to_string() + "6a444db28358f329be0a1e57e5382a17ba052c89fd5e7be33e0888f2ed0891d5".to_string() ) } @@ -1779,7 +1779,7 @@ mod tests { .unwrap() .unwrap() ), - "6e158be6c6752fc5afe2b840627c5a025eaebf10ff4fa1e1a9694ad451e60f99".to_string() + "9c9d4576485ce813c2f823159fa5dfc8242a63c7cb232f6ed8be3106ff85fc84".to_string() ) } @@ -1904,7 +1904,7 @@ mod tests { .unwrap() .unwrap() ), - "67abea32d6a5c6d39b7d1fb2a2bc58c8d95c32458923669913bcba6b30ab9ae3".to_string() + "7c5b874af2a5026de28d72642963fc3524d1597ebbc0bab3f16337c6eb29f8e9".to_string() ) } diff --git a/packages/rs-drive/src/drive/balances/calculate_total_credits_balance/v0/mod.rs b/packages/rs-drive/src/drive/balances/calculate_total_credits_balance/v0/mod.rs index b60c19c507..f70b2e2cc2 100644 --- a/packages/rs-drive/src/drive/balances/calculate_total_credits_balance/v0/mod.rs +++ b/packages/rs-drive/src/drive/balances/calculate_total_credits_balance/v0/mod.rs @@ -42,6 +42,15 @@ impl Drive { drive_version, )?; + let total_specialized_balances = self.grove_get_sum_tree_total_value( + SubtreePath::empty(), + Into::<&[u8; 1]>::into(RootTree::PreFundedSpecializedBalances), + DirectQueryType::StatefulDirectQuery, + transaction, + &mut drive_operations, + drive_version, + )?; + let total_in_pools = self.grove_get_sum_tree_total_value( SubtreePath::empty(), Into::<&[u8; 1]>::into(RootTree::Pools), @@ -55,6 +64,7 @@ impl Drive { total_credits_in_platform, total_in_pools, total_identity_balances, + total_specialized_balances, }) } } diff --git a/packages/rs-drive/src/drive/batch/transitions/document/document_create_transition.rs b/packages/rs-drive/src/drive/batch/transitions/document/document_create_transition.rs index cfdb190263..672b2192a0 100644 --- a/packages/rs-drive/src/drive/batch/transitions/document/document_create_transition.rs +++ b/packages/rs-drive/src/drive/batch/transitions/document/document_create_transition.rs @@ -64,6 +64,14 @@ impl DriveHighLevelDocumentOperationConverter for DocumentCreateTransitionAction }, )); + // We remove from the identity balance an equal amount + ops.push(IdentityOperation( + IdentityOperationType::RemoveFromIdentityBalance { + identity_id: owner_id.into_buffer(), + balance_to_remove: credits, + }, + )); + // We add the contested document // The contested document resides in a special location in grovedb until a time where the // resolution expires, at that point it either will be moved to