From 621b423ac1b61d6bac1651cf1481fb8220080d40 Mon Sep 17 00:00:00 2001
From: Justin Traglia <95511699+jtraglia@users.noreply.github.com>
Date: Wed, 2 Nov 2022 08:50:07 -0500
Subject: [PATCH] p2p/discover: fix handling of distance 256 in lookupDistances
 (#26087)

Noticed that lookupDistances for FINDNODE requests didn't consider 256 a valid
distance. This is actually part of the example in the comment above the
function, surprised that wasn't tested before.
---
 p2p/discover/v5_udp.go          |  2 +-
 p2p/discover/v5_udp_test.go     | 37 +++++++++++++++++++++++++++++++++
 p2p/discover/v5wire/encoding.go |  4 ++--
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/p2p/discover/v5_udp.go b/p2p/discover/v5_udp.go
index 757a3587dc7d..1c66602f8791 100644
--- a/p2p/discover/v5_udp.go
+++ b/p2p/discover/v5_udp.go
@@ -323,7 +323,7 @@ func lookupDistances(target, dest enode.ID) (dists []uint) {
 	td := enode.LogDist(target, dest)
 	dists = append(dists, uint(td))
 	for i := 1; len(dists) < lookupRequestLimit; i++ {
-		if td+i < 256 {
+		if td+i <= 256 {
 			dists = append(dists, uint(td+i))
 		}
 		if td-i > 0 {
diff --git a/p2p/discover/v5_udp_test.go b/p2p/discover/v5_udp_test.go
index 30d610a4dd8c..ca63688afa13 100644
--- a/p2p/discover/v5_udp_test.go
+++ b/p2p/discover/v5_udp_test.go
@@ -34,6 +34,7 @@ import (
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/stretchr/testify/require"
 )
 
 // Real sockets, real crypto: this test checks end-to-end connectivity for UDPv5.
@@ -519,6 +520,42 @@ func TestUDPv5_talkRequest(t *testing.T) {
 	}
 }
 
+// This test checks that lookupDistances works.
+func TestUDPv5_lookupDistances(t *testing.T) {
+	test := newUDPV5Test(t)
+	lnID := test.table.self().ID()
+
+	t.Run("target distance of 1", func(t *testing.T) {
+		node := nodeAtDistance(lnID, 1, intIP(0))
+		dists := lookupDistances(lnID, node.ID())
+		require.Equal(t, []uint{1, 2, 3}, dists)
+	})
+
+	t.Run("target distance of 2", func(t *testing.T) {
+		node := nodeAtDistance(lnID, 2, intIP(0))
+		dists := lookupDistances(lnID, node.ID())
+		require.Equal(t, []uint{2, 3, 1}, dists)
+	})
+
+	t.Run("target distance of 128", func(t *testing.T) {
+		node := nodeAtDistance(lnID, 128, intIP(0))
+		dists := lookupDistances(lnID, node.ID())
+		require.Equal(t, []uint{128, 129, 127}, dists)
+	})
+
+	t.Run("target distance of 255", func(t *testing.T) {
+		node := nodeAtDistance(lnID, 255, intIP(0))
+		dists := lookupDistances(lnID, node.ID())
+		require.Equal(t, []uint{255, 256, 254}, dists)
+	})
+
+	t.Run("target distance of 256", func(t *testing.T) {
+		node := nodeAtDistance(lnID, 256, intIP(0))
+		dists := lookupDistances(lnID, node.ID())
+		require.Equal(t, []uint{256, 255, 254}, dists)
+	})
+}
+
 // This test checks that lookup works.
 func TestUDPv5_lookup(t *testing.T) {
 	t.Parallel()
diff --git a/p2p/discover/v5wire/encoding.go b/p2p/discover/v5wire/encoding.go
index 6f8f3466e9de..e41d7f4c451e 100644
--- a/p2p/discover/v5wire/encoding.go
+++ b/p2p/discover/v5wire/encoding.go
@@ -65,7 +65,7 @@ type (
 	handshakeAuthData struct {
 		h struct {
 			SrcID      enode.ID
-			SigSize    byte // ignature data
+			SigSize    byte // signature data
 			PubkeySize byte // offset of
 		}
 		// Trailing variable-size data.
@@ -529,7 +529,7 @@ func (c *Codec) decodeHandshake(fromAddr string, head *Header) (n *enode.Node, a
 	if err != nil {
 		return nil, auth, nil, errInvalidAuthKey
 	}
-	// Derive sesssion keys.
+	// Derive session keys.
 	session := deriveKeys(sha256.New, c.privkey, ephkey, auth.h.SrcID, c.localnode.ID(), cdata)
 	session = session.keysFlipped()
 	return n, auth, session, nil