Merge pull request #16 from status-im/discotest

Fixed findNode msg handling, minor optimization, added a test

ethp2p/discovery.nim

@@ -46,7 +46,7 @@ type
     privKey: PrivateKey
     address: Address
     bootstrapNodes*: seq[Node]
-    thisNode: Node
+    thisNode*: Node
     kademlia: KademliaProtocol[DiscoveryProtocol]
     socket: AsyncSocket
 
@@ -134,7 +134,8 @@ proc sendPong*(d: DiscoveryProtocol, n: Node, token: MDigest[256]) =
   d.send(n, msg)
 
 proc sendFindNode*(d: DiscoveryProtocol, n: Node, targetNodeId: NodeId) =
-  var data = newSeq[byte](32) & @(targetNodeId.toByteArrayBE())
+  var data: array[64, byte]
+  data[32 .. ^1] = targetNodeId.toByteArrayBE()
   let payload = rlp.encode((data, expiration()))
   let msg = pack(cmdFindNode, payload, d.privKey)
   debug ">>> find_node to ", n#, ": ", msg.toHex()
@@ -214,7 +215,7 @@ proc recvFindNode(d: DiscoveryProtocol, node: Node, payload: Bytes) {.inline.} =
   let rlp = rlpFromBytes(payload.toRange)
   debug "<<< find_node from ", node
   let rng = rlp.listElem(0).toBytes
-  let nodeId = readUIntBE[256](rng.toOpenArray())
+  let nodeId = readUIntBE[256](rng[32 .. ^1].toOpenArray())
   d.kademlia.recvFindNode(node, nodeId)
 
 proc expirationValid(rlpEncodedPayload: seq[byte]): bool {.inline.} =

tests/tdiscovery.nim

@@ -1,4 +1,46 @@
-import ../ethp2p/discovery
-import ../ethp2p/peer_pool
+import ../ethp2p/[discovery, kademlia, peer_pool, enode]
+import eth_keys, net, asyncdispatch, sequtils
 
-# TODO:
+import logging, byteutils
+
+addHandler(newConsoleLogger())
+
+proc startDiscoveryNode(privKey: PrivateKey, address: Address, bootnodes: seq[ENode]): Future[DiscoveryProtocol] {.async.} =
+  result = newDiscoveryProtocol(privKey, address, bootnodes)
+  result.open()
+  await result.bootstrap()
+
+proc localAddress(port: int): Address =
+  let port = Port(port)
+  result = Address(udpPort: port, tcpPort: port, ip: parseIpAddress("127.0.0.1"))
+
+let
+  bootNodeKey = initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a617")
+  bootNodeAddr = localAddress(20301)
+  bootENode = initENode(bootNodeKey.getPublicKey, bootNodeAddr)
+
+  nodeKeys = [
+    initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a618"),
+    initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a619"),
+    initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a620")
+  ]
+
+proc nodeIdInNodes(id: NodeId, nodes: openarray[Node]): bool =
+  for n in nodes:
+    if id == n.id: return true
+
+proc test() {.async.} =
+  let bootNode = await startDiscoveryNode(bootNodeKey, bootNodeAddr, @[])
+
+  var nodeAddrs = newSeqOfCap[Address](nodeKeys.len)
+  for i in 0 ..< nodeKeys.len: nodeAddrs.add(localAddress(20302 + i))
+
+  var nodes = await all(zip(nodeKeys, nodeAddrs).mapIt(startDiscoveryNode(it.a, it.b, @[bootENode])))
+  nodes.add(bootNode)
+
+  for i in nodes:
+    for j in nodes:
+      if j != i:
+        doAssert(nodeIdInNodes(i.thisNode.id, j.randomNodes(nodes.len - 1)))
+
+waitFor test()