Allow for multiple enr fields to be inserted + enr update proc

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:

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()