Add record size check on init + add insertFieldPair call + tests

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
+  maxEnrSize = 300  ## Maximum size of an encoded node record, in bytes.
-  minRlpListLen = 4 # for signature, seqId, "id" key, id
+  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,
+  record.raw = ? makeEnrRaw(seqNum, pk, record.pairs)
-      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)
-
   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)

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

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[])
+      keypair = KeyPair.random(rng[])
-      enr = Record.init(100, keys.seckey, none(ValidIpAddress), Port(9000), Port(9000))[]
+      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()