fix: sum tree verification with specialized balances (#1899)

dashpay · Jun 20, 2024 · 966f29f · 966f29f
1 parent a75ca37
commit 966f29f
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 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
@@ -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
@@ -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
@@ -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
@@ -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