From 157724d9b92812cf73a9760b602676139e94d758 Mon Sep 17 00:00:00 2001
From: Lorenzo Delgado <lorenzo@status.im>
Date: Wed, 22 Mar 2023 16:01:02 +0100
Subject: [PATCH] refactor(enr): added waku2 capabilities accessor

---
 tests/v2/test_waku_enr.nim    | 66 ++++++++++++++++++++++-------------
 waku/v2/protocol/waku_enr.nim | 33 ++++++++++++------
 2 files changed, 64 insertions(+), 35 deletions(-)

diff --git a/tests/v2/test_waku_enr.nim b/tests/v2/test_waku_enr.nim
index e693219c9..57632cf33 100644
--- a/tests/v2/test_waku_enr.nim
+++ b/tests/v2/test_waku_enr.nim
@@ -53,10 +53,13 @@ suite "Waku ENR -  Capabilities bitfield":
     check recordRes.isOk()
     let record = recordRes.tryGet()
 
-    let bitfieldRes = record.getCapabilitiesField()
-    check bitfieldRes.isOk()
+    let typedRecord = record.toTyped()
+    require typedRecord.isOk()
 
-    let bitfield = bitfieldRes.tryGet()
+    let bitfieldOpt = typedRecord.value.waku2
+    check bitfieldOpt.isSome()
+
+    let bitfield = bitfieldOpt.get()
     check:
       bitfield.toCapabilities() == @[Capabilities.Relay, Capabilities.Store]
 
@@ -69,12 +72,15 @@ suite "Waku ENR -  Capabilities bitfield":
     let record = Record.init(1, enrkey, wakuFlags=some(caps))
 
     ## When
-    let bitfieldRes = record.getCapabilitiesField()
+    let typedRecord = record.toTyped()
+    require typedRecord.isOk()
+
+    let bitfieldOpt = typedRecord.value.waku2
 
     ## Then
-    check bitfieldRes.isOk()
+    check bitfieldOpt.isSome()
 
-    let bitfield = bitfieldRes.tryGet()
+    let bitfield = bitfieldOpt.get()
     check:
       bitfield.toCapabilities() == @[Capabilities.Relay, Capabilities.Store]
 
@@ -87,14 +93,13 @@ suite "Waku ENR -  Capabilities bitfield":
     let record = EnrBuilder.init(enrPrivKey, enrSeqNum).build().tryGet()
 
     ## When
-    let bitfieldRes = record.getCapabilitiesField()
+    let typedRecord = record.toTyped()
+    require typedRecord.isOk()
+
+    let bitfieldOpt = typedRecord.value.waku2
 
     ## Then
-    check bitfieldRes.isErr()
-
-    let err = bitfieldRes.tryError()
-    check:
-      err == "Key not found in ENR"
+    check bitfieldOpt.isNone()
 
   test "check capabilities on a waku node record":
     ## Given
@@ -107,12 +112,19 @@ suite "Waku ENR -  Capabilities bitfield":
     require waku_enr.fromBase64(record, wakuRecord)
 
     ## Then
+    let typedRecordRes = record.toTyped()
+    require typedRecordRes.isOk()
+
+    let bitfieldOpt = typedRecordRes.value.waku2
+    require bitfieldOpt.isSome()
+
+    let bitfield = bitfieldOpt.get()
     check:
-      record.supportsCapability(Relay) == true
-      record.supportsCapability(Store) == true
-      record.supportsCapability(Filter) == false
-      record.supportsCapability(Lightpush) == false
-      record.getCapabilities() == @[Capabilities.Relay, Capabilities.Store]
+      bitfield.supportsCapability(Capabilities.Relay) == true
+      bitfield.supportsCapability(Capabilities.Store) == true
+      bitfield.supportsCapability(Capabilities.Filter) == false
+      bitfield.supportsCapability(Capabilities.Lightpush) == false
+      bitfield.toCapabilities() == @[Capabilities.Relay, Capabilities.Store]
 
   test "check capabilities on a non-waku node record":
     ## Given
@@ -122,16 +134,22 @@ suite "Waku ENR -  Capabilities bitfield":
     "Y3AyNTZrMaEDhpehBDbZjM_L9ek699Y7vhUJ-eAdMyQW_Fil522Y0fODdGNwgiMog3VkcIIjKA"
 
     ## When
-    var nonWakuEnrRecord: Record
-    require waku_enr.fromURI(nonWakuEnrRecord, nonWakuEnr)
+    var record: Record
+    require waku_enr.fromURI(record, nonWakuEnr)
 
     ## Then
+    let typedRecordRes = record.toTyped()
+    require typedRecordRes.isOk()
+
+    let bitfieldOpt = typedRecordRes.value.waku2
+    check bitfieldOpt.isNone()
+
     check:
-      nonWakuEnrRecord.getCapabilities() == []
-      nonWakuEnrRecord.supportsCapability(Relay) == false
-      nonWakuEnrRecord.supportsCapability(Store) == false
-      nonWakuEnrRecord.supportsCapability(Filter) == false
-      nonWakuEnrRecord.supportsCapability(Lightpush) == false
+      record.getCapabilities() == []
+      record.supportsCapability(Capabilities.Relay) == false
+      record.supportsCapability(Capabilities.Store) == false
+      record.supportsCapability(Capabilities.Filter) == false
+      record.supportsCapability(Capabilities.Lightpush) == false
 
 
 suite "Waku ENR - Multiaddresses":
diff --git a/waku/v2/protocol/waku_enr.nim b/waku/v2/protocol/waku_enr.nim
index 0826ebe00..1d9697ae1 100644
--- a/waku/v2/protocol/waku_enr.nim
+++ b/waku/v2/protocol/waku_enr.nim
@@ -34,8 +34,7 @@ type
 
   ## 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
+  Capabilities*{.pure.} = enum
     Relay = 0,
     Store = 1,
     Filter = 2,
@@ -82,24 +81,36 @@ proc withWakuCapabilities*(builder: var EnrBuilder, caps: openArray[Capabilities
 
 # ENR record accessors (e.g., Record, TypedRecord, etc.)
 
-proc getCapabilitiesField*(r: Record): EnrResult[CapabilitiesBitfield] =
-  let field = ?r.get(CapabilitiesEnrField, seq[uint8])
-  ok(CapabilitiesBitfield(field[0]))
+func waku2*(record: TypedRecord): Option[CapabilitiesBitfield] =
+  let field = record.tryGet(CapabilitiesEnrField, seq[uint8])
+  if field.isNone():
+    return none(CapabilitiesBitfield)
 
+  some(CapabilitiesBitfield(field.get()[0]))
 
 proc supportsCapability*(r: Record, cap: Capabilities): bool =
-  let bitfield = getCapabilitiesField(r)
-  if bitfield.isErr():
+  let recordRes = r.toTyped()
+  if recordRes.isErr():
     return false
 
-  bitfield.value.supportsCapability(cap)
+  let bitfieldOpt = recordRes.value.waku2
+  if bitfieldOpt.isNone():
+    return false
+
+  let bitfield = bitfieldOpt.get()
+  bitfield.supportsCapability(cap)
 
 proc getCapabilities*(r: Record): seq[Capabilities] =
-  let bitfield = getCapabilitiesField(r)
-  if bitfield.isErr():
+  let recordRes = r.toTyped()
+  if recordRes.isErr():
     return @[]
 
-  bitfield.value.toCapabilities()
+  let bitfieldOpt = recordRes.value.waku2
+  if bitfieldOpt.isNone():
+    return @[]
+
+  let bitfield = bitfieldOpt.get()
+  bitfield.toCapabilities()
 
 
 ## Multiaddress