diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim index fad276d..b5508bb 100644 --- a/eth/p2p/discoveryv5/protocol.nim +++ b/eth/p2p/discoveryv5/protocol.nim @@ -116,6 +116,10 @@ const ## call type + DiscoveryConfig* = object + tableIpLimits*: TableIpLimits + bitsPerHop*: int + Protocol* = ref object transp: DatagramTransport localNode*: Node @@ -148,6 +152,11 @@ type DiscResult*[T] = Result[T, cstring] +const + defaultDiscoveryConfig* = DiscoveryConfig( + tableIpLimits: DefaultTableIpLimits, + bitsPerHop: DefaultBitsPerHop) + proc addNode*(d: Protocol, node: Node): bool = ## Add `Node` to discovery routing table. ## @@ -880,18 +889,32 @@ proc ipMajorityLoop(d: Protocol) {.async.} = except CancelledError: trace "ipMajorityLoop canceled" -proc newProtocol*(privKey: PrivateKey, - enrIp: Option[ValidIpAddress], - enrTcpPort, enrUdpPort: Option[Port], - localEnrFields: openArray[(string, seq[byte])] = [], - bootstrapRecords: openArray[Record] = [], - previousRecord = none[enr.Record](), - bindPort: Port, - bindIp = IPv4_any(), - enrAutoUpdate = false, - tableIpLimits = DefaultTableIpLimits, - rng = newRng()): - Protocol = +func init*( + T: type DiscoveryConfig, + tableIpLimit: uint, + bucketIpLimit: uint, + bitsPerHop: int): T = + + DiscoveryConfig( + tableIpLimits: TableIpLimits( + tableIpLimit: tableIpLimit, + bucketIpLimit: bucketIpLimit), + bitsPerHop: bitsPerHop + ) + +proc newProtocol*( + privKey: PrivateKey, + enrIp: Option[ValidIpAddress], + enrTcpPort, enrUdpPort: Option[Port], + localEnrFields: openArray[(string, seq[byte])] = [], + bootstrapRecords: openArray[Record] = [], + previousRecord = none[enr.Record](), + bindPort: Port, + bindIp = IPv4_any(), + enrAutoUpdate = false, + config = defaultDiscoveryConfig, + rng = newRng()): + Protocol = # TODO: Tried adding bindPort = udpPort as parameter but that gave # "Error: internal error: environment misses: udpPort" in nim-beacon-chain. # Anyhow, nim-beacon-chain would also require some changes to support port @@ -933,7 +956,8 @@ proc newProtocol*(privKey: PrivateKey, bootstrapRecords: @bootstrapRecords, ipVote: IpVote.init(), enrAutoUpdate: enrAutoUpdate, - routingTable: RoutingTable.init(node, DefaultBitsPerHop, tableIpLimits, rng), + routingTable: RoutingTable.init( + node, config.bitsPerHop, config.tableIpLimits, rng), rng: rng) template listeningAddress*(p: Protocol): Address = diff --git a/tests/p2p/discv5_test_helper.nim b/tests/p2p/discv5_test_helper.nim index 3ddad9b..397bed5 100644 --- a/tests/p2p/discv5_test_helper.nim +++ b/tests/p2p/discv5_test_helper.nim @@ -9,26 +9,31 @@ export net 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, - bootstrapRecords: openArray[Record] = [], - localEnrFields: openArray[(string, seq[byte])] = [], - previousRecord = none[enr.Record]()): - discv5_protocol.Protocol = +proc initDiscoveryNode*( + rng: ref BrHmacDrbgContext, + privKey: PrivateKey, + address: Address, + bootstrapRecords: openArray[Record] = [], + localEnrFields: openArray[(string, seq[byte])] = [], + previousRecord = none[enr.Record]()): + discv5_protocol.Protocol = # set bucketIpLimit to allow bucket split - let tableIpLimits = TableIpLimits(tableIpLimit: 1000, bucketIpLimit: 24) + let config = DiscoveryConfig.init(1000, 24, 5) - result = newProtocol(privKey, - some(address.ip), - some(address.port), some(address.port), - bindPort = address.port, - bootstrapRecords = bootstrapRecords, - localEnrFields = localEnrFields, - previousRecord = previousRecord, - tableIpLimits = tableIpLimits, - rng = rng) + let protocol = newProtocol( + privKey, + some(address.ip), + some(address.port), some(address.port), + bindPort = address.port, + bootstrapRecords = bootstrapRecords, + localEnrFields = localEnrFields, + previousRecord = previousRecord, + config = config, + rng = rng) - result.open() + protocol.open() + + protocol proc nodeIdInNodes*(id: NodeId, nodes: openArray[Node]): bool = for n in nodes: diff --git a/tests/utp/test_discv5_protocol.nim b/tests/utp/test_discv5_protocol.nim index 3296699..3c238e3 100644 --- a/tests/utp/test_discv5_protocol.nim +++ b/tests/utp/test_discv5_protocol.nim @@ -14,31 +14,8 @@ import ../../eth/p2p/discoveryv5/[enr, node, routing_table], ../../eth/p2p/discoveryv5/protocol as discv5_protocol, ../../eth/utp/utp_discv5_protocol, - ../../eth/keys - -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, - bootstrapRecords: openArray[Record] = [], - localEnrFields: openArray[(string, seq[byte])] = [], - previousRecord = none[enr.Record]()): discv5_protocol.Protocol = - # set bucketIpLimit to allow bucket split - let tableIpLimits = TableIpLimits(tableIpLimit: 1000, bucketIpLimit: 24) - - result = newProtocol(privKey, - some(address.ip), - some(address.port), some(address.port), - bindPort = address.port, - bootstrapRecords = bootstrapRecords, - localEnrFields = localEnrFields, - previousRecord = previousRecord, - tableIpLimits = tableIpLimits, - rng = rng) - - result.open() + ../../eth/keys, + ../p2p/discv5_test_helper proc generateByteArray(rng: var BrHmacDrbgContext, length: int): seq[byte] = var bytes = newSeq[byte](length)