From 4bb2fac98f81033f7c901fd78ce4fde55a68e8ec Mon Sep 17 00:00:00 2001
From: Bo QIU <35757521+boqiu@users.noreply.github.com>
Date: Thu, 5 Dec 2024 14:04:57 +0800
Subject: [PATCH] Optional subscribe FIND_CHUNKS topic (#291)

* sub annouce shard config msg

* Optional subscribe find_chunks topic
---
 node/network/src/config.rs       |  4 ++++
 node/network/src/service.rs      | 13 ++++++++++++-
 node/network/src/types/mod.rs    |  2 +-
 node/network/src/types/topics.rs |  8 --------
 node/src/config/convert.rs       |  1 +
 node/src/config/mod.rs           |  1 +
 tests/sync_test.py               |  6 ++++++
 7 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/node/network/src/config.rs b/node/network/src/config.rs
index d2eb392..616bd8b 100644
--- a/node/network/src/config.rs
+++ b/node/network/src/config.rs
@@ -133,6 +133,9 @@ pub struct Config {
     /// Whether to disable network identity in ENR.
     /// This is for test purpose only.
     pub disable_enr_network_id: bool,
+
+    /// Whether to allow find chunks from peers.
+    pub find_chunks_enabled: bool,
 }
 
 impl Default for Config {
@@ -214,6 +217,7 @@ impl Default for Config {
             peer_db: Default::default(),
             peer_manager: Default::default(),
             disable_enr_network_id: false,
+            find_chunks_enabled: false,
         }
     }
 }
diff --git a/node/network/src/service.rs b/node/network/src/service.rs
index e6b3812..3de768e 100644
--- a/node/network/src/service.rs
+++ b/node/network/src/service.rs
@@ -243,7 +243,18 @@ impl<AppReqId: ReqId> Service<AppReqId> {
         //     }
         // }
 
-        for topic_kind in &crate::types::CORE_TOPICS {
+        let mut topics = vec![
+            GossipKind::NewFile,
+            GossipKind::FindFile,
+            GossipKind::AnnounceFile,
+            GossipKind::AnnounceShardConfig,
+        ];
+        if config.find_chunks_enabled {
+            topics.push(GossipKind::FindChunks);
+            topics.push(GossipKind::AnnounceChunks);
+        }
+
+        for topic_kind in topics {
             if swarm.behaviour_mut().subscribe_kind(topic_kind.clone()) {
                 subscribed_topics.push(topic_kind.clone());
             } else {
diff --git a/node/network/src/types/mod.rs b/node/network/src/types/mod.rs
index c8f93a7..b93cf42 100644
--- a/node/network/src/types/mod.rs
+++ b/node/network/src/types/mod.rs
@@ -11,4 +11,4 @@ pub use pubsub::{
     PubsubMessage, SignedAnnounceChunks, SignedAnnounceFile, SignedAnnounceShardConfig,
     SignedMessage, SnappyTransform,
 };
-pub use topics::{GossipEncoding, GossipKind, GossipTopic, CORE_TOPICS};
+pub use topics::{GossipEncoding, GossipKind, GossipTopic};
diff --git a/node/network/src/types/topics.rs b/node/network/src/types/topics.rs
index a4fc80b..ae087a8 100644
--- a/node/network/src/types/topics.rs
+++ b/node/network/src/types/topics.rs
@@ -15,14 +15,6 @@ pub const ANNOUNCE_FILE_TOPIC: &str = "announce_file";
 pub const ANNOUNCE_CHUNKS_TOPIC: &str = "announce_chunks";
 pub const ANNOUNCE_SHARD_CONFIG_TOPIC: &str = "announce_shard_config";
 
-pub const CORE_TOPICS: [GossipKind; 5] = [
-    GossipKind::NewFile,
-    GossipKind::FindFile,
-    GossipKind::FindChunks,
-    GossipKind::AnnounceFile,
-    GossipKind::AnnounceChunks,
-];
-
 /// A gossipsub topic which encapsulates the type of messages that should be sent and received over
 /// the pubsub protocol and the way the messages should be encoded.
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
diff --git a/node/src/config/convert.rs b/node/src/config/convert.rs
index 660fdfd..0ba2a1e 100644
--- a/node/src/config/convert.rs
+++ b/node/src/config/convert.rs
@@ -111,6 +111,7 @@ impl ZgsConfig {
         network_config.peer_db = self.network_peer_db;
         network_config.peer_manager = self.network_peer_manager.clone();
         network_config.disable_enr_network_id = self.discv5_disable_enr_network_id;
+        network_config.find_chunks_enabled = self.network_find_chunks_enabled;
 
         Ok(network_config)
     }
diff --git a/node/src/config/mod.rs b/node/src/config/mod.rs
index 8e9183b..fccc7f5 100644
--- a/node/src/config/mod.rs
+++ b/node/src/config/mod.rs
@@ -19,6 +19,7 @@ build_config! {
     (network_libp2p_nodes, (Vec<String>), vec![])
     (network_private, (bool), false)
     (network_disable_discovery, (bool), false)
+    (network_find_chunks_enabled, (bool), false)
 
     // discv5
     (discv5_request_timeout_secs, (u64), 5)
diff --git a/tests/sync_test.py b/tests/sync_test.py
index a07f923..45f5bb5 100755
--- a/tests/sync_test.py
+++ b/tests/sync_test.py
@@ -14,6 +14,12 @@ class SyncTest(TestFramework):
     def setup_params(self):
         self.num_nodes = 2
 
+        # enable find chunks topic
+        for i in range(self.num_nodes):
+            self.zgs_node_configs[i] = {
+                "network_find_chunks_enabled": True
+            }
+
     def run_test(self):
         # By default, auto_sync_enabled and sync_file_on_announcement_enabled are both false,
         # and file or chunks sync should be triggered by rpc.