From afc133b3200feb2f0750266a0beb86f464b85d0d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20Chabowski?= <rafal.chabowski@fuel.sh>
Date: Tue, 14 Jan 2025 15:35:17 +0100
Subject: [PATCH] Update `balances()` complexity to support queries from SDK

---
 crates/fuel-core/src/schema/balance.rs | 26 +++++++--
 tests/tests/dos.rs                     | 73 ++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 5 deletions(-)

diff --git a/crates/fuel-core/src/schema/balance.rs b/crates/fuel-core/src/schema/balance.rs
index 92d0820d24b..45fd4c0041d 100644
--- a/crates/fuel-core/src/schema/balance.rs
+++ b/crates/fuel-core/src/schema/balance.rs
@@ -79,11 +79,27 @@ impl BalanceQuery {
         Ok(balance)
     }
 
-    #[graphql(
-        complexity = "query_costs().balance_query + query_costs().storage_iterator \
-        + (query_costs().storage_read * first.unwrap_or_default() as usize) * child_complexity \
-        + (query_costs().storage_read * last.unwrap_or_default() as usize) * child_complexity"
-    )]
+    // TODO: https://github.com/FuelLabs/fuel-core/issues/2496
+    // This is the complexity we want to use with "balances()" query, but it's not
+    // currently possible, because we need to handle queries with ~10k items.
+    // We use the temporary complexity until the SDKs are updated to not
+    // request such a large number of items.
+    // #[graphql(
+    // complexity = "query_costs().balance_query + query_costs().storage_iterator \
+    // + (query_costs().storage_read * first.unwrap_or_default() as usize) \
+    // + (child_complexity * first.unwrap_or_default() as usize) \
+    // + (query_costs().storage_read * last.unwrap_or_default() as usize) \
+    // + (child_complexity * last.unwrap_or_default() as usize)"
+    // )]
+
+    // The 0.66 factor is a Goldilocks approach to balancing the complexity of the query for the SDKs.
+    // Rust SDK sends a query with child_complexity ≅ 11 and we want to support slightly more
+    // than 10k items in a single query (so we target 11k). The total complexity would be 11k * 11 = 121k,
+    // but since our default limit is 80k, we need the 0.66 factor.
+    #[graphql(complexity = "
+        (child_complexity as f32 * first.unwrap_or_default() as f32 * 0.66) as usize + \
+        (child_complexity as f32 * last.unwrap_or_default() as f32 * 0.66) as usize
+    ")]
     async fn balances(
         &self,
         ctx: &Context<'_>,
diff --git a/tests/tests/dos.rs b/tests/tests/dos.rs
index 03092699ef4..dccfb92e691 100644
--- a/tests/tests/dos.rs
+++ b/tests/tests/dos.rs
@@ -11,6 +11,7 @@ use fuel_core::service::{
 };
 use fuel_core_client::client::FuelClient;
 use fuel_core_types::blockchain::header::LATEST_STATE_TRANSITION_VERSION;
+use test_case::test_case;
 use test_helpers::send_graph_ql_query;
 
 #[tokio::test]
@@ -720,3 +721,75 @@ async fn heavy_tasks_doesnt_block_graphql() {
     let health = result.expect("Health check failed");
     assert!(health);
 }
+
+const BALANCES_QUERY_RS_SDK: &str = r#"
+      query {
+        balances(
+          filter: {
+            owner: "6b63804cfbf9856e68e5b6e7aef238dc8311ec55bec04df774003a2c96e0418e"
+          }
+        $COUNT
+        ) {
+          edges {
+            node {
+              owner
+              amount
+              assetId
+            }
+            cursor
+          }
+          pageInfo {
+            hasNextPage
+            hasPreviousPage
+            startCursor
+            endCursor
+          }
+        }
+      }
+    "#;
+
+// Currently, the TS SDK does not query for `cursor` and `owner` fields, which means
+// it's a little less complex. That's why we test the DoS against the more heavy
+// query issues by the Rust SDK and this one is currently unused.
+const _BALANCES_QUERY_TS_SDK: &str = r#"
+      query {
+        balances(
+          filter: {
+            owner: "6b63804cfbf9856e68e5b6e7aef238dc8311ec55bec04df774003a2c96e0418e"
+          }
+        $COUNT
+        ) {
+          edges {
+            node {
+              amount
+              assetId
+            }
+          }
+          pageInfo {
+            hasNextPage
+            hasPreviousPage
+            startCursor
+            endCursor
+          }
+        }
+      }
+    "#;
+
+#[test_case("first: 11000", "balances"; "should handle 11000, forward")]
+#[test_case("last: 11000", "balances"; "should handle 11000, backward")]
+#[test_case("first: 11500", "Query is too complex"; "should bail with 11500, forward")]
+#[test_case("last: 11500", "Query is too complex"; "should bail with 11500, backward")]
+#[tokio::test]
+async fn balances_complexity_handles_amount_queried_by_sdk(
+    count: &str,
+    result_substring: &str,
+) {
+    let query = BALANCES_QUERY_RS_SDK;
+    let query = query.replace("$COUNT", count);
+
+    let node = FuelService::new_node(Config::local_node()).await.unwrap();
+    let url = format!("http://{}/v1/graphql", node.bound_address);
+
+    let result = send_graph_ql_query(&url, &query).await;
+    assert!(result.contains(result_substring), "{:?}", result);
+}