From e034f30b5e95c74c071531d24b5a723c67622f01 Mon Sep 17 00:00:00 2001
From: Agnish Ghosh <iamagnix@gmail.com>
Date: Mon, 15 Jul 2024 13:50:41 +0530
Subject: [PATCH] add: subscribeAllSubnets feature

---
 beacon_chain/nimbus_beacon_node.nim | 13 ++++++++-----
 beacon_chain/spec/network.nim       |  4 ++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/beacon_chain/nimbus_beacon_node.nim b/beacon_chain/nimbus_beacon_node.nim
index 3c36336a1..8eae6b049 100644
--- a/beacon_chain/nimbus_beacon_node.nim
+++ b/beacon_chain/nimbus_beacon_node.nim
@@ -1130,10 +1130,11 @@ proc addCapellaMessageHandlers(
   node.addAltairMessageHandlers(forkDigest, slot)
   node.network.subscribe(getBlsToExecutionChangeTopic(forkDigest), basicParams)
 
-proc fetchCustodySubnetCount* (node: BeaconNode, res: var uint64)=
-  res = CUSTODY_REQUIREMENT
+proc fetchCustodySubnetCount* (node: BeaconNode): uint64=
+  var res = CUSTODY_REQUIREMENT.uint64
   if node.config.subscribeAllSubnets:
-    res = DATA_COLUMN_SIDECAR_SUBNET_COUNT
+    res = DATA_COLUMN_SIDECAR_SUBNET_COUNT.uint64
+  res
 
 proc fetchCustodyColumnCountFromRemotePeer*(node: BeaconNode, pid: PeerId): uint64 =
   # Fetches the custody column count from a remote peer
@@ -1164,7 +1165,8 @@ proc fetchCustodyColumnCountFromRemotePeer*(node: BeaconNode, pid: PeerId): uint
 proc addDenebMessageHandlers(
     node: BeaconNode, forkDigest: ForkDigest, slot: Slot) =
   node.addCapellaMessageHandlers(forkDigest, slot)
-  for topic in dataColumnSidecarTopics(forkDigest):
+  let targetSubnets = node.fetchCustodySubnetCount()
+  for topic in dataColumnSidecarTopics(forkDigest, targetSubnets):
     node.network.subscribe(topic, basicParams)
 
 proc addElectraMessageHandlers(
@@ -1188,7 +1190,8 @@ proc removeCapellaMessageHandlers(node: BeaconNode, forkDigest: ForkDigest) =
 
 proc removeDenebMessageHandlers(node: BeaconNode, forkDigest: ForkDigest) =
   node.removeCapellaMessageHandlers(forkDigest)
-  for topic in dataColumnSidecarTopics(forkDigest):
+  let targetSubnets = node.fetchCustodySubnetCount()
+  for topic in dataColumnSidecarTopics(forkDigest, targetSubnets):
     node.network.unsubscribe(topic)
 
 proc removeElectraMessageHandlers(node: BeaconNode, forkDigest: ForkDigest) =
diff --git a/beacon_chain/spec/network.nim b/beacon_chain/spec/network.nim
index 9bc08c71e..69ec646a4 100644
--- a/beacon_chain/spec/network.nim
+++ b/beacon_chain/spec/network.nim
@@ -244,6 +244,6 @@ func getDataColumnSidecarTopic*(forkDigest: ForkDigest,
 func compute_subnet_for_data_column_sidecar*(column_index: ColumnIndex): uint64 =
     uint64(column_index mod DATA_COLUMN_SIDECAR_SUBNET_COUNT)
 
-iterator dataColumnSidecarTopics*(forkDigest: ForkDigest): string =
-  for subnet_id in 0'u64..<DATA_COLUMN_SIDECAR_SUBNET_COUNT:
+iterator dataColumnSidecarTopics*(forkDigest: ForkDigest, targetSubnetCount: uint64): string =
+  for subnet_id in 0'u64..<targetSubnetCount:
     yield getDataColumnSidecarTopic(forkDigest, subnet_id)
\ No newline at end of file