Allow for passing in previous enr at discovery protocol creation

eth/p2p/discoveryv5/enr.nim

@@ -227,8 +227,8 @@ proc find(r: Record, key: string): Option[int] =
     if k == key:
       return some(i)
 
-proc insertFieldPairs*(record: var Record, pk: PrivateKey,
+proc update*(record: Record, pk: PrivateKey, fieldPairs: openarray[FieldPair]):
-    fieldPairs: openarray[FieldPair]): EnrResult[bool] =
+    EnrResult[Record] =
   var r = record
 
   let pubkey = r.get(PublicKey)
@@ -255,16 +255,14 @@ proc insertFieldPairs*(record: var Record, pk: PrivateKey,
   if updated:
     r.seqNum.inc()
     r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
-    record = r
-    ok(true)
-  else:
-    ok(true)
 
-proc update*(r: var Record, pk: PrivateKey,
+  ok(r)
-                            ip: Option[ValidIpAddress],
+
-                            tcpPort, udpPort: Port,
+proc update*(r: Record, pk: PrivateKey,
-                            extraFields: openarray[FieldPair]):
+                        ip: Option[ValidIpAddress],
-                            EnrResult[bool] =
+                        tcpPort, udpPort: Port,
+                        extraFields: openarray[FieldPair] = []):
+                        EnrResult[Record] =
   var fields = newSeq[FieldPair]()
 
   if ip.isSome():
@@ -282,7 +280,7 @@ proc update*(r: var Record, pk: PrivateKey,
 
   fields.add extraFields
 
-  r.insertFieldPairs(pk, fields)
+  r.update(pk, fields)
 
 proc tryGet*(r: Record, key: string, T: type): Option[T] =
   try:

eth/p2p/discoveryv5/protocol.nim

@@ -693,6 +693,7 @@ proc newProtocol*(privKey: PrivateKey, db: Database,
                   externalIp: Option[ValidIpAddress], tcpPort, udpPort: Port,
                   localEnrFields: openarray[FieldPair] = [],
                   bootstrapRecords: openarray[Record] = [],
+                  previousEnr = none[enr.Record](),
                   bindIp = IPv4_any(), rng = newRng()):
                   Protocol {.raises: [Defect].} =
   # TODO: Tried adding bindPort = udpPort as parameter but that gave
@@ -701,10 +702,16 @@ proc newProtocol*(privKey: PrivateKey, db: Database,
   # remapping through NAT and this API is also subject to change once we
   # introduce support for ipv4 + ipv6 binding/listening.
   let
-    # TODO: Defect or return a result on this call?
+    # TODO:
-    enrRec = enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
+    # - Defect as is now or return a result for enr errors?
-      localEnrFields).expect("Record within size limits")
+    # - In case incorrect key, allow for new enr based on new key (new node id)?
-    node = newNode(enrRec).expect("Properly initialized node")
+    enr = if previousEnr.isSome():
+            previousEnr.get().update(privKey, externalIp, tcpPort, udpPort,
+              localEnrFields).expect("Record within size limits and correct key")
+          else:
+            enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
+              localEnrFields).expect("Record within size limits")
+    node = newNode(enr).expect("Properly initialized record")
 
   # TODO Consider whether this should be a Defect
   doAssert rng != nil, "RNG initialization failed"

tests/p2p/discv5_test_helper.nim

@@ -7,7 +7,8 @@ import
 proc localAddress*(port: int): Address =
   Address(ip: ValidIpAddress.init("127.0.0.1"), port: Port(port))
 
-proc initDiscoveryNode*(rng: ref BrHmacDrbgContext, privKey: PrivateKey, address: Address,
+proc initDiscoveryNode*(rng: ref BrHmacDrbgContext, privKey: PrivateKey,
+                        address: Address,
                         bootstrapRecords: openarray[Record] = [],
                         localEnrFields: openarray[FieldPair] = []):
                         discv5_protocol.Protocol =

tests/p2p/test_discoveryv5.nim

@@ -54,7 +54,6 @@ procSuite "Discovery v5 Tests":
 
     await node1.closeWait()
 
-
   asyncTest "Handshake cleanup":
     let node = initDiscoveryNode(
       rng, PrivateKey.random(rng[]), localAddress(20302))

tests/p2p/test_enr.nim

@@ -87,7 +87,7 @@ suite "ENR":
       let r = initRecord(1, pk, {"maxplus1": repeat(byte 2, 170),})
       check r.isErr()
 
-  test "ENR insert":
+  test "ENR update":
     let
       pk = PrivateKey.fromHex(
         "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
@@ -95,28 +95,31 @@ suite "ENR":
     var r = Record.init(1, pk, none(ValidIpAddress), Port(9000), Port(9000))[]
 
     block: # Insert new k:v pair, update of seqNum should occur.
-      let res = r.insertFieldPairs(pk, [newField])
+      let res = r.update(pk, [newField])
+      check res.isOk()
+      r = res[]
       check:
-        res.isOk()
         r.get("test", uint) == 123
         r.seqNum == 2
 
     block: # Insert same k:v pair, no update of seqNum should occur.
-      let res = r.insertFieldPairs(pk, [newField])
+      let res = r.update(pk, [newField])
+      check res.isOk()
+      r = res[]
       check:
-        res.isOk()
         r.get("test", uint) == 123
         r.seqNum == 2
 
     block: # Insert k:v pair with changed value, update of seqNum should occur.
       let updatedField = toFieldPair("test", 1234'u)
-      let res =  r.insertFieldPairs(pk, [updatedField])
+      let res = r.update(pk, [updatedField])
+      check res.isOk()
+      r = res[]
       check:
-        res.isOk()
         r.get("test", uint) == 1234
         r.seqNum == 3
 
-  test "ENR insert sorted":
+  test "ENR update sorted":
     let pk = PrivateKey.fromHex(
       "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
     var r = initRecord(123, pk, {"abc": 1234'u,
@@ -127,11 +130,11 @@ suite "ENR":
 
     let newField = toFieldPair("test", 123'u)
     let newField2 = toFieldPair("zzz", 123'u)
-    let res = r.insertFieldPairs(pk, [newField, newField2])
+    let res = r.update(pk, [newField, newField2])
     check res.isOk()
-    check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00, zzz: 123)"""
+    check $res[] == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00, zzz: 123)"""
 
-  test "ENR insert size too big":
+  test "ENR update size too big":
     let pk = PrivateKey.fromHex(
       "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
 
@@ -139,10 +142,10 @@ suite "ENR":
     check r.isOk()
 
     let newField = toFieldPair("test", 123'u)
-    let res = r[].insertFieldPairs(pk, [newField])
+    let res = r[].update(pk, [newField])
     check res.isErr()
 
-  test "ENR insert invalid key":
+  test "ENR update invalid key":
     let pk = PrivateKey.fromHex(
       "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
 
@@ -152,5 +155,29 @@ suite "ENR":
     let
       wrongPk = PrivateKey.random(rng[])
       newField = toFieldPair("test", 123'u)
-      res = r[].insertFieldPairs(wrongPk, [newField])
+      res = r[].update(wrongPk, [newField])
     check res.isErr()
+
+  test "ENR update address":
+    let
+      pk = PrivateKey.fromHex(
+        "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
+    var r = Record.init(1, pk, none(ValidIpAddress), Port(9000), Port(9000))[]
+
+    block:
+      let res = r.update(pk, none(ValidIpAddress), Port(9000), Port(9000))
+      check res.isOk()
+      r = res[]
+      check:
+        r.get("tcp", uint) == 9000
+        r.get("udp", uint) == 9000
+        r.seqNum == 1
+
+    block:
+      let res = r.update(pk, none(ValidIpAddress), Port(9001), Port(9002))
+      check res.isOk()
+      r = res[]
+      check:
+        r.get("tcp", uint) == 9001
+        r.get("udp", uint) == 9002
+        r.seqNum == 2