From c3f23e5912efff98fc6c8181db579037e5a19a2c Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Tue, 24 Mar 2020 10:51:34 +0100
Subject: [PATCH] Minor adjustments to store bootnode records + deletion test

---
 eth/p2p/discoveryv5/enr.nim      |  2 ++
 eth/p2p/discoveryv5/node.nim     |  4 ----
 eth/p2p/discoveryv5/protocol.nim | 19 ++++++++++---------
 tests/p2p/test_discoveryv5.nim   | 24 ++++++++++++++++++++++++
 4 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/eth/p2p/discoveryv5/enr.nim b/eth/p2p/discoveryv5/enr.nim
index 1ff34e7..c876478 100644
--- a/eth/p2p/discoveryv5/enr.nim
+++ b/eth/p2p/discoveryv5/enr.nim
@@ -309,6 +309,8 @@ proc `$`*(r: Record): string =
     result &= $v
   result &= ')'
 
+proc `==`*(a, b: Record): bool = a.raw == b.raw
+
 proc read*(rlp: var Rlp, T: typedesc[Record]): T {.inline.} =
   if not result.fromBytes(rlp.rawData.toOpenArray):
     raise newException(ValueError, "Could not deserialize")
diff --git a/eth/p2p/discoveryv5/node.nim b/eth/p2p/discoveryv5/node.nim
index 9a0d8be..1dee8f8 100644
--- a/eth/p2p/discoveryv5/node.nim
+++ b/eth/p2p/discoveryv5/node.nim
@@ -48,10 +48,6 @@ proc newNode*(r: Record): Node =
   result = newNode(initENode(pk, a))
   result.record = r
 
-proc newNodes*(records: openarray[Record]): seq[Node] =
-  for record in records:
-    result.add(newNode(record))
-
 proc hash*(n: Node): hashes.Hash = hash(n.node.pubkey.data)
 proc `==`*(a, b: Node): bool = (a.isNil and b.isNil) or (not a.isNil and not b.isNil and a.node.pubkey == b.node.pubkey)
 
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index b1ab8ab..2d3fb7c 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -38,7 +38,7 @@ type
     awaitedPackets: Table[(NodeId, RequestId), Future[Option[Packet]]]
     lookupLoop: Future[void]
     revalidateLoop: Future[void]
-    bootstrapNodes: seq[Node]
+    bootstrapRecords*: seq[Record]
 
   PendingRequest = object
     node: Node
@@ -304,7 +304,7 @@ proc sendPing(d: Protocol, toNode: Node): RequestId =
   d.send(toNode, data)
   return reqId
 
-proc ping(d: Protocol, toNode: Node): Future[Option[PongPacket]] {.async.} =
+proc ping*(d: Protocol, toNode: Node): Future[Option[PongPacket]] {.async.} =
   let reqId = d.sendPing(toNode)
   let resp = await d.waitPacket(toNode, reqId)
 
@@ -392,7 +392,7 @@ proc lookupRandom*(d: Protocol): Future[seq[Node]]
     raise newException(RandomSourceDepleted, "Could not randomize bytes")
   d.lookup(id)
 
-proc revalidateNode(d: Protocol, n: Node)
+proc revalidateNode*(d: Protocol, n: Node)
     {.async, raises:[Defect, Exception].} = # TODO: Exception
   trace "Ping to revalidate node", node = $n
   let pong = await d.ping(n)
@@ -408,8 +408,8 @@ proc revalidateNode(d: Protocol, n: Node)
     # For now we never remove bootstrap nodes. It might make sense to actually
     # do so and to retry them only in case we drop to a really low amount of
     # peers in the DHT
-    if n notin d.bootstrapNodes:
-      trace "Revalidation of node failed, removing node", node = $n
+    if n.record notin d.bootstrapRecords:
+      trace "Revalidation of node failed, removing node", record = n.record
       d.routingTable.removeNode(n)
       # Remove shared secrets when removing the node from routing table.
       # This might be to direct, so we could keep these longer. But better
@@ -462,19 +462,20 @@ proc newProtocol*(privKey: PrivateKey, db: Database,
     whoareyouMagic: whoareyouMagic(node.id),
     idHash: sha256.digest(node.id.toByteArrayBE).data,
     codec: Codec(localNode: node, privKey: privKey, db: db),
-    bootstrapNodes: newNodes(bootstrapRecords))
+    bootstrapRecords: @bootstrapRecords)
 
   result.routingTable.init(node)
 
 proc open*(d: Protocol) =
-  debug "Starting discovery node", node = $d.localNode,
+  info "Starting discovery node", node = $d.localNode,
     uri = toURI(d.localNode.record)
   # TODO allow binding to specific IP / IPv6 / etc
   let ta = initTAddress(IPv4_any(), d.localNode.node.address.udpPort)
   d.transp = newDatagramTransport(processClient, udata = d, local = ta)
 
-  for node in d.bootstrapNodes:
-    d.addNode(node)
+  for record in d.bootstrapRecords:
+    debug "Adding bootstrap node", uri = toURI(record)
+    d.addNode(record)
 
 proc start*(d: Protocol) =
   # Might want to move these to a separate proc if this turns out to be needed.
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
index 6ea5e07..bc7558a 100644
--- a/tests/p2p/test_discoveryv5.nim
+++ b/tests/p2p/test_discoveryv5.nim
@@ -133,6 +133,30 @@ suite "Discovery v5 Tests":
 
     await node.closeWait()
 
+  asyncTest "Node deletion":
+    let
+      bootnode = initDiscoveryNode(newPrivateKey(), localAddress(20301), @[])
+      node1 = initDiscoveryNode(newPrivateKey(), localAddress(20302),
+        @[bootnode.localNode.record])
+      node2 = initDiscoveryNode(newPrivateKey(), localAddress(20303),
+        @[bootnode.localNode.record])
+      pong1 = await discv5_protocol.ping(node1, bootnode.localNode)
+      pong2 = await discv5_protocol.ping(node1, node2.localNode)
+
+    check pong1.isSome() and pong2.isSome()
+
+    await bootnode.closeWait()
+    await node2.closeWait()
+
+    await node1.revalidateNode(bootnode.localNode)
+    await node1.revalidateNode(node2.localNode)
+
+    check node1.getNode(bootnode.localNode.id) == bootnode.localNode
+    check node1.getNode(node2.localNode.id) == nil
+
+    await node1.closeWait()
+
+
   asyncTest "Handshake cleanup":
     let node = initDiscoveryNode(newPrivateKey(), localAddress(20302), @[])
     var tag: PacketTag