Remove subnet filter from Peer.TrackedSubnets()

network/metrics.go

@@ -12,11 +12,12 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/network/peer"
 	"github.com/ava-labs/avalanchego/utils"
-	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/set"
 )
 
 type metrics struct {
+	trackedSubnets set.Set[ids.ID]
+
 	numTracked                      prometheus.Gauge
 	numPeers                        prometheus.Gauge
 	numSubnetPeers                  *prometheus.GaugeVec
@@ -41,8 +42,9 @@ type metrics struct {
 	peerConnectedStartTimesSum float64
 }
 
-func newMetrics(namespace string, registerer prometheus.Registerer, initialSubnetIDs set.Set[ids.ID]) (*metrics, error) {
+func newMetrics(namespace string, registerer prometheus.Registerer, trackedSubnets set.Set[ids.ID]) (*metrics, error) {
 	m := &metrics{
+		trackedSubnets: trackedSubnets,
 		numPeers: prometheus.NewGauge(prometheus.GaugeOpts{
 			Namespace: namespace,
 			Name:      "peers",
@@ -169,11 +171,7 @@ func newMetrics(namespace string, registerer prometheus.Registerer, initialSubne
 	)
 
 	// init subnet tracker metrics with tracked subnets
-	for subnetID := range initialSubnetIDs {
-		// no need to track primary network ID
-		if subnetID == constants.PrimaryNetworkID {
-			continue
-		}
+	for subnetID := range trackedSubnets {
 		// initialize to 0
 		subnetIDStr := subnetID.String()
 		m.numSubnetPeers.WithLabelValues(subnetIDStr).Set(0)
@@ -189,8 +187,10 @@ func (m *metrics) markConnected(peer peer.Peer) {
 	m.connected.Inc()
 
 	trackedSubnets := peer.TrackedSubnets()
-	for subnetID := range trackedSubnets {
-		m.numSubnetPeers.WithLabelValues(subnetID.String()).Inc()
+	for subnetID := range m.trackedSubnets {
+		if trackedSubnets.Contains(subnetID) {
+			m.numSubnetPeers.WithLabelValues(subnetID.String()).Inc()
+		}
 	}
 
 	m.lock.Lock()
@@ -206,8 +206,10 @@ func (m *metrics) markDisconnected(peer peer.Peer) {
 	m.disconnected.Inc()
 
 	trackedSubnets := peer.TrackedSubnets()
-	for subnetID := range trackedSubnets {
-		m.numSubnetPeers.WithLabelValues(subnetID.String()).Dec()
+	for subnetID := range m.trackedSubnets {
+		if trackedSubnets.Contains(subnetID) {
+			m.numSubnetPeers.WithLabelValues(subnetID.String()).Dec()
+		}
 	}
 
 	m.lock.Lock()

network/network.go

@@ -460,8 +460,12 @@ func (n *network) Connected(nodeID ids.NodeID) {
 
 	peerVersion := peer.Version()
 	n.router.Connected(nodeID, peerVersion, constants.PrimaryNetworkID)
-	for subnetID := range peer.TrackedSubnets() {
-		n.router.Connected(nodeID, peerVersion, subnetID)
+
+	trackedSubnets := peer.TrackedSubnets()
+	for subnetID := range n.peerConfig.MySubnets {
+		if trackedSubnets.Contains(subnetID) {
+			n.router.Connected(nodeID, peerVersion, subnetID)
+		}
 	}
 }
 

network/peer/peer.go

@@ -31,9 +31,14 @@ import (
 	"github.com/ava-labs/avalanchego/version"
 )
 
-// maxBloomSaltLen restricts the allowed size of the bloom salt to prevent
-// excessively expensive bloom filter contains checks.
-const maxBloomSaltLen = 32
+const (
+	// maxBloomSaltLen restricts the allowed size of the bloom salt to prevent
+	// excessively expensive bloom filter contains checks.
+	maxBloomSaltLen = 32
+	// maxNumTrackedSubnets limits how many subnets a peer can track to prevent
+	// excessive memory usage.
+	maxNumTrackedSubnets = 16
+)
 
 var (
 	errClosed = errors.New("closed")
@@ -131,8 +136,8 @@ type peer struct {
 	// version is the claimed version the peer is running that we received in
 	// the Handshake message.
 	version *version.Application
-	// trackedSubnets is the subset of subnetIDs the peer sent us in the Handshake
-	// message that we are also tracking.
+	// trackedSubnets are the subnetIDs the peer sent us in the Handshake
+	// message.
 	trackedSubnets set.Set[ids.ID]
 	// options of ACPs provided in the Handshake message.
 	supportedACPs set.Set[uint32]
@@ -263,9 +268,8 @@ func (p *peer) Info() Info {
 		publicIPStr = p.ip.IPPort.String()
 	}
 
-	uptimes := make(map[ids.ID]json.Uint32, p.trackedSubnets.Len())
-
-	for subnetID := range p.trackedSubnets {
+	uptimes := make(map[ids.ID]json.Uint32, p.MySubnets.Len())
+	for subnetID := range p.MySubnets {
 		uptime, exist := p.ObservedUptime(subnetID)
 		if !exist {
 			continue
@@ -802,8 +806,12 @@ func (p *peer) getUptimes() (uint32, []*p2p.SubnetUptime) {
 		primaryUptime = 0
 	}
 
-	subnetUptimes := make([]*p2p.SubnetUptime, 0, p.trackedSubnets.Len())
-	for subnetID := range p.trackedSubnets {
+	subnetUptimes := make([]*p2p.SubnetUptime, 0, p.MySubnets.Len())
+	for subnetID := range p.MySubnets {
+		if !p.trackedSubnets.Contains(subnetID) {
+			continue
+		}
+
 		subnetUptime, err := p.UptimeCalculator.CalculateUptimePercent(p.id, subnetID)
 		if err != nil {
 			p.Log.Debug("failed to get peer uptime percentage",
@@ -948,6 +956,17 @@ func (p *peer) handleHandshake(msg *p2p.Handshake) {
 	}
 
 	// handle subnet IDs
+	if numTrackedSubnets := len(msg.TrackedSubnets); numTrackedSubnets > maxNumTrackedSubnets {
+		p.Log.Debug("message with invalid field",
+			zap.Stringer("nodeID", p.id),
+			zap.Stringer("messageOp", message.HandshakeOp),
+			zap.String("field", "TrackedSubnets"),
+			zap.Int("numTrackedSubnets", numTrackedSubnets),
+		)
+		p.StartClose()
+		return
+	}
+
 	for _, subnetIDBytes := range msg.TrackedSubnets {
 		subnetID, err := ids.ToID(subnetIDBytes)
 		if err != nil {
@@ -958,8 +977,7 @@ func (p *peer) handleHandshake(msg *p2p.Handshake) {
 			p.StartClose()
 			return
 		}
-		// add only if we also track this subnet
-		if p.MySubnets.Contains(subnetID) {
+		if subnetID != constants.PrimaryNetworkID {
 			p.trackedSubnets.Add(subnetID)
 		}
 	}