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