diff --git a/p2p/discover/table.go b/p2p/discover/table.go index e6dafb0dca..2b7a28708b 100644 --- a/p2p/discover/table.go +++ b/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() diff --git a/p2p/discover/table_test.go b/p2p/discover/table_test.go index 2781dd4225..3ba3422251 100644 --- a/p2p/discover/table_test.go +++ b/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 diff --git a/p2p/discover/table_util_test.go b/p2p/discover/table_util_test.go index 8f3813bdcf..d6309dfd6c 100644 --- a/p2p/discover/table_util_test.go +++ b/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) } } diff --git a/p2p/discover/v4_lookup_test.go b/p2p/discover/v4_lookup_test.go index 1f9ad69d0a..8867a5a8ac 100644 --- a/p2p/discover/v4_lookup_test.go +++ b/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() diff --git a/p2p/discover/v4_udp_test.go b/p2p/discover/v4_udp_test.go index 53ecb1bc6e..361e379626 100644 --- a/p2p/discover/v4_udp_test.go +++ b/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. diff --git a/p2p/discover/v5_udp.go b/p2p/discover/v5_udp.go index 6ba7a90618..8b3e33d37c 100644 --- a/p2p/discover/v5_udp.go +++ b/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 } diff --git a/p2p/discover/v5_udp_test.go b/p2p/discover/v5_udp_test.go index 18d8aeac6d..eaa969ea8b 100644 --- a/p2p/discover/v5_udp_test.go +++ b/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)