Some discv5 cleanup / refactor (#326)

* Use Address in HandshakeKey

* Refactor + comments + dcli use queryRandom

* Rename types.nim to messages.nim

eth/p2p/discoveryv5/dcli.nim

@@ -166,6 +166,12 @@ proc setupNat(conf: DiscoveryConf): tuple[ip: Option[ValidIpAddress],
       if extPorts.isSome:
         (result.tcpPort, result.udpPort) = extPorts.get()
 
+proc discover(d: protocol.Protocol) {.async.} =
+  while true:
+    let discovered = await d.queryRandom()
+    info "Lookup finished", nodes = discovered.len
+    await sleepAsync(30.seconds)
+
 proc run(config: DiscoveryConf) =
   let
     (ip, tcpPort, udpPort) = setupNat(config)
@@ -206,7 +212,7 @@ proc run(config: DiscoveryConf) =
       echo "No Talk Response message returned"
   of noCommand:
     d.start()
-    runForever()
+    waitfor(discover(d))
 
 when isMainModule:
   let config = DiscoveryConf.load()

eth/p2p/discoveryv5/encoding.nim

@@ -1,7 +1,12 @@
+## Discovery v5 packet encoding as specified at
+## https://github.com/ethereum/devp2p/blob/master/discv5/discv5-wire.md#packet-encoding
+## And handshake/sessions as specified at
+## https://github.com/ethereum/devp2p/blob/master/discv5/discv5-theory.md#sessions
+##
 import
-  std/[tables, options],
+  std/[tables, options, hashes, net],
   nimcrypto, stint, chronicles, bearssl, stew/[results, byteutils],
-  eth/[rlp, keys], types, node, enr, hkdf, sessions
+  eth/[rlp, keys], messages, node, enr, hkdf, sessions
 
 from stew/objects import checkedEnumAssign
 
@@ -68,6 +73,10 @@ type
       node*: Option[Node]
       srcIdHs*: NodeId
 
+  HandshakeKey* = object
+    nodeId*: NodeId
+    address*: Address
+
   Codec* = object
     localNode*: Node
     privKey*: PrivateKey
@@ -76,6 +85,13 @@ type
 
   DecodeResult*[T] = Result[T, cstring]
 
+func `==`*(a, b: HandshakeKey): bool =
+  (a.nodeId == b.nodeId) and (a.address == b.address)
+
+func hash*(key: HandshakeKey): Hash =
+  result = key.nodeId.hash !& key.address.hash
+  result = !$result
+
 proc idHash(challengeData, ephkey: openarray[byte], nodeId: NodeId):
     MDigest[256] =
   var ctx: sha256
@@ -238,7 +254,7 @@ proc encodeWhoareyouPacket*(rng: var BrHmacDrbgContext, c: var Codec,
       recordSeq: recordSeq,
       challengeData: @iv & header)
     challenge = Challenge(whoareyouData: whoareyouData, pubkey: pubkey)
-    key = HandShakeKey(nodeId: toId, address: $toAddr)
+    key = HandShakeKey(nodeId: toId, address: toAddr)
 
   c.handshakes[key] = challenge
 
@@ -455,7 +471,7 @@ proc decodeHandshakePacket(c: var Codec, fromAddr: Address, nonce: AESGCMNonce,
   if header.len < staticHeaderSize + authdataHeadSize + int(sigSize) + int(ephKeySize):
     return err("Invalid header for handshake message packet")
 
-  let key = HandShakeKey(nodeId: srcId, address: $fromAddr)
+  let key = HandShakeKey(nodeId: srcId, address: fromAddr)
   var challenge: Challenge
   if not c.handshakes.pop(key, challenge):
     return err("No challenge found: timed out or unsolicited packet")

eth/p2p/discoveryv5/lru.nim

@@ -8,10 +8,10 @@ type
     table: Table[K, DoublyLinkedNode[(K, V)]] # DoublyLinkedNode is alraedy ref
     capacity: int
 
-proc init*[K, V](T: type LRUCache[K, V], capacity: int): LRUCache[K, V] =
+func init*[K, V](T: type LRUCache[K, V], capacity: int): LRUCache[K, V] =
   LRUCache[K, V](capacity: capacity) # Table and list init is done default
 
-proc get*[K, V](lru: var LRUCache[K, V], key: K): Option[V] =
+func get*[K, V](lru: var LRUCache[K, V], key: K): Option[V] =
   let node = lru.table.getOrDefault(key, nil)
   if node.isNil:
     return none(V)
@@ -20,7 +20,7 @@ proc get*[K, V](lru: var LRUCache[K, V], key: K): Option[V] =
   lru.list.prepend(node)
   return some(node.value[1])
 
-proc put*[K, V](lru: var LRUCache[K, V], key: K, value: V) =
+func put*[K, V](lru: var LRUCache[K, V], key: K, value: V) =
   let node = lru.table.getOrDefault(key, nil)
   if not node.isNil:
     lru.list.remove(node)
@@ -32,10 +32,10 @@ proc put*[K, V](lru: var LRUCache[K, V], key: K, value: V) =
   lru.list.prepend((key, value))
   lru.table[key] = lru.list.head
 
-proc del*[K, V](lru: var LRUCache[K, V], key: K) =
+func del*[K, V](lru: var LRUCache[K, V], key: K) =
   var node: DoublyLinkedNode[(K, V)]
   if lru.table.pop(key, node):
     lru.list.remove(node)
 
-proc len*[K, V](lru: LRUCache[K, V]): int =
+func len*[K, V](lru: LRUCache[K, V]): int =
   lru.table.len

eth/p2p/discoveryv5/messages.nim

@@ -1,20 +1,15 @@
+## Discovery v5 Protocol Messages as specified at
+## https://github.com/ethereum/devp2p/blob/master/discv5/discv5-wire.md#protocol-messages
+## These messages get RLP encoded.
+##
 import
   std/[hashes, net],
   stew/arrayops,
-  eth/rlp, enr, node
+  eth/rlp, enr
 
 {.push raises: [Defect].}
 
-const
-  aesKeySize* = 128 div 8
-
 type
-  AesKey* = array[aesKeySize, byte]
-
-  HandshakeKey* = object
-    nodeId*: NodeId
-    address*: string # TODO: Replace with Address, need hash
-
   MessageKind* = enum
     # TODO This is needed only to make Nim 1.2.6 happy
     #      Without it, the `MessageKind` type cannot be used as
@@ -139,7 +134,3 @@ proc append*(writer: var RlpWriter, ip: IpAddress) =
 
 proc hash*(reqId: RequestId): Hash =
   hash(reqId.id)
-
-proc hash*(key: HandshakeKey): Hash =
-  result = key.nodeId.hash !& key.address.hash
-  result = !$result

eth/p2p/discoveryv5/node.nim

@@ -48,7 +48,7 @@ func newNode*(r: Record): Result[Node, cstring] =
     ok(Node(id: pk.get().toNodeId(), pubkey: pk.get(), record: r,
        address: none(Address)))
 
-proc update*(n: Node, pk: PrivateKey, ip: Option[ValidIpAddress],
+func update*(n: Node, pk: PrivateKey, ip: Option[ValidIpAddress],
     tcpPort, udpPort: Option[Port] = none[Port](),
     extraFields: openarray[FieldPair] = []): Result[void, cstring] =
   ? n.record.update(pk, ip, tcpPort, udpPort, extraFields)

eth/p2p/discoveryv5/protocol.nim

@@ -77,7 +77,7 @@ import
   stew/shims/net as stewNet, json_serialization/std/net,
   stew/endians2, chronicles, chronos, stint, bearssl, metrics,
   eth/[rlp, keys, async_utils],
-  types, encoding, node, routing_table, enr, random2, sessions, ip_vote
+  messages, encoding, node, routing_table, enr, random2, sessions, ip_vote
 
 import nimcrypto except toHex
 
@@ -341,7 +341,7 @@ proc handleMessage(d: Protocol, srcId: NodeId, fromAddr: Address,
 
 proc sendWhoareyou(d: Protocol, toId: NodeId, a: Address,
     requestNonce: AESGCMNonce, node: Option[Node]) {.raises: [Exception].} =
-  let key = HandShakeKey(nodeId: toId, address: $a)
+  let key = HandShakeKey(nodeId: toId, address: a)
   if not d.codec.hasHandshake(key):
     let
       recordSeq = if node.isSome(): node.get().record.seqNum

eth/p2p/discoveryv5/sessions.nim

@@ -1,22 +1,28 @@
+## Session cache as mentioned at
+## https://github.com/ethereum/devp2p/blob/master/discv5/discv5-theory.md#session-cache
+##
 import
   std/options,
   stint, stew/endians2, stew/shims/net,
-  types, node, lru
+  node, lru
 
 export lru
 
 {.push raises: [Defect].}
 
-const keySize = sizeof(NodeId) +
-                16 + # max size of ip address (ipv6)
-                2 # Sizeof port
+const
+  aesKeySize* = 128 div 8
+  keySize = sizeof(NodeId) +
+            16 + # max size of ip address (ipv6)
+            2 # Sizeof port
 
 type
+  AesKey* = array[aesKeySize, byte]
   SessionKey* = array[keySize, byte]
   SessionValue* = array[sizeof(AesKey) + sizeof(AesKey), byte]
   Sessions* = LRUCache[SessionKey, SessionValue]
 
-proc makeKey(id: NodeId, address: Address): SessionKey =
+func makeKey(id: NodeId, address: Address): SessionKey =
   var pos = 0
   result[pos ..< pos+sizeof(id)] = toBytes(id)
   pos.inc(sizeof(id))
@@ -28,13 +34,13 @@ proc makeKey(id: NodeId, address: Address): SessionKey =
   pos.inc(sizeof(address.ip.address_v6))
   result[pos ..< pos+sizeof(address.port)] = toBytes(address.port.uint16)
 
-proc store*(s: var Sessions, id: NodeId, address: Address, r, w: AesKey) =
+func store*(s: var Sessions, id: NodeId, address: Address, r, w: AesKey) =
   var value: array[sizeof(r) + sizeof(w), byte]
   value[0 .. 15] = r
   value[16 .. ^1] = w
   s.put(makeKey(id, address), value)
 
-proc load*(s: var Sessions, id: NodeId, address: Address, r, w: var AesKey): bool =
+func load*(s: var Sessions, id: NodeId, address: Address, r, w: var AesKey): bool =
   let res = s.get(makeKey(id, address))
   if res.isSome():
     let val = res.get()
@@ -44,5 +50,5 @@ proc load*(s: var Sessions, id: NodeId, address: Address, r, w: var AesKey): boo
   else:
     return false
 
-proc del*(s: var Sessions, id: NodeId, address: Address) =
+func del*(s: var Sessions, id: NodeId, address: Address) =
   s.del(makeKey(id, address))

tests/fuzzing/discoveryv5/fuzz_decode_message.nim

@@ -1,6 +1,6 @@
 import
   testutils/fuzzing, stew/byteutils,
-  eth/rlp, eth/p2p/discoveryv5/[encoding, types]
+  eth/rlp, eth/p2p/discoveryv5/[encoding, messages]
 
 test:
   block:

tests/p2p/test_discoveryv5.nim

@@ -2,7 +2,7 @@ import
   std/tables,
   chronos, chronicles, stint, testutils/unittests,
   stew/shims/net, eth/keys, bearssl,
-  eth/p2p/discoveryv5/[enr, node, routing_table, encoding, sessions, types],
+  eth/p2p/discoveryv5/[enr, node, routing_table, encoding, sessions, messages],
   eth/p2p/discoveryv5/protocol as discv5_protocol,
   ./discv5_test_helper
 
@@ -610,7 +610,7 @@ procSuite "Discovery v5 Tests":
     # Check handshake duplicates
     check receiveNode.codec.handshakes.len == 1
     # Check if it is for the first packet that a handshake is stored
-    let key = HandShakeKey(nodeId: sendNode.id, address: $a)
+    let key = HandShakeKey(nodeId: sendNode.id, address: a)
     check receiveNode.codec.handshakes[key].whoareyouData.requestNonce ==
       firstRequestNonce
 

tests/p2p/test_discoveryv5_encoding.nim

@@ -2,7 +2,7 @@ import
   std/[unittest, options, sequtils, tables],
   stint, stew/byteutils, stew/shims/net,
   eth/[rlp, keys],
-  eth/p2p/discoveryv5/[types, encoding, enr, node, sessions]
+  eth/p2p/discoveryv5/[messages, encoding, enr, node, sessions]
 
 let rng = newRng()
 
@@ -310,7 +310,7 @@ suite "Discovery v5.1 Packet Encodings Test Vectors":
         challengeData: hexToSeqByte(whoareyouChallengeData))
       pubkey = some(privKeyA.toPublicKey())
       challenge = Challenge(whoareyouData: whoareyouData, pubkey: pubkey)
-      key = HandShakeKey(nodeId: nodeA.id, address: $(nodeA.address.get()))
+      key = HandShakeKey(nodeId: nodeA.id, address: nodeA.address.get())
 
     check: not codecB.handshakes.hasKeyOrPut(key, challenge)
 
@@ -357,7 +357,7 @@ suite "Discovery v5.1 Packet Encodings Test Vectors":
         challengeData: hexToSeqByte(whoareyouChallengeData))
       pubkey = none(PublicKey)
       challenge = Challenge(whoareyouData: whoareyouData, pubkey: pubkey)
-      key = HandShakeKey(nodeId: nodeA.id, address: $(nodeA.address.get()))
+      key = HandShakeKey(nodeId: nodeA.id, address: nodeA.address.get())
 
     check: not codecB.handshakes.hasKeyOrPut(key, challenge)
 
@@ -472,7 +472,7 @@ suite "Discovery v5.1 Additional Encode/Decode":
 
     let decoded = codecB.decodePacket(nodeA.address.get(), data)
 
-    let key = HandShakeKey(nodeId: nodeB.id, address: $nodeB.address.get())
+    let key = HandShakeKey(nodeId: nodeB.id, address: nodeB.address.get())
     var challenge: Challenge
 
     check: