[PeerDAS] rework ENR custody_subnet_count and add tests (#14077)

* [PeerDAS] rework ENR custody_subnet_count related code * update according to proposed spec change * Run gazelle
2026-01-10 05:47:59 -05:00 · 2024-06-12 13:35:12 -07:00
parent ce4b6b5464
commit 2863dcb9ea
5 changed files with 32 additions and 23 deletions
--- a/beacon-chain/p2p/BUILD.bazel
+++ b/beacon-chain/p2p/BUILD.bazel
@@ -82,7 +82,6 @@ go_library(
        "@com_github_ethereum_go_ethereum//p2p/discover:go_default_library",
        "@com_github_ethereum_go_ethereum//p2p/enode:go_default_library",
        "@com_github_ethereum_go_ethereum//p2p/enr:go_default_library",
-        "@com_github_ferranbt_fastssz//:go_default_library",
        "@com_github_holiman_uint256//:go_default_library",
        "@com_github_kr_pretty//:go_default_library",
        "@com_github_libp2p_go_libp2p//:go_default_library",
@@ -155,6 +154,7 @@ go_test(
        "//beacon-chain/p2p/types:go_default_library",
        "//beacon-chain/startup:go_default_library",
        "//cmd/beacon-chain/flags:go_default_library",
+        "//config/features:go_default_library",
        "//config/fieldparams:go_default_library",
        "//config/params:go_default_library",
        "//consensus-types/primitives:go_default_library",
--- a/beacon-chain/p2p/custody.go
+++ b/beacon-chain/p2p/custody.go
@@ -1,13 +1,13 @@
 package p2p

 import (
-	ssz "github.com/ferranbt/fastssz"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
 	"github.com/prysmaticlabs/prysm/v5/cmd/beacon-chain/flags"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
-	"github.com/sirupsen/logrus"
 )

 func (s *Service) GetValidCustodyPeers(peers []peer.ID) ([]peer.ID, error) {
@@ -66,18 +66,16 @@ func (s *Service) CustodyCountFromRemotePeer(pid peer.ID) uint64 {
 	}

 	// Load the `custody_subnet_count`
-	custodyObj := CustodySubnetCount(make([]byte, 8))
-	if err := peerRecord.Load(&custodyObj); err != nil {
+	var csc CustodySubnetCount
+	if err := peerRecord.Load(&csc); err != nil {
 		log.WithField("peerID", pid).Error("Cannot load the custody_subnet_count from peer")
 		return peerCustodyCountCount
 	}

-	// Unmarshal the custody count from the peer's ENR.
-	peerCustodyCountFromRecord := ssz.UnmarshallUint64(custodyObj)
 	log.WithFields(logrus.Fields{
 		"peerID":       pid,
-		"custodyCount": peerCustodyCountFromRecord,
+		"custodyCount": csc,
 	}).Debug("Custody count read from peer's ENR")

-	return peerCustodyCountFromRecord
+	return uint64(csc)
 }
--- a/beacon-chain/p2p/discovery.go
+++ b/beacon-chain/p2p/discovery.go
@@ -14,8 +14,8 @@ import (
 	"github.com/libp2p/go-libp2p/core/peer"
 	ma "github.com/multiformats/go-multiaddr"
 	"github.com/pkg/errors"
-	ssz "github.com/prysmaticlabs/fastssz"
 	"github.com/prysmaticlabs/go-bitfield"
+
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/v5/cmd/beacon-chain/flags"
 	"github.com/prysmaticlabs/prysm/v5/config/features"
@@ -43,16 +43,21 @@ const (
 	udp6
 )

+const (
+	quickProtocolEnrKey      = "quic"
+	custodySubnetCountEnrKey = "csc"
+)
+
 type (
 	quicProtocol       uint16
-	CustodySubnetCount []byte
+	CustodySubnetCount uint64
 )

 // quicProtocol is the "quic" key, which holds the QUIC port of the node.
-func (quicProtocol) ENRKey() string { return "quic" }
+func (quicProtocol) ENRKey() string { return quickProtocolEnrKey }

 // https://github.com/ethereum/consensus-specs/blob/dev/specs/_features/eip7594/p2p-interface.md#the-discovery-domain-discv5
-func (CustodySubnetCount) ENRKey() string { return "custody_subnet_count" }
+func (CustodySubnetCount) ENRKey() string { return custodySubnetCountEnrKey }

 // RefreshPersistentSubnets checks that we are tracking our local persistent subnets for a variety of gossip topics.
 // This routine checks for our attestation, sync committee and data column subnets and updates them if they have
@@ -270,16 +275,10 @@ func (s *Service) createLocalNode(
 	}

 	if features.Get().EnablePeerDAS {
-		var custodyBytes []byte
-		custodyBytes = ssz.MarshalUint64(custodyBytes, params.BeaconConfig().CustodyRequirement)
-		custodySubnetEntry := CustodySubnetCount(custodyBytes)
-
+		custodySubnetEntry := CustodySubnetCount(params.BeaconConfig().CustodyRequirement)
 		if flags.Get().SubscribeToAllSubnets {
-			var allCustodyBytes []byte
-			allCustodyBytes = ssz.MarshalUint64(allCustodyBytes, params.BeaconConfig().DataColumnSidecarSubnetCount)
-			custodySubnetEntry = CustodySubnetCount(allCustodyBytes)
+			custodySubnetEntry = CustodySubnetCount(params.BeaconConfig().DataColumnSidecarSubnetCount)
 		}
-
 		localNode.Set(custodySubnetEntry)
 	}

--- a/beacon-chain/p2p/discovery_test.go
+++ b/beacon-chain/p2p/discovery_test.go
@@ -20,6 +20,8 @@ import (
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/prysmaticlabs/go-bitfield"
+	logTest "github.com/sirupsen/logrus/hooks/test"
+
 	mock "github.com/prysmaticlabs/prysm/v5/beacon-chain/blockchain/testing"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p/peers"
@@ -27,6 +29,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p/peers/scorers"
 	testp2p "github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p/testing"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/startup"
+	"github.com/prysmaticlabs/prysm/v5/config/features"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/wrapper"
 	leakybucket "github.com/prysmaticlabs/prysm/v5/container/leaky-bucket"
@@ -37,7 +40,6 @@ import (
 	"github.com/prysmaticlabs/prysm/v5/testing/assert"
 	"github.com/prysmaticlabs/prysm/v5/testing/require"
 	"github.com/prysmaticlabs/prysm/v5/time/slots"
-	logTest "github.com/sirupsen/logrus/hooks/test"
 )

 var discoveryWaitTime = 1 * time.Second
@@ -131,6 +133,11 @@ func TestStartDiscV5_DiscoverAllPeers(t *testing.T) {
 }

 func TestCreateLocalNode(t *testing.T) {
+	resetFn := features.InitWithReset(&features.Flags{
+		EnablePeerDAS: true,
+	})
+	defer resetFn()
+
 	testCases := []struct {
 		name          string
 		cfg           *Config
@@ -227,6 +234,11 @@ func TestCreateLocalNode(t *testing.T) {
 			syncSubnets := new([]byte)
 			require.NoError(t, localNode.Node().Record().Load(enr.WithEntry(syncCommsSubnetEnrKey, syncSubnets)))
 			require.DeepSSZEqual(t, []byte{0}, *syncSubnets)
+
+			// Check custody_subnet_count config.
+			custodySubnetCount := new(uint64)
+			require.NoError(t, localNode.Node().Record().Load(enr.WithEntry(custodySubnetCountEnrKey, custodySubnetCount)))
+			require.Equal(t, uint64(1), *custodySubnetCount)
 		})
 	}
 }
--- a/go.mod
+++ b/go.mod
@@ -19,7 +19,6 @@ require (
 	github.com/emicklei/dot v0.11.0
 	github.com/ethereum/c-kzg-4844 v1.0.2-0.20240507203752-26d3b4156f7a
 	github.com/ethereum/go-ethereum v1.13.5
-	github.com/ferranbt/fastssz v0.0.0-20210120143747-11b9eff30ea9
 	github.com/fjl/memsize v0.0.0-20190710130421-bcb5799ab5e5
 	github.com/fsnotify/fsnotify v1.6.0
 	github.com/ghodss/yaml v1.0.0
@@ -136,6 +135,7 @@ require (
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/dop251/goja v0.0.0-20230806174421-c933cf95e127 // indirect
 	github.com/elastic/gosigar v0.14.2 // indirect
+	github.com/ferranbt/fastssz v0.0.0-20210120143747-11b9eff30ea9 // indirect
 	github.com/flynn/noise v1.1.0 // indirect
 	github.com/francoispqt/gojay v1.2.13 // indirect
 	github.com/getsentry/sentry-go v0.25.0 // indirect