p2p/discover: fix out-of-bounds slicing for chunked neighbors packets

The code assumed that Table.closest always returns at least 13 nodes.
This is not true for small tables (e.g. during bootstrap).

1 files changed, 13 insertions, 18 deletions

diff --git a/p2p/discover/udp_test.go b/p2p/discover/udp_test.go
index 7bf6df5ab..11fa31d7c 100644
--- a/p2p/discover/udp_test.go
+++ b/p2p/discover/udp_test.go
@@ -164,26 +164,21 @@ func TestUDP_findnode(t *testing.T) {
 	// check that closest neighbors are returned.
 	test.packetIn(nil, findnodePacket, &findnode{Target: testTarget, Expiration: futureExp})
 	expected := test.table.closest(targetHash, bucketSize)
-	test.waitPacketOut(func(p *neighbors) {
-		if len(p.Nodes) != 13 {
-			t.Errorf("wrong number of results: got %d, want %d", len(p.Nodes), bucketSize)
-		}
-		for i := range p.Nodes {
-			if p.Nodes[i].ID != expected.entries[i].ID {
-				t.Errorf("result mismatch at %d:\n  got:  %v\n  want: %v", i, p.Nodes[i], expected.entries[i])
+
+	waitNeighbors := func(want []*Node) {
+		test.waitPacketOut(func(p *neighbors) {
+			if len(p.Nodes) != len(want) {
+				t.Errorf("wrong number of results: got %d, want %d", len(p.Nodes), bucketSize)
 			}
-		}
-	})
-	test.waitPacketOut(func(p *neighbors) {
-		if len(p.Nodes) != 3 {
-			t.Errorf("wrong number of results: got %d, want %d", len(p.Nodes), bucketSize)
-		}
-		for i := range p.Nodes {
-			if p.Nodes[i].ID != expected.entries[i + 13].ID {
-				t.Errorf("result mismatch at %d:\n  got:  %v\n  want: %v", i, p.Nodes[i], expected.entries[i])
+			for i := range p.Nodes {
+				if p.Nodes[i].ID != want[i].ID {
+					t.Errorf("result mismatch at %d:\n  got:  %v\n  want: %v", i, p.Nodes[i], expected.entries[i])
+				}
 			}
-		}
-	})
+		})
+	}
+	waitNeighbors(expected.entries[:maxNeighbors])
+	waitNeighbors(expected.entries[maxNeighbors:])
 }
 
 func TestUDP_findnodeMultiReply(t *testing.T) {