p2p/discover: add liveness check in collectTableNodes (#28686)

* p2p/discover: add liveness check in collectTableNodes

* p2p/discover: fix test

* p2p/discover: rename to appendLiveNodes

* p2p/discover: add dedup logic back

* p2p/discover: simplify

* p2p/discover: fix issue found by test

Author: fjl

p2p/discover/table.go

@@ -459,6 +459,26 @@ func (tab *Table) findnodeByID(target enode.ID, nresults int, preferLive bool) *
 	return nodes
 }
 
+// appendLiveNodes adds nodes at the given distance to the result slice.
+func (tab *Table) appendLiveNodes(dist uint, result []*enode.Node) []*enode.Node {
+	if dist > 256 {
+		return result
+	}
+	if dist == 0 {
+		return append(result, tab.self())
+	}
+
+	tab.mutex.Lock()
+	defer tab.mutex.Unlock()
+	for _, n := range tab.bucketAtDistance(int(dist)).entries {
+		if n.livenessChecks >= 1 {
+			node := n.Node // avoid handing out pointer to struct field
+			result = append(result, &node)
+		}
+	}
+	return result
+}
+
 // len returns the number of nodes in the table.
 func (tab *Table) len() (n int) {
 	tab.mutex.Lock()

p2p/discover/table_test.go

@@ -199,7 +199,7 @@ func TestTable_findnodeByID(t *testing.T) {
 		tab, db := newTestTable(transport)
 		defer db.Close()
 		defer tab.close()
-		fillTable(tab, test.All)
+		fillTable(tab, test.All, true)
 
 		// check that closest(Target, N) returns nodes
 		result := tab.findnodeByID(test.Target, test.N, false).entries

p2p/discover/table_util_test.go

@@ -109,8 +109,11 @@ func fillBucket(tab *Table, n *node) (last *node) {
 
 // fillTable adds nodes the table to the end of their corresponding bucket
 // if the bucket is not full. The caller must not hold tab.mutex.
-func fillTable(tab *Table, nodes []*node) {
+func fillTable(tab *Table, nodes []*node, setLive bool) {
 	for _, n := range nodes {
+		if setLive {
+			n.livenessChecks = 1
+		}
 		tab.addSeenNode(n)
 	}
 }

p2p/discover/v4_lookup_test.go

@@ -40,7 +40,7 @@ func TestUDPv4_Lookup(t *testing.T) {
 	}
 
 	// Seed table with initial node.
-	fillTable(test.table, []*node{wrapNode(lookupTestnet.node(256, 0))})
+	fillTable(test.table, []*node{wrapNode(lookupTestnet.node(256, 0))}, true)
 
 	// Start the lookup.
 	resultC := make(chan []*enode.Node, 1)
@@ -74,7 +74,7 @@ func TestUDPv4_LookupIterator(t *testing.T) {
 	for i := range lookupTestnet.dists[256] {
 		bootnodes[i] = wrapNode(lookupTestnet.node(256, i))
 	}
-	fillTable(test.table, bootnodes)
+	fillTable(test.table, bootnodes, true)
 	go serveTestnet(test, lookupTestnet)
 
 	// Create the iterator and collect the nodes it yields.
@@ -109,7 +109,7 @@ func TestUDPv4_LookupIteratorClose(t *testing.T) {
 	for i := range lookupTestnet.dists[256] {
 		bootnodes[i] = wrapNode(lookupTestnet.node(256, i))
 	}
-	fillTable(test.table, bootnodes)
+	fillTable(test.table, bootnodes, true)
 	go serveTestnet(test, lookupTestnet)
 
 	it := test.udp.RandomNodes()

p2p/discover/v4_udp_test.go

@@ -269,7 +269,7 @@ func TestUDPv4_findnode(t *testing.T) {
 		}
 		nodes.push(n, numCandidates)
 	}
-	fillTable(test.table, nodes.entries)
+	fillTable(test.table, nodes.entries, false)
 
 	// ensure there's a bond with the test node,
 	// findnode won't be accepted otherwise.

p2p/discover/v5_udp.go

@@ -851,6 +851,7 @@ func (t *UDPv5) handleFindnode(p *v5wire.Findnode, fromID enode.ID, fromAddr *ne
 
 // collectTableNodes creates a FINDNODE result set for the given distances.
 func (t *UDPv5) collectTableNodes(rip net.IP, distances []uint, limit int) []*enode.Node {
+	var bn []*enode.Node
 	var nodes []*enode.Node
 	var processed = make(map[uint]struct{})
 	for _, dist := range distances {
@@ -859,21 +860,11 @@ func (t *UDPv5) collectTableNodes(rip net.IP, distances []uint, limit int) []*en
 		if seen || dist > 256 {
 			continue
 		}
-
-		// Get the nodes.
-		var bn []*enode.Node
-		if dist == 0 {
-			bn = []*enode.Node{t.Self()}
-		} else if dist <= 256 {
-			t.tab.mutex.Lock()
-			bn = unwrapNodes(t.tab.bucketAtDistance(int(dist)).entries)
-			t.tab.mutex.Unlock()
-		}
 		processed[dist] = struct{}{}
 
-		// Apply some pre-checks to avoid sending invalid nodes.
-		for _, n := range bn {
-			// TODO livenessChecks > 1
+		for _, n := range t.tab.appendLiveNodes(dist, bn[:0]) {
+			// Apply some pre-checks to avoid sending invalid nodes.
+			// Note liveness is checked by appendLiveNodes.
 			if netutil.CheckRelayIP(rip, n.IP()) != nil {
 				continue
 			}

p2p/discover/v5_udp_test.go

@@ -159,9 +159,9 @@ func TestUDPv5_findnodeHandling(t *testing.T) {
 	nodes253 := nodesAtDistance(test.table.self().ID(), 253, 16)
 	nodes249 := nodesAtDistance(test.table.self().ID(), 249, 4)
 	nodes248 := nodesAtDistance(test.table.self().ID(), 248, 10)
-	fillTable(test.table, wrapNodes(nodes253))
-	fillTable(test.table, wrapNodes(nodes249))
-	fillTable(test.table, wrapNodes(nodes248))
+	fillTable(test.table, wrapNodes(nodes253), true)
+	fillTable(test.table, wrapNodes(nodes249), true)
+	fillTable(test.table, wrapNodes(nodes248), true)
 
 	// Requesting with distance zero should return the node's own record.
 	test.packetIn(&v5wire.Findnode{ReqID: []byte{0}, Distances: []uint{0}})
@@ -589,7 +589,7 @@ func TestUDPv5_lookup(t *testing.T) {
 
 	// Seed table with initial node.
 	initialNode := lookupTestnet.node(256, 0)
-	fillTable(test.table, []*node{wrapNode(initialNode)})
+	fillTable(test.table, []*node{wrapNode(initialNode)}, true)
 
 	// Start the lookup.
 	resultC := make(chan []*enode.Node, 1)