From 3bbe757fe37c307fba49c9279f6de94bb8a6ed95 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Tue, 7 Jul 2020 17:19:15 +0200
Subject: [PATCH 1/8] Add record size check on init + add insertFieldPair call
 + tests

---
 eth/p2p/discoveryv5/enr.nim      | 127 ++++++++++++++++++++++++-------
 eth/p2p/discoveryv5/protocol.nim |   3 +-
 tests/p2p/test_enr.nim           |  94 +++++++++++++++++++++--
 3 files changed, 189 insertions(+), 35 deletions(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index a33ebbf..e059d99 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/eth/p2p/discoveryv5/enr.nim
@@ -1,4 +1,4 @@
-# ENR implemetation according to spec:
+# ENR implementation according to specification in EIP-778:
 # https://github.com/ethereum/EIPs/blob/master/EIPS/eip-778.md
 
 import
@@ -11,8 +11,9 @@ export options
 {.push raises: [Defect].}
 
 const
-  maxEnrSize = 300
-  minRlpListLen = 4 # for signature, seqId, "id" key, id
+  maxEnrSize = 300  ## Maximum size of an encoded node record, in bytes.
+  minRlpListLen = 4 ## Minimum node record RLP list has: signature, seqId,
+  ## "id" key and value.
 
 type
   FieldPair* = (string, Field)
@@ -63,8 +64,49 @@ template toField[T](v: T): Field =
   else:
     {.error: "Unsupported field type".}
 
+proc `==`(a, b: Field): bool =
+  if a.kind == b.kind:
+    case a.kind
+    of kString:
+      return a.str == b.str
+    of kNum:
+      return a.num == b.num
+    of kBytes:
+      return a.bytes == b.bytes
+  else:
+    return false
+
+proc makeEnrRaw(seqNum: uint64, pk: PrivateKey,
+    pairs: openarray[FieldPair]): EnrResult[seq[byte]] =
+  proc append(w: var RlpWriter, seqNum: uint64,
+      pairs: openarray[FieldPair]): seq[byte] =
+    w.append(seqNum)
+    for (k, v) in pairs:
+      w.append(k)
+      case v.kind
+      of kString: w.append(v.str)
+      of kNum: w.append(v.num)
+      of kBytes: w.append(v.bytes)
+    w.finish()
+
+  let toSign = block:
+    var w = initRlpList(pairs.len * 2 + 1)
+    w.append(seqNum, pairs)
+
+  let sig = signNR(pk, toSign)
+
+  var raw = block:
+    var w = initRlpList(pairs.len * 2 + 2)
+    w.append(sig.toRaw())
+    w.append(seqNum, pairs)
+
+  if raw.len > maxEnrSize:
+    err("Record exceeds maximum size")
+  else:
+    ok(raw)
+
 proc makeEnrAux(seqNum: uint64, pk: PrivateKey,
-    pairs: openarray[(string, Field)]): EnrResult[Record] =
+    pairs: openarray[FieldPair]): EnrResult[Record] =
   var record: Record
   record.pairs = @pairs
   record.seqNum = seqNum
@@ -76,35 +118,18 @@ proc makeEnrAux(seqNum: uint64, pk: PrivateKey,
     Field(kind: kBytes, bytes: @(pubkey.toRawCompressed()))))
 
   # Sort by key
-  record.pairs.sort() do(a, b: (string, Field)) -> int:
+  record.pairs.sort() do(a, b: FieldPair) -> int:
     cmp(a[0], b[0])
 
-  proc append(w: var RlpWriter, seqNum: uint64,
-      pairs: openarray[(string, Field)]): seq[byte] =
-    w.append(seqNum)
-    for (k, v) in pairs:
-      w.append(k)
-      case v.kind
-      of kString: w.append(v.str)
-      of kNum: w.append(v.num)
-      of kBytes: w.append(v.bytes)
-    w.finish()
-
-  let toSign = block:
-    var w = initRlpList(record.pairs.len * 2 + 1)
-    w.append(seqNum, record.pairs)
-
-  let sig = signNR(pk, toSign)
-
-  record.raw = block:
-    var w = initRlpList(record.pairs.len * 2 + 2)
-    w.append(sig.toRaw())
-    w.append(seqNum, record.pairs)
-
+  record.raw = ? makeEnrRaw(seqNum, pk, record.pairs)
   ok(record)
 
 macro initRecord*(seqNum: uint64, pk: PrivateKey,
     pairs: untyped{nkTableConstr}): untyped =
+  ## Initialize a `Record` with given sequence number, private key and k:v
+  ## pairs.
+  ##
+  ## Can fail in case the record exceeds the `maxEnrSize`.
   for c in pairs:
     c.expectKind(nnkExprColonExpr)
     c[1] = newCall(bindSym"toField", c[1])
@@ -121,6 +146,10 @@ proc init*(T: type Record, seqNum: uint64,
                            tcpPort, udpPort: Port,
                            extraFields: openarray[FieldPair] = []):
                            EnrResult[T] =
+  ## Initialize a `Record` with given sequence number, private key, optional
+  ## ip-address, tcp port, udp port, and optional custom k:v pairs.
+  ##
+  ## Can fail in case the record exceeds the `maxEnrSize`.
   var fields = newSeq[FieldPair]()
 
   if ip.isSome():
@@ -190,6 +219,50 @@ proc get*(r: Record, T: type PublicKey): Option[T] =
     if pk.isOk:
       return some pk[]
 
+proc find(r: Record, key: string): Option[int] =
+  ## Search for key in record key:value pairs.
+  ##
+  ## Returns some(index of key) if key is found in record. Else return none.
+  for i, (k, v) in r.pairs:
+    if k == key:
+      return some(i)
+
+proc insertFieldPair*(record: var Record, pk: PrivateKey, fieldPair: FieldPair):
+    EnrResult[bool] =
+  ## Insert or modify a k:v pair to the `Record`.
+  ##
+  ## If k:v pair doesn't exist yet, it will be inserted. If it does exist, it
+  ## will be updated if the value is different, else nothing is done. In case of
+  ## an insert or update the seqNum will be be incremented and a new signature
+  ## will be applied.
+  ##
+  ## Can fail in case of wrong PrivateKey or if the size of the resulting record
+  ## is > maxEnrSize. `record` will not be altered in these cases.
+  var r = record
+
+  let pubkey = r.get(PublicKey)
+  if pubkey.isNone() or pubkey.get() != pk.toPublicKey():
+    return err("Public key does not correspond with given private key")
+
+  let index = r.find(fieldPair[0])
+  if(index.isSome()):
+    if r.pairs[index.get()][1] == fieldPair[1]:
+      # Exact k:v pair is already in record, nothing to do here.
+      return ok(true)
+    else:
+      # Need to update the value.
+      r.pairs[index.get()] = fieldPair
+  else:
+    # Add new k:v pair.
+    r.pairs.insert(fieldPair,
+      lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0]))
+
+  r.seqNum.inc()
+
+  r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
+  record = r
+  ok(true)
+
 proc tryGet*(r: Record, key: string, T: type): Option[T] =
   try:
     return some get(r, key, T)
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index abe2357..2d3105b 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -701,8 +701,9 @@ 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?
     enrRec = enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
-      localEnrFields).expect("Properly intialized private key")
+      localEnrFields).expect("Record within size limits")
     node = newNode(enrRec).expect("Properly initialized node")
 
   # TODO Consider whether this should be a Defect
diff --git a/tests/p2p/test_enr.nim b/tests/p2p/test_enr.nim
index 85f83af..b95858e 100644
--- a/tests/p2p/test_enr.nim
+++ b/tests/p2p/test_enr.nim
@@ -1,5 +1,5 @@
 import
-  unittest, options,
+  unittest, options, sequtils,
   nimcrypto/utils, stew/shims/net,
   eth/p2p/enode, eth/p2p/discoveryv5/enr, eth/keys
 
@@ -35,14 +35,15 @@ suite "ENR":
 
   test "Create from ENode address":
     let
-      keys = KeyPair.random(rng[])
+      keypair = KeyPair.random(rng[])
       ip = ValidIpAddress.init("10.20.30.40")
-      enr = Record.init(100, keys.seckey, some(ip), Port(9000), Port(9000), @[])[]
+      enr = Record.init(
+        100, keypair.seckey, some(ip), Port(9000), Port(9000),@[])[]
       typedEnr = get enr.toTypedRecord()
 
     check:
       typedEnr.secp256k1.isSome()
-      typedEnr.secp256k1.get == keys.pubkey.toRawCompressed()
+      typedEnr.secp256k1.get == keypair.pubkey.toRawCompressed()
 
       typedEnr.ip.isSome()
       typedEnr.ip.get() == [byte 10, 20, 30, 40]
@@ -55,13 +56,14 @@ suite "ENR":
 
   test "ENR without address":
     let
-      keys = KeyPair.random(rng[])
-      enr = Record.init(100, keys.seckey, none(ValidIpAddress), Port(9000), Port(9000))[]
+      keypair = KeyPair.random(rng[])
+      enr = Record.init(
+        100, keypair.seckey, none(ValidIpAddress), Port(9000), Port(9000))[]
       typedEnr = get enr.toTypedRecord()
 
     check:
       typedEnr.secp256k1.isSome()
-      typedEnr.secp256k1.get() == keys.pubkey.toRawCompressed()
+      typedEnr.secp256k1.get() == keypair.pubkey.toRawCompressed()
 
       typedEnr.ip.isNone()
       typedEnr.tcp.isSome()
@@ -73,3 +75,81 @@ suite "ENR":
       typedEnr.ip6.isNone()
       typedEnr.tcp6.isNone()
       typedEnr.udp6.isNone()
+
+  test "ENR init size too big":
+    let pk = PrivateKey.fromHex(
+      "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
+    block: # This gives ENR of 300 bytes encoded
+      let r = initRecord(1, pk, {"maxvalue": repeat(byte 2, 169),})
+      check r.isOk()
+
+    block: # This gives ENR of 301 bytes encoded
+      let r = initRecord(1, pk, {"maxplus1": repeat(byte 2, 170),})
+      check r.isErr()
+
+  test "ENR insert":
+    let
+      pk = PrivateKey.fromHex(
+        "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
+      newField = toFieldPair("test", 123'u)
+    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.insertFieldPair(pk, newField)
+      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.insertFieldPair(pk, newField)
+      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.insertFieldPair(pk, updatedField)
+      check:
+        res.isOk()
+        r.get("test", uint) == 1234
+        r.seqNum == 3
+
+  test "ENR insert sorted":
+    let pk = PrivateKey.fromHex(
+      "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
+    var r = initRecord(123, pk, {"abc": 1234'u,
+                                 "z": [byte 0],
+                                 "123": "abc",
+                                 "a12": 1'u})[]
+    check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, z: 0x00)"""
+
+    let newField = toFieldPair("test", 123'u)
+    let res = r.insertFieldPair(pk, newField)
+    check res.isOk()
+    check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00)"""
+
+  test "ENR insert size too big":
+    let pk = PrivateKey.fromHex(
+      "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
+
+    var r = initRecord(1, pk, {"maxvalue": repeat(byte 2, 169),})
+    check r.isOk()
+
+    let newField = toFieldPair("test", 123'u)
+    let res = r[].insertFieldPair(pk, newField)
+    check res.isErr()
+
+  test "ENR insert invalid key":
+    let pk = PrivateKey.fromHex(
+      "5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
+
+    var r = initRecord(1, pk, {"abc": 1'u,})
+    check r.isOk()
+
+    let
+      wrongPk = PrivateKey.random(rng[])
+      newField = toFieldPair("test", 123'u)
+      res = r[].insertFieldPair(wrongPk, newField)
+    check res.isErr()

From 57302fcf52cea6fedd189e50bbeea41bd5c8a253 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Tue, 7 Jul 2020 22:48:26 +0200
Subject: [PATCH 2/8] Allow for multiple enr fields to be inserted + enr update
 proc

---
 eth/p2p/discoveryv5/enr.nim | 71 ++++++++++++++++++++++++-------------
 tests/p2p/test_enr.nim      | 15 ++++----
 2 files changed, 54 insertions(+), 32 deletions(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index e059d99..3d98af9 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/eth/p2p/discoveryv5/enr.nim
@@ -227,41 +227,62 @@ proc find(r: Record, key: string): Option[int] =
     if k == key:
       return some(i)
 
-proc insertFieldPair*(record: var Record, pk: PrivateKey, fieldPair: FieldPair):
-    EnrResult[bool] =
-  ## Insert or modify a k:v pair to the `Record`.
-  ##
-  ## If k:v pair doesn't exist yet, it will be inserted. If it does exist, it
-  ## will be updated if the value is different, else nothing is done. In case of
-  ## an insert or update the seqNum will be be incremented and a new signature
-  ## will be applied.
-  ##
-  ## Can fail in case of wrong PrivateKey or if the size of the resulting record
-  ## is > maxEnrSize. `record` will not be altered in these cases.
+proc insertFieldPairs*(record: var Record, pk: PrivateKey,
+    fieldPairs: openarray[FieldPair]): EnrResult[bool] =
   var r = record
 
   let pubkey = r.get(PublicKey)
   if pubkey.isNone() or pubkey.get() != pk.toPublicKey():
     return err("Public key does not correspond with given private key")
 
-  let index = r.find(fieldPair[0])
-  if(index.isSome()):
-    if r.pairs[index.get()][1] == fieldPair[1]:
-      # Exact k:v pair is already in record, nothing to do here.
-      return ok(true)
+  var updated = false
+  for fieldPair in fieldPairs:
+    let index = r.find(fieldPair[0])
+    if(index.isSome()):
+      if r.pairs[index.get()][1] == fieldPair[1]:
+        # Exact k:v pair is already in record, nothing to do here.
+        continue
+      else:
+        # Need to update the value.
+        r.pairs[index.get()] = fieldPair
+        updated = true
     else:
-      # Need to update the value.
-      r.pairs[index.get()] = fieldPair
+      # Add new k:v pair.
+      r.pairs.insert(fieldPair,
+        lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0]))
+      updated = true
+
+  if updated:
+    r.seqNum.inc()
+    r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
+    record = r
+    ok(true)
   else:
-    # Add new k:v pair.
-    r.pairs.insert(fieldPair,
-      lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0]))
+    ok(true)
 
-  r.seqNum.inc()
+proc update*(r: var Record, pk: PrivateKey,
+                            ip: Option[ValidIpAddress],
+                            tcpPort, udpPort: Port,
+                            extraFields: openarray[FieldPair]):
+                            EnrResult[bool] =
+  var fields = newSeq[FieldPair]()
 
-  r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
-  record = r
-  ok(true)
+  if ip.isSome():
+    let
+      ipExt = ip.get()
+      isV6 = ipExt.family == IPv6
+
+    fields.add(if isV6: ("ip6", ipExt.address_v6.toField)
+               else: ("ip", ipExt.address_v4.toField))
+    fields.add(((if isV6: "tcp6" else: "tcp"), tcpPort.uint16.toField))
+    fields.add(((if isV6: "udp6" else: "udp"), udpPort.uint16.toField))
+  else:
+    fields.add(("tcp", tcpPort.uint16.toField))
+    fields.add(("udp", udpPort.uint16.toField))
+
+  fields.add extraFields
+
+  r.insertFieldPairs(pk, fields)
 
 proc tryGet*(r: Record, key: string, T: type): Option[T] =
   try:
diff --git a/tests/p2p/test_enr.nim b/tests/p2p/test_enr.nim
index b95858e..720ba60 100644
--- a/tests/p2p/test_enr.nim
+++ b/tests/p2p/test_enr.nim
@@ -95,14 +95,14 @@ 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.insertFieldPair(pk, newField)
+      let res = r.insertFieldPairs(pk, [newField])
       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.insertFieldPair(pk, newField)
+      let res = r.insertFieldPairs(pk, [newField])
       check:
         res.isOk()
         r.get("test", uint) == 123
@@ -110,7 +110,7 @@ suite "ENR":
 
     block: # Insert k:v pair with changed value, update of seqNum should occur.
       let updatedField = toFieldPair("test", 1234'u)
-      let res =  r.insertFieldPair(pk, updatedField)
+      let res =  r.insertFieldPairs(pk, [updatedField])
       check:
         res.isOk()
         r.get("test", uint) == 1234
@@ -126,9 +126,10 @@ suite "ENR":
     check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, z: 0x00)"""
 
     let newField = toFieldPair("test", 123'u)
-    let res = r.insertFieldPair(pk, newField)
+    let newField2 = toFieldPair("zzz", 123'u)
+    let res = r.insertFieldPairs(pk, [newField, newField2])
     check res.isOk()
-    check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00)"""
+    check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00, zzz: 123)"""
 
   test "ENR insert size too big":
     let pk = PrivateKey.fromHex(
@@ -138,7 +139,7 @@ suite "ENR":
     check r.isOk()
 
     let newField = toFieldPair("test", 123'u)
-    let res = r[].insertFieldPair(pk, newField)
+    let res = r[].insertFieldPairs(pk, [newField])
     check res.isErr()
 
   test "ENR insert invalid key":
@@ -151,5 +152,5 @@ suite "ENR":
     let
       wrongPk = PrivateKey.random(rng[])
       newField = toFieldPair("test", 123'u)
-      res = r[].insertFieldPair(wrongPk, newField)
+      res = r[].insertFieldPairs(wrongPk, [newField])
     check res.isErr()

From 72420d7f171daaeba3e681d539ff402b35058aa0 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Tue, 7 Jul 2020 23:39:32 +0200
Subject: [PATCH 3/8] Allow for passing in previous enr at discovery protocol
 creation

---
 eth/p2p/discoveryv5/enr.nim      | 22 ++++++-------
 eth/p2p/discoveryv5/protocol.nim | 15 ++++++---
 tests/p2p/discv5_test_helper.nim |  3 +-
 tests/p2p/test_discoveryv5.nim   |  1 -
 tests/p2p/test_enr.nim           | 55 ++++++++++++++++++++++++--------
 5 files changed, 64 insertions(+), 32 deletions(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index 3d98af9..def2606 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/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,
-    fieldPairs: openarray[FieldPair]): EnrResult[bool] =
+proc update*(record: Record, pk: PrivateKey, fieldPairs: openarray[FieldPair]):
+    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,
-                            ip: Option[ValidIpAddress],
-                            tcpPort, udpPort: Port,
-                            extraFields: openarray[FieldPair]):
-                            EnrResult[bool] =
+  ok(r)
+
+proc update*(r: Record, pk: PrivateKey,
+                        ip: Option[ValidIpAddress],
+                        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:
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index 2d3105b..e10842e 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/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?
-    enrRec = enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
-      localEnrFields).expect("Record within size limits")
-    node = newNode(enrRec).expect("Properly initialized node")
+    # TODO:
+    # - Defect as is now or return a result for enr errors?
+    # - In case incorrect key, allow for new enr based on new key (new node id)?
+    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"
diff --git a/tests/p2p/discv5_test_helper.nim b/tests/p2p/discv5_test_helper.nim
index 299eea4..794525b 100644
--- a/tests/p2p/discv5_test_helper.nim
+++ b/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 =
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
index 4b56c85..c2a4547 100644
--- a/tests/p2p/test_discoveryv5.nim
+++ b/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))
diff --git a/tests/p2p/test_enr.nim b/tests/p2p/test_enr.nim
index 720ba60..c0936db 100644
--- a/tests/p2p/test_enr.nim
+++ b/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

From 0fb21e72d72040b930c7e8aaaedba89f722fbb53 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Wed, 8 Jul 2020 11:45:58 +0200
Subject: [PATCH 4/8] Add updateEnr and allow for table constructor usage in
 newProtocol

---
 eth/p2p/discoveryv5/protocol.nim | 15 +++++++++++----
 tests/p2p/discv5_test_helper.nim |  2 +-
 tests/p2p/test_discoveryv5.nim   | 13 ++++++-------
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index e10842e..9b5299d 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -73,7 +73,7 @@
 ## This might be a concern for mobile devices.
 
 import
-  std/[tables, sets, options, math, random], bearssl,
+  std/[tables, sets, options, math, random, sequtils], bearssl,
   stew/shims/net as stewNet, json_serialization/std/net,
   stew/[byteutils, endians2], chronicles, chronos, stint,
   eth/[rlp, keys, async_utils], types, encoding, node, routing_table, enr
@@ -172,6 +172,12 @@ proc nodesDiscovered*(d: Protocol): int {.inline.} = d.routingTable.len
 func privKey*(d: Protocol): lent PrivateKey =
   d.privateKey
 
+proc updateEnr*(
+    d: Protocol, enrFields: openarray[(string, seq[byte])]): DiscResult[void] =
+  let fields = mapIt(enrFields, toFieldPair(it[0], it[1]))
+  d.localNode.record = ? d.localNode.record.update(d.privateKey, fields)
+  ok()
+
 proc send(d: Protocol, a: Address, data: seq[byte]) =
   let ta = initTAddress(a.ip, a.port)
   try:
@@ -691,7 +697,7 @@ proc lookupLoop(d: Protocol) {.async, raises: [Exception, Defect].} =
 
 proc newProtocol*(privKey: PrivateKey, db: Database,
                   externalIp: Option[ValidIpAddress], tcpPort, udpPort: Port,
-                  localEnrFields: openarray[FieldPair] = [],
+                  localEnrFields: openarray[(string, seq[byte])] = [],
                   bootstrapRecords: openarray[Record] = [],
                   previousEnr = none[enr.Record](),
                   bindIp = IPv4_any(), rng = newRng()):
@@ -702,15 +708,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
+    extraFields = mapIt(localEnrFields, toFieldPair(it[0], it[1]))
     # TODO:
     # - Defect as is now or return a result for enr errors?
     # - In case incorrect key, allow for new enr based on new key (new node id)?
     enr = if previousEnr.isSome():
             previousEnr.get().update(privKey, externalIp, tcpPort, udpPort,
-              localEnrFields).expect("Record within size limits and correct key")
+              extraFields).expect("Record within size limits and correct key")
           else:
             enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
-              localEnrFields).expect("Record within size limits")
+              extraFields).expect("Record within size limits")
     node = newNode(enr).expect("Properly initialized record")
 
   # TODO Consider whether this should be a Defect
diff --git a/tests/p2p/discv5_test_helper.nim b/tests/p2p/discv5_test_helper.nim
index 794525b..dbff83e 100644
--- a/tests/p2p/discv5_test_helper.nim
+++ b/tests/p2p/discv5_test_helper.nim
@@ -10,7 +10,7 @@ proc localAddress*(port: int): Address =
 proc initDiscoveryNode*(rng: ref BrHmacDrbgContext, privKey: PrivateKey,
                         address: Address,
                         bootstrapRecords: openarray[Record] = [],
-                        localEnrFields: openarray[FieldPair] = []):
+                        localEnrFields: openarray[(string, seq[byte])] = []):
                         discv5_protocol.Protocol =
   var db = DiscoveryDB.init(newMemoryDB())
   result = newProtocol(privKey, db,
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
index c2a4547..b776904 100644
--- a/tests/p2p/test_discoveryv5.nim
+++ b/tests/p2p/test_discoveryv5.nim
@@ -333,11 +333,12 @@ procSuite "Discovery v5 Tests":
       # resolve in previous test block
       let pong = await targetNode.ping(mainNode.localNode)
       check pong.isOk()
-      # TODO: need to add some logic to update ENRs properly
+
       targetSeqNum.inc()
-      let r = enr.Record.init(targetSeqNum, targetKey,
-        some(targetAddress.ip), targetAddress.port, targetAddress.port)[]
-      targetNode.localNode.record = r
+      # need to add something to get the enr sequence number incremented
+      let update = targetNode.updateEnr({"addsomefield": @[byte 1]})
+      check update.isOk()
+
       let n = await mainNode.resolve(targetId)
       check:
         n.isSome()
@@ -348,9 +349,7 @@ procSuite "Discovery v5 Tests":
     # close targetNode, resolve should lookup, check if we get updated ENR.
     block:
       targetSeqNum.inc()
-      let r = enr.Record.init(targetSeqNum, targetKey, some(targetAddress.ip),
-        targetAddress.port, targetAddress.port)[]
-      targetNode.localNode.record = r
+      let update = targetNode.updateEnr({"addsomefield": @[byte 2]})
 
       # ping node so that its ENR gets added
       check (await targetNode.ping(lookupNode.localNode)).isOk()

From 95a09fdf7f004acf1d848347b107ddc70b54a17a Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Wed, 8 Jul 2020 12:14:00 +0200
Subject: [PATCH 5/8] Get rid of some duplicate code

---
 eth/p2p/discoveryv5/enr.nim | 52 ++++++++++++++++---------------------
 1 file changed, 22 insertions(+), 30 deletions(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index def2606..06d9525 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/eth/p2p/discoveryv5/enr.nim
@@ -76,6 +76,8 @@ proc `==`(a, b: Field): bool =
   else:
     return false
 
+proc cmp(a, b: FieldPair): int = cmp(a[0], b[0])
+
 proc makeEnrRaw(seqNum: uint64, pk: PrivateKey,
     pairs: openarray[FieldPair]): EnrResult[seq[byte]] =
   proc append(w: var RlpWriter, seqNum: uint64,
@@ -118,8 +120,9 @@ proc makeEnrAux(seqNum: uint64, pk: PrivateKey,
     Field(kind: kBytes, bytes: @(pubkey.toRawCompressed()))))
 
   # Sort by key
-  record.pairs.sort() do(a, b: FieldPair) -> int:
-    cmp(a[0], b[0])
+  record.pairs.sort(cmp)
+  # TODO: Should deduplicate on keys here also. Should we error on that or just
+  # deal with it?
 
   record.raw = ? makeEnrRaw(seqNum, pk, record.pairs)
   ok(record)
@@ -140,18 +143,8 @@ macro initRecord*(seqNum: uint64, pk: PrivateKey,
 template toFieldPair*(key: string, value: auto): FieldPair =
   (key, toField(value))
 
-proc init*(T: type Record, seqNum: uint64,
-                           pk: PrivateKey,
-                           ip: Option[ValidIpAddress],
-                           tcpPort, udpPort: Port,
-                           extraFields: openarray[FieldPair] = []):
-                           EnrResult[T] =
-  ## Initialize a `Record` with given sequence number, private key, optional
-  ## ip-address, tcp port, udp port, and optional custom k:v pairs.
-  ##
-  ## Can fail in case the record exceeds the `maxEnrSize`.
-  var fields = newSeq[FieldPair]()
-
+proc addAddress(fields: var seq[FieldPair], ip: Option[ValidIpAddress],
+    tcpPort, udpPort: Port) =
   if ip.isSome():
     let
       ipExt = ip.get()
@@ -165,6 +158,19 @@ proc init*(T: type Record, seqNum: uint64,
     fields.add(("tcp", tcpPort.uint16.toField))
     fields.add(("udp", udpPort.uint16.toField))
 
+proc init*(T: type Record, seqNum: uint64,
+                           pk: PrivateKey,
+                           ip: Option[ValidIpAddress],
+                           tcpPort, udpPort: Port,
+                           extraFields: openarray[FieldPair] = []):
+                           EnrResult[T] =
+  ## Initialize a `Record` with given sequence number, private key, optional
+  ## ip-address, tcp port, udp port, and optional custom k:v pairs.
+  ##
+  ## Can fail in case the record exceeds the `maxEnrSize`.
+  var fields = newSeq[FieldPair]()
+
+  fields.addAddress(ip, tcpPort, udpPort)
   fields.add extraFields
   makeEnrAux(seqNum, pk, fields)
 
@@ -248,8 +254,7 @@ proc update*(record: Record, pk: PrivateKey, fieldPairs: openarray[FieldPair]):
         updated = true
     else:
       # Add new k:v pair.
-      r.pairs.insert(fieldPair,
-        lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0]))
+      r.pairs.insert(fieldPair, lowerBound(r.pairs, fieldPair, cmp))
       updated = true
 
   if updated:
@@ -265,21 +270,8 @@ proc update*(r: Record, pk: PrivateKey,
                         EnrResult[Record] =
   var fields = newSeq[FieldPair]()
 
-  if ip.isSome():
-    let
-      ipExt = ip.get()
-      isV6 = ipExt.family == IPv6
-
-    fields.add(if isV6: ("ip6", ipExt.address_v6.toField)
-               else: ("ip", ipExt.address_v4.toField))
-    fields.add(((if isV6: "tcp6" else: "tcp"), tcpPort.uint16.toField))
-    fields.add(((if isV6: "udp6" else: "udp"), udpPort.uint16.toField))
-  else:
-    fields.add(("tcp", tcpPort.uint16.toField))
-    fields.add(("udp", udpPort.uint16.toField))
-
+  fields.addAddress(ip, tcpPort, udpPort)
   fields.add extraFields
-
   r.update(pk, fields)
 
 proc tryGet*(r: Record, key: string, T: type): Option[T] =

From d3db83fa0a821f3f4a543c9df3f6a9b6f71b2639 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Wed, 8 Jul 2020 13:13:29 +0200
Subject: [PATCH 6/8] Extra test on newProtocol

---
 eth/p2p/discoveryv5/enr.nim      |  2 ++
 eth/p2p/discoveryv5/protocol.nim | 11 +++++++----
 tests/p2p/discv5_test_helper.nim |  6 ++++--
 tests/p2p/test_discoveryv5.nim   | 34 ++++++++++++++++++++++++++++----
 4 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index 06d9525..eb7c2ca 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/eth/p2p/discoveryv5/enr.nim
@@ -258,6 +258,8 @@ proc update*(record: Record, pk: PrivateKey, fieldPairs: openarray[FieldPair]):
       updated = true
 
   if updated:
+    if r.seqNum == high(r.seqNum): # highly unlikely
+      return err("Maximum sequence number reached")
     r.seqNum.inc()
     r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
 
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index 9b5299d..9396247 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -172,7 +172,10 @@ proc nodesDiscovered*(d: Protocol): int {.inline.} = d.routingTable.len
 func privKey*(d: Protocol): lent PrivateKey =
   d.privateKey
 
-proc updateEnr*(
+func getRecord*(d: Protocol): Record =
+  d.localNode.record
+
+proc updateRecord*(
     d: Protocol, enrFields: openarray[(string, seq[byte])]): DiscResult[void] =
   let fields = mapIt(enrFields, toFieldPair(it[0], it[1]))
   d.localNode.record = ? d.localNode.record.update(d.privateKey, fields)
@@ -699,7 +702,7 @@ proc newProtocol*(privKey: PrivateKey, db: Database,
                   externalIp: Option[ValidIpAddress], tcpPort, udpPort: Port,
                   localEnrFields: openarray[(string, seq[byte])] = [],
                   bootstrapRecords: openarray[Record] = [],
-                  previousEnr = none[enr.Record](),
+                  previousRecord = none[enr.Record](),
                   bindIp = IPv4_any(), rng = newRng()):
                   Protocol {.raises: [Defect].} =
   # TODO: Tried adding bindPort = udpPort as parameter but that gave
@@ -712,8 +715,8 @@ proc newProtocol*(privKey: PrivateKey, db: Database,
     # TODO:
     # - Defect as is now or return a result for enr errors?
     # - In case incorrect key, allow for new enr based on new key (new node id)?
-    enr = if previousEnr.isSome():
-            previousEnr.get().update(privKey, externalIp, tcpPort, udpPort,
+    enr = if previousRecord.isSome():
+            previousRecord.get().update(privKey, externalIp, tcpPort, udpPort,
               extraFields).expect("Record within size limits and correct key")
           else:
             enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
diff --git a/tests/p2p/discv5_test_helper.nim b/tests/p2p/discv5_test_helper.nim
index dbff83e..cabde2f 100644
--- a/tests/p2p/discv5_test_helper.nim
+++ b/tests/p2p/discv5_test_helper.nim
@@ -10,14 +10,16 @@ proc localAddress*(port: int): Address =
 proc initDiscoveryNode*(rng: ref BrHmacDrbgContext, privKey: PrivateKey,
                         address: Address,
                         bootstrapRecords: openarray[Record] = [],
-                        localEnrFields: openarray[(string, seq[byte])] = []):
+                        localEnrFields: openarray[(string, seq[byte])] = [],
+                        previousRecord = none[enr.Record]()):
                         discv5_protocol.Protocol =
   var db = DiscoveryDB.init(newMemoryDB())
   result = newProtocol(privKey, db,
                        some(address.ip),
                        address.port, address.port,
                        bootstrapRecords = bootstrapRecords,
-                       localEnrFields = localEnrFields, rng = rng)
+                       localEnrFields = localEnrFields,
+                       previousRecord = previousRecord, rng = rng)
 
   result.open()
 
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
index b776904..369931e 100644
--- a/tests/p2p/test_discoveryv5.nim
+++ b/tests/p2p/test_discoveryv5.nim
@@ -1,7 +1,7 @@
 import
   chronos, chronicles, tables, stint, testutils/unittests,
-  stew/shims/net, eth/keys, bearssl,
-  eth/p2p/discoveryv5/[enr, node, types, routing_table, encoding],
+  stew/shims/net, eth/[keys, trie/db], bearssl,
+  eth/p2p/discoveryv5/[enr, node, types, routing_table, encoding, discovery_db],
   eth/p2p/discoveryv5/protocol as discv5_protocol,
   ./discv5_test_helper
 
@@ -336,7 +336,7 @@ procSuite "Discovery v5 Tests":
 
       targetSeqNum.inc()
       # need to add something to get the enr sequence number incremented
-      let update = targetNode.updateEnr({"addsomefield": @[byte 1]})
+      let update = targetNode.updateRecord({"addsomefield": @[byte 1]})
       check update.isOk()
 
       let n = await mainNode.resolve(targetId)
@@ -349,7 +349,7 @@ procSuite "Discovery v5 Tests":
     # close targetNode, resolve should lookup, check if we get updated ENR.
     block:
       targetSeqNum.inc()
-      let update = targetNode.updateEnr({"addsomefield": @[byte 2]})
+      let update = targetNode.updateRecord({"addsomefield": @[byte 2]})
 
       # ping node so that its ENR gets added
       check (await targetNode.ping(lookupNode.localNode)).isOk()
@@ -392,3 +392,29 @@ procSuite "Discovery v5 Tests":
     check discoveredFiltered.len == 1 and discoveredFiltered.contains(targetNode)
 
     await lookupNode.closeWait()
+
+  test "New protocol with enr":
+    let
+      privKey = PrivateKey.random(rng[])
+      ip = some(ValidIpAddress.init("127.0.0.1"))
+      port = Port(20301)
+      db = DiscoveryDB.init(newMemoryDB())
+      node = newProtocol(privKey, db, ip, port, port, rng = rng)
+      noUpdatesNode = newProtocol(privKey, db, ip, port, port, rng = rng,
+        previousRecord = some(node.getRecord()))
+      updatesNode = newProtocol(privKey, db, ip, port, Port(20302), rng = rng,
+        previousRecord = some(noUpdatesNode.getRecord()))
+      moreUpdatesNode = newProtocol(privKey, db, ip, port, port, rng = rng,
+        localEnrFields = {"addfield": @[byte 0]},
+        previousRecord = some(updatesNode.getRecord()))
+    check:
+      node.getRecord().seqNum == 1
+      noUpdatesNode.getRecord().seqNum == 1
+      updatesNode.getRecord().seqNum == 2
+      moreUpdatesNode.getRecord().seqNum == 3
+
+    # Defect (for now?) on incorrect key use
+    expect ResultDefect:
+      let incorrectKeyUpdates = newProtocol(PrivateKey.random(rng[]),
+        db, ip, port, port, rng = rng,
+        previousRecord = some(updatesNode.getRecord()))

From 4f3df3c9b233866d5dc2db7425106e4ebacd38e9 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Wed, 8 Jul 2020 14:23:43 +0200
Subject: [PATCH 7/8] Change update back to use var parameter of Record

---
 eth/p2p/discoveryv5/enr.nim      | 17 +++++++-------
 eth/p2p/discoveryv5/protocol.nim | 30 ++++++++++++------------
 tests/p2p/test_enr.nim           | 39 ++++++++++++++------------------
 3 files changed, 42 insertions(+), 44 deletions(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index eb7c2ca..95d50c2 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/eth/p2p/discoveryv5/enr.nim
@@ -233,8 +233,8 @@ proc find(r: Record, key: string): Option[int] =
     if k == key:
       return some(i)
 
-proc update*(record: Record, pk: PrivateKey, fieldPairs: openarray[FieldPair]):
-    EnrResult[Record] =
+proc update*(record: var Record, pk: PrivateKey,
+    fieldPairs: openarray[FieldPair]): EnrResult[void] =
   var r = record
 
   let pubkey = r.get(PublicKey)
@@ -262,14 +262,15 @@ proc update*(record: Record, pk: PrivateKey, fieldPairs: openarray[FieldPair]):
       return err("Maximum sequence number reached")
     r.seqNum.inc()
     r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
+    record = r
 
-  ok(r)
+  ok()
 
-proc update*(r: Record, pk: PrivateKey,
-                        ip: Option[ValidIpAddress],
-                        tcpPort, udpPort: Port,
-                        extraFields: openarray[FieldPair] = []):
-                        EnrResult[Record] =
+proc update*(r: var Record, pk: PrivateKey,
+                            ip: Option[ValidIpAddress],
+                            tcpPort, udpPort: Port,
+                            extraFields: openarray[FieldPair] = []):
+                            EnrResult[void] =
   var fields = newSeq[FieldPair]()
 
   fields.addAddress(ip, tcpPort, udpPort)
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index 9396247..8c78f9c 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -178,8 +178,9 @@ func getRecord*(d: Protocol): Record =
 proc updateRecord*(
     d: Protocol, enrFields: openarray[(string, seq[byte])]): DiscResult[void] =
   let fields = mapIt(enrFields, toFieldPair(it[0], it[1]))
-  d.localNode.record = ? d.localNode.record.update(d.privateKey, fields)
-  ok()
+  d.localNode.record.update(d.privateKey, fields)
+  # TODO: Would it make sense to actively ping ("broadcast") to all the peers
+  # we stored a handshake with in order to get that ENR updated?
 
 proc send(d: Protocol, a: Address, data: seq[byte]) =
   let ta = initTAddress(a.ip, a.port)
@@ -710,18 +711,19 @@ proc newProtocol*(privKey: PrivateKey, db: Database,
   # Anyhow, nim-beacon-chain would also require some changes to support port
   # remapping through NAT and this API is also subject to change once we
   # introduce support for ipv4 + ipv6 binding/listening.
-  let
-    extraFields = mapIt(localEnrFields, toFieldPair(it[0], it[1]))
-    # TODO:
-    # - Defect as is now or return a result for enr errors?
-    # - In case incorrect key, allow for new enr based on new key (new node id)?
-    enr = if previousRecord.isSome():
-            previousRecord.get().update(privKey, externalIp, tcpPort, udpPort,
-              extraFields).expect("Record within size limits and correct key")
-          else:
-            enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
-              extraFields).expect("Record within size limits")
-    node = newNode(enr).expect("Properly initialized record")
+  let extraFields = mapIt(localEnrFields, toFieldPair(it[0], it[1]))
+  # TODO:
+  # - Defect as is now or return a result for enr errors?
+  # - In case incorrect key, allow for new enr based on new key (new node id)?
+  var record: Record
+  if previousRecord.isSome():
+    record = previousRecord.get()
+    record.update(privKey, externalIp, tcpPort, udpPort,
+      extraFields).expect("Record within size limits and correct key")
+  else:
+    record = enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
+     extraFields).expect("Record within size limits")
+  let node = newNode(record).expect("Properly initialized record")
 
   # TODO Consider whether this should be a Defect
   doAssert rng != nil, "RNG initialization failed"
diff --git a/tests/p2p/test_enr.nim b/tests/p2p/test_enr.nim
index c0936db..ac1927d 100644
--- a/tests/p2p/test_enr.nim
+++ b/tests/p2p/test_enr.nim
@@ -95,26 +95,23 @@ 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.update(pk, [newField])
-      check res.isOk()
-      r = res[]
+      let updated = r.update(pk, [newField])
+      check updated.isOk()
       check:
         r.get("test", uint) == 123
         r.seqNum == 2
 
     block: # Insert same k:v pair, no update of seqNum should occur.
-      let res = r.update(pk, [newField])
-      check res.isOk()
-      r = res[]
+      let updated = r.update(pk, [newField])
+      check updated.isOk()
       check:
         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.update(pk, [updatedField])
-      check res.isOk()
-      r = res[]
+      let updated = r.update(pk, [updatedField])
+      check updated.isOk()
       check:
         r.get("test", uint) == 1234
         r.seqNum == 3
@@ -130,9 +127,9 @@ suite "ENR":
 
     let newField = toFieldPair("test", 123'u)
     let newField2 = toFieldPair("zzz", 123'u)
-    let res = r.update(pk, [newField, newField2])
-    check res.isOk()
-    check $res[] == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00, zzz: 123)"""
+    let updated = r.update(pk, [newField, newField2])
+    check updated.isOk()
+    check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00, zzz: 123)"""
 
   test "ENR update size too big":
     let pk = PrivateKey.fromHex(
@@ -142,8 +139,8 @@ suite "ENR":
     check r.isOk()
 
     let newField = toFieldPair("test", 123'u)
-    let res = r[].update(pk, [newField])
-    check res.isErr()
+    let updated = r[].update(pk, [newField])
+    check updated.isErr()
 
   test "ENR update invalid key":
     let pk = PrivateKey.fromHex(
@@ -155,8 +152,8 @@ suite "ENR":
     let
       wrongPk = PrivateKey.random(rng[])
       newField = toFieldPair("test", 123'u)
-      res = r[].update(wrongPk, [newField])
-    check res.isErr()
+      updated = r[].update(wrongPk, [newField])
+    check updated.isErr()
 
   test "ENR update address":
     let
@@ -165,18 +162,16 @@ suite "ENR":
     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[]
+      let updated = r.update(pk, none(ValidIpAddress), Port(9000), Port(9000))
+      check updated.isOk()
       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[]
+      let updated = r.update(pk, none(ValidIpAddress), Port(9001), Port(9002))
+      check updated.isOk()
       check:
         r.get("tcp", uint) == 9001
         r.get("udp", uint) == 9002

From eeb958e83446be96c7dce0275975702a314cf377 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Wed, 8 Jul 2020 14:56:56 +0200
Subject: [PATCH 8/8] Add comments [skip ci]

---
 eth/p2p/discoveryv5/enr.nim | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index 95d50c2..0d4a177 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/eth/p2p/discoveryv5/enr.nim
@@ -165,7 +165,7 @@ proc init*(T: type Record, seqNum: uint64,
                            extraFields: openarray[FieldPair] = []):
                            EnrResult[T] =
   ## Initialize a `Record` with given sequence number, private key, optional
-  ## ip-address, tcp port, udp port, and optional custom k:v pairs.
+  ## ip address, tcp port, udp port, and optional custom k:v pairs.
   ##
   ## Can fail in case the record exceeds the `maxEnrSize`.
   var fields = newSeq[FieldPair]()
@@ -235,6 +235,14 @@ proc find(r: Record, key: string): Option[int] =
 
 proc update*(record: var Record, pk: PrivateKey,
     fieldPairs: openarray[FieldPair]): EnrResult[void] =
+  ## Update a `Record` k:v pairs.
+  ##
+  ## In case any of the k:v pairs is updated or added (new), the sequence number
+  ## of the `Record` will be incremented and a new signature will be applied.
+  ##
+  ## Can fail in case of wrong `PrivateKey`, if the size of the resulting record
+  ## exceeds `maxEnrSize` or if maximum sequence number is reached. The `Record`
+  ## will not be altered in these cases.
   var r = record
 
   let pubkey = r.get(PublicKey)
@@ -271,6 +279,15 @@ proc update*(r: var Record, pk: PrivateKey,
                             tcpPort, udpPort: Port,
                             extraFields: openarray[FieldPair] = []):
                             EnrResult[void] =
+  ## Update a `Record` with given ip address, tcp port, udp port and optional
+  ## custom k:v pairs.
+  ##
+  ## In case any of the k:v pairs is updated or added (new), the sequence number
+  ## of the `Record` will be incremented and a new signature will be applied.
+  ##
+  ## Can fail in case of wrong `PrivateKey`, if the size of the resulting record
+  ## exceeds `maxEnrSize` or if maximum sequence number is reached. The `Record`
+  ## will not be altered in these cases.
   var fields = newSeq[FieldPair]()
 
   fields.addAddress(ip, tcpPort, udpPort)