From ff0b1a330ae69fb43e8ff117ea6ec3e55de3ce39 Mon Sep 17 00:00:00 2001
From: Justin Traglia <95511699+jtraglia@users.noreply.github.com>
Date: Tue, 15 Nov 2022 05:36:16 -0600
Subject: [PATCH] Allow distance of 256 in lookupDistances (#556)

---
 eth/p2p/discoveryv5/protocol.nim |  2 +-
 tests/p2p/test_discoveryv5.nim   | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index 6254092..0b4956e 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -623,7 +623,7 @@ proc lookupDistances*(target, dest: NodeId): seq[uint16] =
   result.add(td)
   var i = 1
   while result.len < lookupRequestLimit:
-    if tdAsInt + i < 256:
+    if tdAsInt + i <= 256:
       result.add(td + uint16(i))
     if tdAsInt - i > 0:
       result.add(td - uint16(i))
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
index 920b2fa..0ca0c34 100644
--- a/tests/p2p/test_discoveryv5.nim
+++ b/tests/p2p/test_discoveryv5.nim
@@ -582,6 +582,9 @@ suite "Discovery v5 Tests":
       check test.len == 1
 
   test "Calculate lookup distances":
+    let rng = newRng()
+    let node = generateNode(PrivateKey.random(rng[]))
+
     # Log distance between zeros is zero
     let dist = lookupDistances(u256(0), u256(0))
     check dist == @[0'u16, 1, 2]
@@ -590,6 +593,22 @@ suite "Discovery v5 Tests":
     let dist1 = lookupDistances(u256(0), u256(1))
     check dist1 == @[1'u16, 2, 3]
 
+    # Ensure that distances are in expected order
+    let dist2 = lookupDistances(node.id, idAtDistance(node.id, 128))
+    check dist2 == @[128'u16, 129, 127]
+
+    # Test target distance of 256
+    let dist3 = lookupDistances(node.id, idAtDistance(node.id, 256))
+    check dist3 == @[256'u16, 255, 254]
+
+    # Test target distance of 255
+    let dist4 = lookupDistances(node.id, idAtDistance(node.id, 255))
+    check dist4 == @[255'u16, 256, 254]
+
+    # Test target distance of 254
+    let dist5 = lookupDistances(node.id, idAtDistance(node.id, 254))
+    check dist5 == @[254'u16, 255, 253]
+
   asyncTest "Handshake cleanup: different ids":
     # Node to test the handshakes on.
     let receiveNode = initDiscoveryNode(