From c93c812bc1c3b2eb8e12948bc3d42a87640c9af0 Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Mon, 23 Sep 2024 15:36:34 -0400
Subject: [PATCH] fix: use byte array to decode ENRs uint8 fields

---
 waku/v2/discv5/discover.go          | 15 ++++-----------
 waku/v2/discv5/filters.go           | 12 +++---------
 waku/v2/peermanager/peer_manager.go |  5 ++---
 waku/v2/protocol/enr/enr.go         | 17 +++++++++++++++++
 4 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/waku/v2/discv5/discover.go b/waku/v2/discv5/discover.go
index 582e46d24..e762d07ad 100644
--- a/waku/v2/discv5/discover.go
+++ b/waku/v2/discv5/discover.go
@@ -21,7 +21,6 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/ethereum/go-ethereum/p2p/enode"
-	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/ethereum/go-ethereum/p2p/nat"
 )
 
@@ -253,19 +252,13 @@ func (d *DiscoveryV5) Stop() {
 }
 
 func isWakuNode(node *enode.Node) bool {
-	enrField := new(wenr.WakuEnrBitfield)
-	if err := node.Record().Load(enr.WithEntry(wenr.WakuENRField, &enrField)); err != nil {
-		if !enr.IsNotFound(err) {
-			utils.Logger().Named("discv5").Error("could not retrieve waku2 ENR field for enr ", zap.Any("node", node))
-		}
+	enrField, err := wenr.GetWakuEnrBitField(node)
+	if err != nil {
+		utils.Logger().Named("discv5").Error("could not retrieve waku2 ENR field for enr ", zap.Error(err))
 		return false
 	}
 
-	if enrField != nil {
-		return *enrField != uint8(0) // #RFC 31 requirement
-	}
-
-	return false
+	return enrField != uint8(0) // #RFC 31 requirement
 }
 
 func (d *DiscoveryV5) evaluateNode() func(node *enode.Node) bool {
diff --git a/waku/v2/discv5/filters.go b/waku/v2/discv5/filters.go
index e34502fa5..40557fbcc 100644
--- a/waku/v2/discv5/filters.go
+++ b/waku/v2/discv5/filters.go
@@ -4,7 +4,6 @@ import (
 	wenr "github.com/waku-org/go-waku/waku/v2/protocol/enr"
 
 	"github.com/ethereum/go-ethereum/p2p/enode"
-	"github.com/ethereum/go-ethereum/p2p/enr"
 )
 
 // FilterPredicate is to create a Predicate using a custom function
@@ -36,16 +35,11 @@ func FilterShard(cluster, index uint16) Predicate {
 func FilterCapabilities(flags wenr.WakuEnrBitfield) Predicate {
 	return func(iterator enode.Iterator) enode.Iterator {
 		predicate := func(node *enode.Node) bool {
-			enrField := new(wenr.WakuEnrBitfield)
-			if err := node.Record().Load(enr.WithEntry(wenr.WakuENRField, &enrField)); err != nil {
+			enrField, err := wenr.GetWakuEnrBitField(node)
+			if err != nil {
 				return false
 			}
-
-			if enrField == nil {
-				return false
-			}
-
-			return *enrField&flags == flags
+			return enrField&flags == flags
 		}
 		return enode.Filter(iterator, predicate)
 	}
diff --git a/waku/v2/peermanager/peer_manager.go b/waku/v2/peermanager/peer_manager.go
index 2ac489a04..c0f1c3fb2 100644
--- a/waku/v2/peermanager/peer_manager.go
+++ b/waku/v2/peermanager/peer_manager.go
@@ -6,7 +6,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/ethereum/go-ethereum/p2p/enr"
 	pubsub "github.com/libp2p/go-libp2p-pubsub"
 	"github.com/libp2p/go-libp2p/core/event"
 	"github.com/libp2p/go-libp2p/core/host"
@@ -535,8 +534,8 @@ func (pm *PeerManager) processPeerENR(p *service.PeerData) []protocol.ID {
 	}
 	supportedProtos := []protocol.ID{}
 	//Identify and specify protocols supported by the peer based on the discovered peer's ENR
-	var enrField wenr.WakuEnrBitfield
-	if err := p.ENR.Record().Load(enr.WithEntry(wenr.WakuENRField, &enrField)); err == nil {
+	enrField, err := wenr.GetWakuEnrBitField(p.ENR)
+	if err == nil {
 		for proto, protoENR := range pm.wakuprotoToENRFieldMap {
 			protoENRField := protoENR.waku2ENRBitField
 			if protoENRField&enrField != 0 {
diff --git a/waku/v2/protocol/enr/enr.go b/waku/v2/protocol/enr/enr.go
index fe1e54625..27c332c3c 100644
--- a/waku/v2/protocol/enr/enr.go
+++ b/waku/v2/protocol/enr/enr.go
@@ -28,6 +28,23 @@ const ShardingBitVectorEnrField = "rsv"
 // WakuEnrBitfield is a8-bit flag field to indicate Waku capabilities. Only the 4 LSBs are currently defined according to RFC31 (https://rfc.vac.dev/spec/31/).
 type WakuEnrBitfield = uint8
 
+func GetWakuEnrBitField(node *enode.Node) (WakuEnrBitfield, error) {
+	enrField := []byte{}
+	err := node.Record().Load(enr.WithEntry(WakuENRField, &enrField))
+	if err != nil {
+		if enr.IsNotFound(err) {
+			return 0, nil
+		}
+		return 0, err
+	}
+
+	if len(enrField) == 0 {
+		return 0, err
+	}
+
+	return WakuEnrBitfield(enrField[0]), nil
+}
+
 // NewWakuEnrBitfield creates a WakuEnrBitField whose value will depend on which protocols are enabled in the node
 func NewWakuEnrBitfield(lightpush, filter, store, relay bool) WakuEnrBitfield {
 	var v uint8