Add updateEnr and allow for table constructor usage in newProtocol

This commit is contained in:
kdeme 2020-07-08 11:45:58 +02:00
parent 72420d7f17
commit 0fb21e72d7
No known key found for this signature in database
GPG Key ID: 4E8DD21420AF43F5
3 changed files with 18 additions and 12 deletions

View File

@ -73,7 +73,7 @@
## This might be a concern for mobile devices. ## This might be a concern for mobile devices.
import 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/shims/net as stewNet, json_serialization/std/net,
stew/[byteutils, endians2], chronicles, chronos, stint, stew/[byteutils, endians2], chronicles, chronos, stint,
eth/[rlp, keys, async_utils], types, encoding, node, routing_table, enr 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 = func privKey*(d: Protocol): lent PrivateKey =
d.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]) = proc send(d: Protocol, a: Address, data: seq[byte]) =
let ta = initTAddress(a.ip, a.port) let ta = initTAddress(a.ip, a.port)
try: try:
@ -691,7 +697,7 @@ proc lookupLoop(d: Protocol) {.async, raises: [Exception, Defect].} =
proc newProtocol*(privKey: PrivateKey, db: Database, proc newProtocol*(privKey: PrivateKey, db: Database,
externalIp: Option[ValidIpAddress], tcpPort, udpPort: Port, externalIp: Option[ValidIpAddress], tcpPort, udpPort: Port,
localEnrFields: openarray[FieldPair] = [], localEnrFields: openarray[(string, seq[byte])] = [],
bootstrapRecords: openarray[Record] = [], bootstrapRecords: openarray[Record] = [],
previousEnr = none[enr.Record](), previousEnr = none[enr.Record](),
bindIp = IPv4_any(), rng = newRng()): 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 # remapping through NAT and this API is also subject to change once we
# introduce support for ipv4 + ipv6 binding/listening. # introduce support for ipv4 + ipv6 binding/listening.
let let
extraFields = mapIt(localEnrFields, toFieldPair(it[0], it[1]))
# TODO: # TODO:
# - Defect as is now or return a result for enr errors? # - 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)? # - In case incorrect key, allow for new enr based on new key (new node id)?
enr = if previousEnr.isSome(): enr = if previousEnr.isSome():
previousEnr.get().update(privKey, externalIp, tcpPort, udpPort, 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: else:
enr.Record.init(1, privKey, externalIp, tcpPort, udpPort, 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") node = newNode(enr).expect("Properly initialized record")
# TODO Consider whether this should be a Defect # TODO Consider whether this should be a Defect

View File

@ -10,7 +10,7 @@ proc localAddress*(port: int): Address =
proc initDiscoveryNode*(rng: ref BrHmacDrbgContext, privKey: PrivateKey, proc initDiscoveryNode*(rng: ref BrHmacDrbgContext, privKey: PrivateKey,
address: Address, address: Address,
bootstrapRecords: openarray[Record] = [], bootstrapRecords: openarray[Record] = [],
localEnrFields: openarray[FieldPair] = []): localEnrFields: openarray[(string, seq[byte])] = []):
discv5_protocol.Protocol = discv5_protocol.Protocol =
var db = DiscoveryDB.init(newMemoryDB()) var db = DiscoveryDB.init(newMemoryDB())
result = newProtocol(privKey, db, result = newProtocol(privKey, db,

View File

@ -333,11 +333,12 @@ procSuite "Discovery v5 Tests":
# resolve in previous test block # resolve in previous test block
let pong = await targetNode.ping(mainNode.localNode) let pong = await targetNode.ping(mainNode.localNode)
check pong.isOk() check pong.isOk()
# TODO: need to add some logic to update ENRs properly
targetSeqNum.inc() targetSeqNum.inc()
let r = enr.Record.init(targetSeqNum, targetKey, # need to add something to get the enr sequence number incremented
some(targetAddress.ip), targetAddress.port, targetAddress.port)[] let update = targetNode.updateEnr({"addsomefield": @[byte 1]})
targetNode.localNode.record = r check update.isOk()
let n = await mainNode.resolve(targetId) let n = await mainNode.resolve(targetId)
check: check:
n.isSome() n.isSome()
@ -348,9 +349,7 @@ procSuite "Discovery v5 Tests":
# close targetNode, resolve should lookup, check if we get updated ENR. # close targetNode, resolve should lookup, check if we get updated ENR.
block: block:
targetSeqNum.inc() targetSeqNum.inc()
let r = enr.Record.init(targetSeqNum, targetKey, some(targetAddress.ip), let update = targetNode.updateEnr({"addsomefield": @[byte 2]})
targetAddress.port, targetAddress.port)[]
targetNode.localNode.record = r
# ping node so that its ENR gets added # ping node so that its ENR gets added
check (await targetNode.ping(lookupNode.localNode)).isOk() check (await targetNode.ping(lookupNode.localNode)).isOk()