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