From dde973e2d432482414ac7565618f912a99feccc5 Mon Sep 17 00:00:00 2001
From: tersec <tersec@users.noreply.github.com>
Date: Tue, 12 Jan 2021 13:43:15 +0100
Subject: [PATCH] allow always-on subscription to all attestation subnets when
 gossiping (#2225)

* allow always-on subscription to all attestation subnets when gossiping

* in subscribe-all-subnets mode, consider all subnets to be stability subnets for ENR purposes
---
 beacon_chain/conf.nim               |  5 +++++
 beacon_chain/nimbus_beacon_node.nim | 29 ++++++++++++++++++++---------
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/beacon_chain/conf.nim b/beacon_chain/conf.nim
index 67c0da260..6c8dd085a 100644
--- a/beacon_chain/conf.nim
+++ b/beacon_chain/conf.nim
@@ -115,6 +115,11 @@ type
       desc: "Node agent string which is used as identifier in network"
       name: "agent-string" }: string
 
+    subscribeAllSubnets* {.
+      defaultValue: false,
+      desc: "Subscribe to all attestation subnet topics when gossiping"
+      name: "subscribe-all-subnets" }: bool
+
     case cmd* {.
       command
       defaultValue: noCommand }: BNStartUpCmd
diff --git a/beacon_chain/nimbus_beacon_node.nim b/beacon_chain/nimbus_beacon_node.nim
index a4f1e6051..14d5773d6 100644
--- a/beacon_chain/nimbus_beacon_node.nim
+++ b/beacon_chain/nimbus_beacon_node.nim
@@ -377,6 +377,7 @@ func getStabilitySubnets(stabilitySubnets: auto): set[uint8] =
 
 proc cycleAttestationSubnets(node: BeaconNode, wallSlot: Slot) =
   static: doAssert RANDOM_SUBNETS_PER_VALIDATOR == 1
+  doAssert not node.config.subscribeAllSubnets
 
   # Only know RANDAO mix, which determines shuffling seed, one epoch in
   # advance. When node.chainDag.headState.data.data.slot.epoch is ahead
@@ -434,13 +435,20 @@ proc cycleAttestationSubnets(node: BeaconNode, wallSlot: Slot) =
 
   node.installAttestationSubnetHandlers(newSubnets)
 
-  let stabilitySubnets =
-    getStabilitySubnets(node.attestationSubnets.stabilitySubnets)
-  if stabilitySubnets != prevStabilitySubnets:
-    node.updateStabilitySubnetMetadata(stabilitySubnets)
+  if not node.config.subscribeAllSubnets:
+    # In subscribeAllSubnets mode, this only gets set once, at initial subnet
+    # attestation handler creation, since they're all considered as stability
+    # subnets in that case.
+    let stabilitySubnets =
+      getStabilitySubnets(node.attestationSubnets.stabilitySubnets)
+    if stabilitySubnets != prevStabilitySubnets:
+      node.updateStabilitySubnetMetadata(stabilitySubnets)
 
 proc getAttestationSubnetHandlers(node: BeaconNode) =
   var initialSubnets: set[uint8]
+  # If/when this stops subscribing to all attestation subnet handlers, the
+  # subscribeAllSubnets implementation needs modification from its current
+  # no-op/exploit-defaults version.
   for i in 0'u8 ..< ATTESTATION_SUBNET_COUNT:
     initialSubnets.incl i
 
@@ -458,7 +466,10 @@ proc getAttestationSubnetHandlers(node: BeaconNode) =
       expiration: wallEpoch + getStabilitySubnetLength())
 
   node.updateStabilitySubnetMetadata(
-    node.attestationSubnets.stabilitySubnets.getStabilitySubnets)
+    if node.config.subscribeAllSubnets:
+      initialSubnets
+    else:
+      node.attestationSubnets.stabilitySubnets.getStabilitySubnets)
 
   # Sets the "current" and "future" attestation subnets. One of these gets
   # replaced by get_attestation_subnet_changes() immediately. Symmetric so
@@ -474,7 +485,6 @@ proc getAttestationSubnetHandlers(node: BeaconNode) =
   node.installAttestationSubnetHandlers(initialSubnets)
 
 proc addMessageHandlers(node: BeaconNode) =
-  # As a side-effect, this gets the attestation subnets too.
   node.network.subscribe(node.topicBeaconBlocks, enableTopicMetrics = true)
   node.network.subscribe(getAttesterSlashingsTopic(node.forkDigest))
   node.network.subscribe(getProposerSlashingsTopic(node.forkDigest))
@@ -482,7 +492,6 @@ proc addMessageHandlers(node: BeaconNode) =
   node.network.subscribe(getAggregateAndProofsTopic(node.forkDigest), enableTopicMetrics = true)
   node.getAttestationSubnetHandlers()
 
-
 func getTopicSubscriptionEnabled(node: BeaconNode): bool =
   node.attestationSubnets.enabled
 
@@ -548,8 +557,10 @@ proc updateGossipStatus(node: BeaconNode, slot: Slot) =
       syncQueueLen
     node.removeMessageHandlers()
 
-  # Subscription or unsubscription might have occurred; recheck.
-  if node.getTopicSubscriptionEnabled:
+  # Subscription or unsubscription might have occurred; recheck. Since Nimbus
+  # initially subscribes to all subnets, simply do not ever cycle attestation
+  # subnets and they'll all remain subscribed.
+  if node.getTopicSubscriptionEnabled and not node.config.subscribeAllSubnets:
     # This exits early all but one call each epoch.
     node.cycleAttestationSubnets(slot)