From 65ed0c4850e7944c179dfef1f5240dc608e7e0ab Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Fri, 9 Oct 2020 10:46:28 +0200
Subject: [PATCH] Activate discovery v5.0 handshake tests again

---
 tests/p2p/discv5_test_helper.nim |  11 ---
 tests/p2p/test_discoveryv5.nim   | 112 +++++++++++++++++--------------
 2 files changed, 62 insertions(+), 61 deletions(-)

diff --git a/tests/p2p/discv5_test_helper.nim b/tests/p2p/discv5_test_helper.nim
index 500ed57..90b360d 100644
--- a/tests/p2p/discv5_test_helper.nim
+++ b/tests/p2p/discv5_test_helper.nim
@@ -36,17 +36,6 @@ proc nodeIdInNodes*(id: NodeId, nodes: openarray[Node]): bool =
   for n in nodes:
     if id == n.id: return true
 
-# proc randomPacket*(rng: var BrHmacDrbgContext, tag: PacketTag): seq[byte] =
-#   var
-#     authTag: AuthTag
-#     msg: array[44, byte]
-
-#   brHmacDrbgGenerate(rng, authTag)
-#   brHmacDrbgGenerate(rng, msg)
-#   result.add(tag)
-#   result.add(rlp.encode(authTag))
-#   result.add(msg)
-
 proc generateNode*(privKey: PrivateKey, port: int = 20302,
     localEnrFields: openarray[FieldPair] = []): Node =
   let port = Port(port)
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
index a06341a..8df28d8 100644
--- a/tests/p2p/test_discoveryv5.nim
+++ b/tests/p2p/test_discoveryv5.nim
@@ -73,56 +73,6 @@ procSuite "Discovery v5 Tests":
 
     await node1.closeWait()
 
-  # asyncTest "Handshake cleanup":
-  #   let node = initDiscoveryNode(
-  #     rng, PrivateKey.random(rng[]), localAddress(20302))
-  #   var tag: PacketTag
-  #   let a = localAddress(20303)
-
-  #   for i in 0 ..< 5:
-  #     brHmacDrbgGenerate(rng[], tag)
-  #     node.receive(a, randomPacket(rng[], tag))
-
-  #   # Checking different nodeIds but same address
-  #   check node.codec.handshakes.len == 5
-  #   # TODO: Could get rid of the sleep by storing the timeout future of the
-  #   # handshake
-  #   await sleepAsync(handshakeTimeout)
-  #   # Checking handshake cleanup
-  #   check node.codec.handshakes.len == 0
-
-  #   await node.closeWait()
-
-  # asyncTest "Handshake different address":
-  #   let node = initDiscoveryNode(
-  #     rng, PrivateKey.random(rng[]), localAddress(20302))
-  #   var tag: PacketTag
-
-  #   for i in 0 ..< 5:
-  #     let a = localAddress(20303 + i)
-  #     node.receive(a, randomPacket(rng[], tag))
-
-  #   check node.codec.handshakes.len == 5
-
-  #   await node.closeWait()
-
-  # asyncTest "Handshake duplicates":
-  #   let node = initDiscoveryNode(
-  #     rng, PrivateKey.random(rng[]), localAddress(20302))
-  #   var tag: PacketTag
-  #   let a = localAddress(20303)
-
-  #   for i in 0 ..< 5:
-  #     node.receive(a, randomPacket(rng[], tag))
-
-  #   # Checking handshake duplicates
-  #   check node.codec.handshakes.len == 1
-
-  #   # TODO: add check that gets the Whoareyou value and checks if its authTag
-  #   # is that of the first packet.
-
-  #   await node.closeWait()
-
   test "Distance check":
     const
       targetId = "0x0000"
@@ -585,3 +535,65 @@ procSuite "Discovery v5 Tests":
         records = [recordInvalidDistance]
         test = verifyNodesRecords(records, fromNode, 0'u32)
       check test.len == 0
+
+  when not UseDiscv51:
+    proc randomPacket(rng: var BrHmacDrbgContext, tag: PacketTag): seq[byte] =
+      var
+        authTag: AuthTag
+        msg: array[44, byte]
+
+      brHmacDrbgGenerate(rng, authTag)
+      brHmacDrbgGenerate(rng, msg)
+      result.add(tag)
+      result.add(rlp.encode(authTag))
+      result.add(msg)
+
+    asyncTest "Handshake cleanup":
+      let node = initDiscoveryNode(
+        rng, PrivateKey.random(rng[]), localAddress(20302))
+      var tag: PacketTag
+      let a = localAddress(20303)
+
+      for i in 0 ..< 5:
+        brHmacDrbgGenerate(rng[], tag)
+        node.receive(a, randomPacket(rng[], tag))
+
+      # Checking different nodeIds but same address
+      check node.codec.handshakes.len == 5
+      # TODO: Could get rid of the sleep by storing the timeout future of the
+      # handshake
+      await sleepAsync(handshakeTimeout)
+      # Checking handshake cleanup
+      check node.codec.handshakes.len == 0
+
+      await node.closeWait()
+
+    asyncTest "Handshake different address":
+      let node = initDiscoveryNode(
+        rng, PrivateKey.random(rng[]), localAddress(20302))
+      var tag: PacketTag
+
+      for i in 0 ..< 5:
+        let a = localAddress(20303 + i)
+        node.receive(a, randomPacket(rng[], tag))
+
+      check node.codec.handshakes.len == 5
+
+      await node.closeWait()
+
+    asyncTest "Handshake duplicates":
+      let node = initDiscoveryNode(
+        rng, PrivateKey.random(rng[]), localAddress(20302))
+      var tag: PacketTag
+      let a = localAddress(20303)
+
+      for i in 0 ..< 5:
+        node.receive(a, randomPacket(rng[], tag))
+
+      # Checking handshake duplicates
+      check node.codec.handshakes.len == 1
+
+      # TODO: add check that gets the Whoareyou value and checks if its authTag
+      # is that of the first packet.
+
+      await node.closeWait()