From 075815b072baa9a6f74458057b910b3e5f08da8e Mon Sep 17 00:00:00 2001
From: Lorenzo Delgado <lorenzo@status.im>
Date: Fri, 10 Mar 2023 11:49:41 +0100
Subject: [PATCH] feat(enr): added enr builder waku capabilities extension

---
 tests/common/test_enr_builder.nim |  2 +-
 tests/v2/test_waku_enr.nim        | 33 ++++++++++++++++++++++++++++---
 waku/v2/protocol/waku_enr.nim     | 26 +++++++++++++++++-------
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/tests/common/test_enr_builder.nim b/tests/common/test_enr_builder.nim
index e9ef45520..d6aa4a234 100644
--- a/tests/common/test_enr_builder.nim
+++ b/tests/common/test_enr_builder.nim
@@ -1,7 +1,7 @@
 {.used.}
 
 import
-  std/options ,
+  std/options,
   stew/results,
   stew/shims/net,
   testutils/unittests
diff --git a/tests/v2/test_waku_enr.nim b/tests/v2/test_waku_enr.nim
index f89d2be36..5c2feda5f 100644
--- a/tests/v2/test_waku_enr.nim
+++ b/tests/v2/test_waku_enr.nim
@@ -37,7 +37,31 @@ suite "Waku ENR -  Capabilities bitfield":
     check:
       caps == @[Capabilities.Relay, Capabilities.Filter, Capabilities.Lightpush]
 
-  test "encode and extract capabilities from record":
+  test "encode and extract capabilities from record (EnrBuilder ext)":
+    ## Given
+    let
+      enrSeqNum = 1u64
+      enrPrivKey = generatesecp256k1key()
+
+    ## When
+    var builder = EnrBuilder.init(enrPrivKey, seqNum = enrSeqNum)
+    builder.withWakuCapabilities(Capabilities.Relay, Capabilities.Store)
+
+    let recordRes = builder.build()
+
+    ## Then
+    check recordRes.isOk()
+    let record = recordRes.tryGet()
+
+    let bitfieldRes = record.getCapabilitiesField()
+    check bitfieldRes.isOk()
+
+    let bitfield = bitfieldRes.tryGet()
+    check:
+      bitfield.toCapabilities() == @[Capabilities.Relay, Capabilities.Store]
+
+  test "encode and extract capabilities from record (deprecated)":
+    # TODO: Remove after removing the `Record.init()` proc
     ## Given
     let enrkey = generatesecp256k1key()
     let caps = CapabilitiesBitfield.init(Capabilities.Relay, Capabilities.Store)
@@ -56,8 +80,11 @@ suite "Waku ENR -  Capabilities bitfield":
 
   test "cannot extract capabilities from record":
     ## Given
-    let enrkey = generatesecp256k1key()
-    let record = Record.init(1, enrkey, wakuFlags=none(CapabilitiesBitfield))
+    let
+      enrSeqNum = 1u64
+      enrPrivKey = generatesecp256k1key()
+
+    let record = EnrBuilder.init(enrPrivKey, enrSeqNum).build().tryGet()
 
     ## When
     let bitfieldRes = record.getCapabilitiesField()
diff --git a/waku/v2/protocol/waku_enr.nim b/waku/v2/protocol/waku_enr.nim
index 563349573..ee6b01870 100644
--- a/waku/v2/protocol/waku_enr.nim
+++ b/waku/v2/protocol/waku_enr.nim
@@ -12,9 +12,10 @@ import
   stew/[endians2, results],
   stew/shims/net,
   eth/keys,
-  eth/p2p/discoveryv5/enr,
   libp2p/[multiaddress, multicodec],
   libp2p/crypto/crypto
+import
+  ../../common/enr
 
 export enr, crypto, multiaddress, net
 
@@ -23,16 +24,17 @@ const
   CapabilitiesEnrField* = "waku2"
 
 
-## Capabilities
+## Node capabilities
 
 type
-  ## 8-bit flag field to indicate Waku capabilities.
+  ## 8-bit flag field to indicate Waku node capabilities.
   ## Only the 4 LSBs are currently defined according
   ## to RFC31 (https://rfc.vac.dev/spec/31/).
   CapabilitiesBitfield* = distinct uint8
 
   ## See: https://rfc.vac.dev/spec/31/#waku2-enr-key
   ## each enum numbers maps to a bit (where 0 is the LSB)
+  # TODO: Make this enum {.pure.}
   Capabilities* = enum
     Relay = 0,
     Store = 1,
@@ -66,9 +68,19 @@ func toCapabilities*(bitfield: CapabilitiesBitfield): seq[Capabilities] =
   toSeq(Capabilities.low..Capabilities.high).filterIt(supportsCapability(bitfield, it))
 
 
-## TODO: Turn into an EnrBuilder extension
-func toFieldPair*(caps: CapabilitiesBitfield): FieldPair =
-  toFieldPair(CapabilitiesEnrField, @[caps.uint8])
+# ENR builder extension
+
+proc withWakuCapabilities*(builder: var EnrBuilder, caps: CapabilitiesBitfield) =
+  builder.addFieldPair(CapabilitiesEnrField, @[caps.uint8])
+
+proc withWakuCapabilities*(builder: var EnrBuilder, caps: varargs[Capabilities]) =
+  withWakuCapabilities(builder, CapabilitiesBitfield.init(caps))
+
+proc withWakuCapabilities*(builder: var EnrBuilder, caps: openArray[Capabilities]) =
+  withWakuCapabilities(builder, CapabilitiesBitfield.init(@caps))
+
+
+# ENR record accessors (e.g., Record, TypedRecord, etc.)
 
 proc getCapabilitiesField*(r: Record): EnrResult[CapabilitiesBitfield] =
   let field = ?r.get(CapabilitiesEnrField, seq[uint8])
@@ -201,7 +213,7 @@ func init*(T: type enr.Record,
 
   # `waku2` field
   if wakuFlags.isSome():
-    wakuEnrFields.add(toFieldPair(wakuFlags.get()))
+    wakuEnrFields.add(toFieldPair(CapabilitiesEnrField, @[wakuFlags.get().uint8]))
 
   # `multiaddrs` field
   if multiaddrs.len > 0: