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