From 6a300f75af0a658467456d821c7733534a3c3573 Mon Sep 17 00:00:00 2001
From: Agnish Ghosh <iamagnix@gmail.com>
Date: Fri, 11 Oct 2024 00:54:59 +0530
Subject: [PATCH] include supernode and nodeid in dc quarantine

---
 .../data_column_quarantine.nim                | 40 +++++++++++++------
 beacon_chain/nimbus_beacon_node.nim           |  3 ++
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/beacon_chain/consensus_object_pools/data_column_quarantine.nim b/beacon_chain/consensus_object_pools/data_column_quarantine.nim
index 40a973f94..9914ebfb8 100644
--- a/beacon_chain/consensus_object_pools/data_column_quarantine.nim
+++ b/beacon_chain/consensus_object_pools/data_column_quarantine.nim
@@ -9,6 +9,7 @@
 
 import
   std/tables,
+  eth/p2p/discoveryv5/[node],
   ../spec/[helpers, eip7594_helpers]
 
 from std/sequtils import mapIt
@@ -23,6 +24,8 @@ type
   DataColumnQuarantine* = object
     data_columns*:
       OrderedTable[(Eth2Digest, ColumnIndex), ref DataColumnSidecar]
+    supernode*: bool
+    nodeid*: NodeId
     onDataColumnSidecarCallback*: OnDataColumnSidecarCallback
 
   DataColumnFetchRecord* = object
@@ -115,24 +118,35 @@ func checkForInitialDcSidecars*(quarantine: DataColumnQuarantine,
 
 func hasDataColumns*(quarantine: DataColumnQuarantine,
     blck: deneb.SignedBeaconBlock | electra.SignedBeaconBlock): bool =
-  var counter = 0
-  for i in 0..<NUMBER_OF_COLUMNS:
-    if len(blck.message.body.blob_kzg_commitments) != 0:
-      if (blck.root, ColumnIndex i) in quarantine.data_columns:
-        inc counter
-    else:
+  let
+    localSubnetCount = 
+      if quarantine.supernode:
+        DATA_COLUMN_SIDECAR_SUBNET_COUNT.uint64
+      else:
+        CUSTODY_REQUIREMENT.uint64
+    localCustodyColumns =
+      get_custody_columns(quarantine.nodeid,
+                          max(SAMPLES_PER_SLOT.uint64,
+                              localSubnetCount))
+  for i in localCustodyColumns:
+    if (blck.root, ColumnIndex i) notin quarantine.data_columns:
       return false
-  if counter == DATA_COLUMN_SIDECAR_SUBNET_COUNT or 
-      (counter == max(SAMPLES_PER_SLOT, CUSTODY_REQUIREMENT) and 
-      counter < max(SAMPLES_PER_SLOT, CUSTODY_REQUIREMENT) + 1):
-    return true
-  else:
-    false
+  true
 
 func dataColumnFetchRecord*(quarantine: DataColumnQuarantine,
     blck: deneb.SignedBeaconBlock | electra.SignedBeaconBlock): DataColumnFetchRecord =
+  let
+    localSubnetCount = 
+      if quarantine.supernode:
+        DATA_COLUMN_SIDECAR_SUBNET_COUNT.uint64
+      else:
+        CUSTODY_REQUIREMENT.uint64
+    localCustodyColumns =
+      get_custody_columns(quarantine.nodeid,
+                          max(SAMPLES_PER_SLOT.uint64,
+                              localSubnetCount))
   var indices: seq[ColumnIndex]
-  for i in 0..<NUMBER_OF_COLUMNS:
+  for i in localCustodyColumns:
     let idx = ColumnIndex(i)
     if not quarantine.data_columns.hasKey(
         (blck.root, idx)):
diff --git a/beacon_chain/nimbus_beacon_node.nim b/beacon_chain/nimbus_beacon_node.nim
index 2beada095..69dbded16 100644
--- a/beacon_chain/nimbus_beacon_node.nim
+++ b/beacon_chain/nimbus_beacon_node.nim
@@ -385,6 +385,9 @@ proc initFullNode(
       onProposerSlashingAdded, onAttesterSlashingAdded))
     blobQuarantine = newClone(BlobQuarantine.init(onBlobSidecarAdded))
     dataColumnQuarantine = newClone(DataColumnQuarantine.init())
+  dataColumnQuarantine[].supernode = supernode
+  dataColumnQuarantine[].nodeid = node.network.nodeId
+  let
     consensusManager = ConsensusManager.new(
       dag, attestationPool, quarantine, node.elManager,
       ActionTracker.init(node.network.nodeId, config.subscribeAllSubnets),