From b9d3bf7b56878d8f4ee82875882a39117965a796 Mon Sep 17 00:00:00 2001
From: QuantumExplorer <quantum@dash.org>
Date: Thu, 19 Sep 2024 16:53:28 +0700
Subject: [PATCH] fix(platform): add limits to identity balances and proposed
 block counts queries (#2148)

---
 packages/rs-drive-abci/src/error/query.rs                 | 4 ++++
 .../proposed_block_counts_by_evonode_ids/v0/mod.rs        | 8 ++++++++
 packages/rs-drive/src/drive/identity/fetch/mod.rs         | 2 +-
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/packages/rs-drive-abci/src/error/query.rs b/packages/rs-drive-abci/src/error/query.rs
index 0a6e931703..ba60b25379 100644
--- a/packages/rs-drive-abci/src/error/query.rs
+++ b/packages/rs-drive-abci/src/error/query.rs
@@ -40,6 +40,10 @@ pub enum QueryError {
     #[error("invalid argument error: {0}")]
     InvalidArgument(String),
 
+    /// Too many elements Error
+    #[error("too many elements error: {0}")]
+    TooManyElements(String),
+
     /// Not found Error
     #[error("not found error: {0}")]
     NotFound(String),
diff --git a/packages/rs-drive-abci/src/query/validator_queries/proposed_block_counts_by_evonode_ids/v0/mod.rs b/packages/rs-drive-abci/src/query/validator_queries/proposed_block_counts_by_evonode_ids/v0/mod.rs
index 28ccfbbe9d..6780b1a690 100644
--- a/packages/rs-drive-abci/src/query/validator_queries/proposed_block_counts_by_evonode_ids/v0/mod.rs
+++ b/packages/rs-drive-abci/src/query/validator_queries/proposed_block_counts_by_evonode_ids/v0/mod.rs
@@ -21,6 +21,14 @@ impl<C> Platform<C> {
         platform_state: &PlatformState,
         platform_version: &PlatformVersion,
     ) -> Result<QueryValidationResult<GetEvonodesProposedEpochBlocksResponseV0>, Error> {
+        if ids.len() > self.config.drive.max_query_limit as usize {
+            return Ok(QueryValidationResult::new_with_error(
+                QueryError::TooManyElements(format!(
+                    "this query only supports up to {} ids at a time",
+                    ids.len()
+                )),
+            ));
+        }
         let evonode_ids = check_validation_result_with_data!(ids
             .into_iter()
             .map(|evonode_id_vec| {
diff --git a/packages/rs-drive/src/drive/identity/fetch/mod.rs b/packages/rs-drive/src/drive/identity/fetch/mod.rs
index 78bdd9c90f..79270a3569 100644
--- a/packages/rs-drive/src/drive/identity/fetch/mod.rs
+++ b/packages/rs-drive/src/drive/identity/fetch/mod.rs
@@ -180,7 +180,7 @@ impl Drive {
             path: vec![vec![RootTree::Balances as u8]],
             query: SizedQuery {
                 query,
-                limit: None,
+                limit: Some(self.config.max_query_limit),
                 offset: None,
             },
         };