From 8a06cab9e43d803df1e860a9a675093ca3be579a Mon Sep 17 00:00:00 2001
From: quake <quake.wang@gmail.com>
Date: Tue, 10 Oct 2023 08:46:03 +0900
Subject: [PATCH] chore: only fetch blocks for filter matched

---
 .../components/send_blocks_proof.rs           |  4 ++--
 src/protocols/light_client/mod.rs             |  2 +-
 src/protocols/light_client/peers.rs           | 23 +++++++++++++++----
 .../light_client/send_blocks_proof.rs         |  8 +++----
 src/utils/network.rs                          |  2 +-
 5 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/protocols/light_client/components/send_blocks_proof.rs b/src/protocols/light_client/components/send_blocks_proof.rs
index 102fe96..f5caae2 100644
--- a/src/protocols/light_client/components/send_blocks_proof.rs
+++ b/src/protocols/light_client/components/send_blocks_proof.rs
@@ -34,7 +34,7 @@ impl<'a> SendBlocksProofProcess<'a> {
         let status = self.execute_internally();
         self.protocol
             .peers()
-            .update_blocks_proof_request(self.peer_index, None);
+            .update_blocks_proof_request(self.peer_index, None, false);
         status
     }
 
@@ -118,7 +118,7 @@ impl<'a> SendBlocksProofProcess<'a> {
             ));
 
             // Get blocks
-            {
+            if original_request.should_get_blocks() {
                 let block_hashes: Vec<packed::Byte32> =
                     headers.iter().map(|header| header.hash()).collect();
                 {
diff --git a/src/protocols/light_client/mod.rs b/src/protocols/light_client/mod.rs
index eaf0d66..830b646 100644
--- a/src/protocols/light_client/mod.rs
+++ b/src/protocols/light_client/mod.rs
@@ -750,7 +750,7 @@ impl LightClientProtocol {
                         .as_bytes();
 
                     self.peers
-                        .update_blocks_proof_request(*peer_index, Some(content));
+                        .update_blocks_proof_request(*peer_index, Some(content), false);
                     if let Err(err) = nc.send_message(
                         SupportProtocols::LightClient.protocol_id(),
                         *peer_index,
diff --git a/src/protocols/light_client/peers.rs b/src/protocols/light_client/peers.rs
index 261306c..8cc6072 100644
--- a/src/protocols/light_client/peers.rs
+++ b/src/protocols/light_client/peers.rs
@@ -145,6 +145,7 @@ pub(crate) struct ProveState {
 pub(crate) struct BlocksProofRequest {
     content: packed::GetBlocksProof,
     when_sent: u64,
+    should_get_blocks: bool,
 }
 
 #[derive(Clone)]
@@ -394,8 +395,16 @@ impl ProveState {
 }
 
 impl BlocksProofRequest {
-    pub(crate) fn new(content: packed::GetBlocksProof, when_sent: u64) -> Self {
-        Self { content, when_sent }
+    pub(crate) fn new(
+        content: packed::GetBlocksProof,
+        when_sent: u64,
+        should_get_blocks: bool,
+    ) -> Self {
+        Self {
+            content,
+            when_sent,
+            should_get_blocks,
+        }
     }
 
     pub(crate) fn last_hash(&self) -> Byte32 {
@@ -430,6 +439,10 @@ impl BlocksProofRequest {
             false
         }
     }
+
+    pub(crate) fn should_get_blocks(&self) -> bool {
+        self.should_get_blocks
+    }
 }
 
 impl BlocksRequest {
@@ -1480,10 +1493,12 @@ impl Peers {
         &self,
         index: PeerIndex,
         request: Option<packed::GetBlocksProof>,
+        should_get_blocks: bool,
     ) {
         if let Some(mut peer) = self.inner.get_mut(&index) {
-            peer.blocks_proof_request =
-                request.map(|content| BlocksProofRequest::new(content, unix_time_as_millis()));
+            peer.blocks_proof_request = request.map(|content| {
+                BlocksProofRequest::new(content, unix_time_as_millis(), should_get_blocks)
+            });
         }
     }
     pub(crate) fn update_blocks_request(&self, index: PeerIndex, hashes: Option<Vec<Byte32>>) {
diff --git a/src/tests/protocols/light_client/send_blocks_proof.rs b/src/tests/protocols/light_client/send_blocks_proof.rs
index 76de814..3ffde42 100644
--- a/src/tests/protocols/light_client/send_blocks_proof.rs
+++ b/src/tests/protocols/light_client/send_blocks_proof.rs
@@ -124,7 +124,7 @@ async fn last_state_is_changed() {
             .unwrap();
         protocol
             .peers()
-            .update_blocks_proof_request(peer_index, Some(content));
+            .update_blocks_proof_request(peer_index, Some(content), true);
     }
 
     num += 1;
@@ -218,7 +218,7 @@ async fn unexpected_response() {
             .unwrap();
         protocol
             .peers()
-            .update_blocks_proof_request(peer_index, Some(content));
+            .update_blocks_proof_request(peer_index, Some(content), true);
     }
 
     // Run the test.
@@ -325,7 +325,7 @@ async fn get_blocks_with_chunks() {
             .unwrap();
         protocol
             .peers()
-            .update_blocks_proof_request(peer_index, Some(content));
+            .update_blocks_proof_request(peer_index, Some(content), true);
     }
 
     // Run the test.
@@ -675,7 +675,7 @@ async fn test_send_blocks_proof(param: TestParameter) {
             .unwrap();
         protocol
             .peers()
-            .update_blocks_proof_request(peer_index, Some(content));
+            .update_blocks_proof_request(peer_index, Some(content), true);
     }
 
     // Run the test.
diff --git a/src/utils/network.rs b/src/utils/network.rs
index 661dee2..10da7bc 100644
--- a/src/utils/network.rs
+++ b/src/utils/network.rs
@@ -45,7 +45,7 @@ pub(crate) fn prove_or_download_matched_blocks(
                     .set(content.clone())
                     .build()
                     .as_bytes();
-                peers.update_blocks_proof_request(*peer_index, Some(content));
+                peers.update_blocks_proof_request(*peer_index, Some(content), true);
                 if let Err(err) = nc.send_message(
                     SupportProtocols::LightClient.protocol_id(),
                     *peer_index,