Set Custody Count Correctly (#14004)

* Set Custody Count Correctly * Fix Discovery Count
prysmaticlabs · Oct 7, 2024 · 59e571c · 59e571c
1 parent 278a2f6
commit 59e571c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git a/beacon-chain/p2p/discovery.go b/beacon-chain/p2p/discovery.go
@@ -14,6 +14,7 @@ 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"
@@ -44,14 +45,14 @@ const (
 
 type (
 	quicProtocol       uint16
-	custodySubnetCount uint64
+	CustodySubnetCount []byte
 )
 
 // quicProtocol is the "quic" key, which holds the QUIC port of the node.
 func (quicProtocol) ENRKey() string { return "quic" }
 
 // 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 "custody_subnet_count" }
 
 // 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
@@ -269,10 +270,14 @@ func (s *Service) createLocalNode(
 	}
 
 	if features.Get().EnablePeerDAS {
-		custodySubnetEntry := custodySubnetCount(params.BeaconConfig().CustodyRequirement)
+		var custodyBytes []byte
+		custodyBytes = ssz.MarshalUint64(custodyBytes, params.BeaconConfig().CustodyRequirement)
+		custodySubnetEntry := CustodySubnetCount(custodyBytes)
 
 		if flags.Get().SubscribeToAllSubnets {
-			custodySubnetEntry = custodySubnetCount(params.BeaconConfig().DataColumnSidecarSubnetCount)
+			var allCustodyBytes []byte
+			allCustodyBytes = ssz.MarshalUint64(allCustodyBytes, params.BeaconConfig().DataColumnSidecarSubnetCount)
+			custodySubnetEntry = CustodySubnetCount(allCustodyBytes)
 		}
 
 		localNode.Set(custodySubnetEntry)

diff --git a/beacon-chain/sync/BUILD.bazel b/beacon-chain/sync/BUILD.bazel
@@ -132,7 +132,6 @@ go_library(
         "@com_github_ethereum_go_ethereum//common/hexutil:go_default_library",
         "@com_github_ethereum_go_ethereum//common/math:go_default_library",
         "@com_github_ethereum_go_ethereum//crypto:go_default_library",
-        "@com_github_ethereum_go_ethereum//p2p/enr:go_default_library",
         "@com_github_hashicorp_golang_lru//:go_default_library",
         "@com_github_libp2p_go_libp2p//core:go_default_library",
         "@com_github_libp2p_go_libp2p//core/host:go_default_library",

diff --git a/beacon-chain/sync/sampling_data_columns.go b/beacon-chain/sync/sampling_data_columns.go
@@ -7,11 +7,12 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/pkg/errors"
+	ssz "github.com/prysmaticlabs/fastssz"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/feed"
 	statefeed "github.com/prysmaticlabs/prysm/v5/beacon-chain/core/feed/state"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
+	"github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p/types"
 	fieldparams "github.com/prysmaticlabs/prysm/v5/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
@@ -57,18 +58,18 @@ func (s *Service) sampleDataColumns(requestedRoot [fieldparams.RootLength]byte,
 		}
 
 		peerCustodiedSubnetCount := params.BeaconConfig().CustodyRequirement
-		peerActualCustodiedSubnetCount := new(uint16)
 
 		if peerRecord != nil {
 			// Load the `custody_subnet_count`
 			// TODO: Do not harcode `custody_subnet_count`
-			entry := enr.WithEntry("custody_subnet_count", peerActualCustodiedSubnetCount)
-			if err := peerRecord.Load(entry); err != nil {
+			custodyBytes := make([]byte, 8)
+			if err := peerRecord.Load(p2p.CustodySubnetCount(custodyBytes)); err != nil {
 				return nil, errors.Wrap(err, "load custody_subnet_count")
 			}
+			actualCustodyCount := ssz.UnmarshallUint64(custodyBytes)
 
-			if uint64(*peerActualCustodiedSubnetCount) > peerCustodiedSubnetCount {
-				peerCustodiedSubnetCount = uint64(*peerActualCustodiedSubnetCount)
+			if actualCustodyCount > peerCustodiedSubnetCount {
+				peerCustodiedSubnetCount = actualCustodyCount
 			}
 		}