p2p/discover: add RandomNodes iterator

p2p/discover/common.go

@@ -25,6 +25,7 @@ import (
 	"github.com/ethereum/go-ethereum/p2p/netutil"
 )
 
+// UDPConn is a network connection on which discovery can operate.
 type UDPConn interface {
 	ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error)
 	WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error)
@@ -32,7 +33,7 @@ type UDPConn interface {
 	LocalAddr() net.Addr
 }
 
-// Config holds Table-related settings.
+// Config holds settings for the discovery listener.
 type Config struct {
 	// These settings are required and configure the UDP listener:
 	PrivateKey *ecdsa.PrivateKey
@@ -50,7 +51,7 @@ func ListenUDP(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) {
 }
 
 // ReadPacket is a packet that couldn't be handled. Those packets are sent to the unhandled
-// channel if configured.
+// channel if configured. This is exported for internal use, do not use this type.
 type ReadPacket struct {
 	Data []byte
 	Addr *net.UDPAddr

p2p/discover/lookup.go

@@ -0,0 +1,209 @@
+// Copyright 2019 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package discover
+
+import (
+	"context"
+
+	"github.com/ethereum/go-ethereum/p2p/enode"
+)
+
+// lookup performs a network search for nodes close to the given target. It approaches the
+// target by querying nodes that are closer to it on each iteration. The given target does
+// not need to be an actual node identifier.
+type lookup struct {
+	tab         *Table
+	queryfunc   func(*node) ([]*node, error)
+	replyCh     chan []*node
+	cancelCh    <-chan struct{}
+	asked, seen map[enode.ID]bool
+	result      nodesByDistance
+	replyBuffer []*node
+	queries     int
+}
+
+type queryFunc func(*node) ([]*node, error)
+
+func newLookup(ctx context.Context, tab *Table, target enode.ID, q queryFunc) *lookup {
+	it := &lookup{
+		tab:       tab,
+		queryfunc: q,
+		asked:     make(map[enode.ID]bool),
+		seen:      make(map[enode.ID]bool),
+		result:    nodesByDistance{target: target},
+		replyCh:   make(chan []*node, alpha),
+		cancelCh:  ctx.Done(),
+		queries:   -1,
+	}
+	// Don't query further if we hit ourself.
+	// Unlikely to happen often in practice.
+	it.asked[tab.self().ID()] = true
+	return it
+}
+
+// run runs the lookup to completion and returns the closest nodes found.
+func (it *lookup) run() []*enode.Node {
+	for it.next() {
+	}
+	return unwrapNodes(it.result.entries)
+}
+
+// next advances the lookup until any new nodes have been found.
+// It returns false when the lookup has ended.
+func (it *lookup) next() bool {
+	for it.startQueries() {
+		select {
+		case nodes := <-it.replyCh:
+			it.replyBuffer = it.replyBuffer[:0]
+			for _, n := range nodes {
+				if n != nil && !it.seen[n.ID()] {
+					it.seen[n.ID()] = true
+					it.result.push(n, bucketSize)
+					it.replyBuffer = append(it.replyBuffer, n)
+				}
+			}
+			it.queries--
+			if len(it.replyBuffer) > 0 {
+				return true
+			}
+		case <-it.cancelCh:
+			it.shutdown()
+		}
+	}
+	return false
+}
+
+func (it *lookup) shutdown() {
+	for it.queries > 0 {
+		<-it.replyCh
+		it.queries--
+	}
+	it.queryfunc = nil
+	it.replyBuffer = nil
+}
+
+func (it *lookup) startQueries() bool {
+	if it.queryfunc == nil {
+		return false
+	}
+
+	// The first query returns nodes from the local table.
+	if it.queries == -1 {
+		it.tab.mutex.Lock()
+		closest := it.tab.closest(it.result.target, bucketSize, false)
+		it.tab.mutex.Unlock()
+		it.queries = 1
+		it.replyCh <- closest.entries
+		return true
+	}
+
+	// Ask the closest nodes that we haven't asked yet.
+	for i := 0; i < len(it.result.entries) && it.queries < alpha; i++ {
+		n := it.result.entries[i]
+		if !it.asked[n.ID()] {
+			it.asked[n.ID()] = true
+			it.queries++
+			go it.query(n, it.replyCh)
+		}
+	}
+	// The lookup ends when no more nodes can be asked.
+	return it.queries > 0
+}
+
+func (it *lookup) query(n *node, reply chan<- []*node) {
+	fails := it.tab.db.FindFails(n.ID(), n.IP())
+	r, err := it.queryfunc(n)
+	if err == errClosed {
+		// Avoid recording failures on shutdown.
+		reply <- nil
+		return
+	} else if len(r) == 0 {
+		fails++
+		it.tab.db.UpdateFindFails(n.ID(), n.IP(), fails)
+		it.tab.log.Trace("Findnode failed", "id", n.ID(), "failcount", fails, "err", err)
+		if fails >= maxFindnodeFailures {
+			it.tab.log.Trace("Too many findnode failures, dropping", "id", n.ID(), "failcount", fails)
+			it.tab.delete(n)
+		}
+	} else if fails > 0 {
+		// Reset failure counter because it counts _consecutive_ failures.
+		it.tab.db.UpdateFindFails(n.ID(), n.IP(), 0)
+	}
+
+	// Grab as many nodes as possible. Some of them might not be alive anymore, but we'll
+	// just remove those again during revalidation.
+	for _, n := range r {
+		it.tab.addSeenNode(n)
+	}
+	reply <- r
+}
+
+// lookupIterator performs lookup operations and iterates over all seen nodes.
+// When a lookup finishes, a new one is created through nextLookup.
+type lookupIterator struct {
+	buffer     []*node
+	nextLookup lookupFunc
+	ctx        context.Context
+	cancel     func()
+	lookup     *lookup
+}
+
+type lookupFunc func(ctx context.Context) *lookup
+
+func newLookupIterator(ctx context.Context, next lookupFunc) *lookupIterator {
+	ctx, cancel := context.WithCancel(ctx)
+	return &lookupIterator{ctx: ctx, cancel: cancel, nextLookup: next}
+}
+
+// Node returns the current node.
+func (it *lookupIterator) Node() *enode.Node {
+	if len(it.buffer) == 0 {
+		return nil
+	}
+	return unwrapNode(it.buffer[0])
+}
+
+// Next moves to the next node.
+func (it *lookupIterator) Next() bool {
+	// Consume next node in lastReply.
+	if len(it.buffer) > 0 {
+		it.buffer = it.buffer[1:]
+	}
+	// Advance the lookup to refill lastReply.
+	for len(it.buffer) == 0 {
+		if it.ctx.Err() != nil {
+			it.lookup = nil
+			it.buffer = nil
+			return false
+		}
+		if it.lookup == nil {
+			it.lookup = it.nextLookup(it.ctx)
+			continue
+		}
+		if !it.lookup.next() {
+			it.lookup = nil
+			continue
+		}
+		it.buffer = it.lookup.replyBuffer
+	}
+	return true
+}
+
+// Close ends the iterator.
+func (it *lookupIterator) Close() {
+	it.cancel()
+}

p2p/discover/table_util_test.go

@@ -17,11 +17,14 @@
 package discover
 
 import (
+	"bytes"
 	"crypto/ecdsa"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"math/rand"
 	"net"
+	"reflect"
 	"sort"
 	"sync"
 
@@ -169,6 +172,28 @@ func hasDuplicates(slice []*node) bool {
 	return false
 }
 
+func checkNodesEqual(got, want []*enode.Node) error {
+	if reflect.DeepEqual(got, want) {
+		return nil
+	}
+	output := new(bytes.Buffer)
+	fmt.Fprintf(output, "got %d nodes:\n", len(got))
+	for _, n := range got {
+		fmt.Fprintf(output, "  %v %v\n", n.ID(), n)
+	}
+	fmt.Fprintf(output, "want %d:\n", len(want))
+	for _, n := range want {
+		fmt.Fprintf(output, "  %v %v\n", n.ID(), n)
+	}
+	return errors.New(output.String())
+}
+
+func sortByID(nodes []*enode.Node) {
+	sort.Slice(nodes, func(i, j int) bool {
+		return string(nodes[i].ID().Bytes()) < string(nodes[j].ID().Bytes())
+	})
+}
+
 func sortedByDistanceTo(distbase enode.ID, slice []*node) bool {
 	return sort.SliceIsSorted(slice, func(i, j int) bool {
 		return enode.DistCmp(distbase, slice[i].ID(), slice[j].ID()) < 0

p2p/discover/v4_lookup_test.go

@@ -20,7 +20,6 @@ import (
 	"crypto/ecdsa"
 	"fmt"
 	"net"
-	"reflect"
 	"sort"
 	"testing"
 
@@ -49,19 +48,7 @@ func TestUDPv4_Lookup(t *testing.T) {
 	}()
 
 	// Answer lookup packets.
-	for done := false; !done; {
+	serveTestnet(test, lookupTestnet)
-		done = test.waitPacketOut(func(p packetV4, to *net.UDPAddr, hash []byte) {
-			n, key := lookupTestnet.nodeByAddr(to)
-			switch p.(type) {
-			case *pingV4:
-				test.packetInFrom(nil, key, to, &pongV4{Expiration: futureExp, ReplyTok: hash})
-			case *findnodeV4:
-				dist := enode.LogDist(n.ID(), lookupTestnet.target.id())
-				nodes := lookupTestnet.nodesAtDistance(dist - 1)
-				test.packetInFrom(nil, key, to, &neighborsV4{Expiration: futureExp, Nodes: nodes})
-			}
-		})
-	}
 
 	// Verify result nodes.
 	results := <-resultC
@@ -78,8 +65,94 @@ func TestUDPv4_Lookup(t *testing.T) {
 	if !sortedByDistanceTo(lookupTestnet.target.id(), wrapNodes(results)) {
 		t.Errorf("result set not sorted by distance to target")
 	}
-	if !reflect.DeepEqual(results, lookupTestnet.closest(bucketSize)) {
+	if err := checkNodesEqual(results, lookupTestnet.closest(bucketSize)); err != nil {
-		t.Errorf("results aren't the closest %d nodes", bucketSize)
+		t.Errorf("results aren't the closest %d nodes\n%v", bucketSize, err)
+	}
+}
+
+func TestUDPv4_LookupIterator(t *testing.T) {
+	t.Parallel()
+	test := newUDPTest(t)
+	defer test.close()
+
+	// Seed table with initial nodes.
+	bootnodes := make([]*node, len(lookupTestnet.dists[256]))
+	for i := range lookupTestnet.dists[256] {
+		bootnodes[i] = wrapNode(lookupTestnet.node(256, i))
+	}
+	fillTable(test.table, bootnodes)
+	go serveTestnet(test, lookupTestnet)
+
+	// Create the iterator and collect the nodes it yields.
+	iter := test.udp.RandomNodes()
+	seen := make(map[enode.ID]*enode.Node)
+	for limit := lookupTestnet.len(); iter.Next() && len(seen) < limit; {
+		seen[iter.Node().ID()] = iter.Node()
+	}
+	iter.Close()
+
+	// Check that all nodes in lookupTestnet were seen by the iterator.
+	results := make([]*enode.Node, 0, len(seen))
+	for _, n := range seen {
+		results = append(results, n)
+	}
+	sortByID(results)
+	want := lookupTestnet.nodes()
+	if err := checkNodesEqual(results, want); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// TestUDPv4_LookupIteratorClose checks that lookupIterator ends when its Close
+// method is called.
+func TestUDPv4_LookupIteratorClose(t *testing.T) {
+	t.Parallel()
+	test := newUDPTest(t)
+	defer test.close()
+
+	// Seed table with initial nodes.
+	bootnodes := make([]*node, len(lookupTestnet.dists[256]))
+	for i := range lookupTestnet.dists[256] {
+		bootnodes[i] = wrapNode(lookupTestnet.node(256, i))
+	}
+	fillTable(test.table, bootnodes)
+	go serveTestnet(test, lookupTestnet)
+
+	it := test.udp.RandomNodes()
+	if ok := it.Next(); !ok || it.Node() == nil {
+		t.Fatalf("iterator didn't return any node")
+	}
+
+	it.Close()
+
+	ncalls := 0
+	for ; ncalls < 100 && it.Next(); ncalls++ {
+		if it.Node() == nil {
+			t.Error("iterator returned Node() == nil node after Next() == true")
+		}
+	}
+	t.Logf("iterator returned %d nodes after close", ncalls)
+	if it.Next() {
+		t.Errorf("Next() == true after close and %d more calls", ncalls)
+	}
+	if n := it.Node(); n != nil {
+		t.Errorf("iterator returned non-nil node after close and %d more calls", ncalls)
+	}
+}
+
+func serveTestnet(test *udpTest, testnet *preminedTestnet) {
+	for done := false; !done; {
+		done = test.waitPacketOut(func(p packetV4, to *net.UDPAddr, hash []byte) {
+			n, key := testnet.nodeByAddr(to)
+			switch p.(type) {
+			case *pingV4:
+				test.packetInFrom(nil, key, to, &pongV4{Expiration: futureExp, ReplyTok: hash})
+			case *findnodeV4:
+				dist := enode.LogDist(n.ID(), testnet.target.id())
+				nodes := testnet.nodesAtDistance(dist - 1)
+				test.packetInFrom(nil, key, to, &neighborsV4{Expiration: futureExp, Nodes: nodes})
+			}
+		})
 	}
 }
 
@@ -148,6 +221,25 @@ type preminedTestnet struct {
 	dists  [hashBits + 1][]*ecdsa.PrivateKey
 }
 
+func (tn *preminedTestnet) len() int {
+	n := 0
+	for _, keys := range tn.dists {
+		n += len(keys)
+	}
+	return n
+}
+
+func (tn *preminedTestnet) nodes() []*enode.Node {
+	result := make([]*enode.Node, 0, tn.len())
+	for dist, keys := range tn.dists {
+		for index := range keys {
+			result = append(result, tn.node(dist, index))
+		}
+	}
+	sortByID(result)
+	return result
+}
+
 func (tn *preminedTestnet) node(dist, index int) *enode.Node {
 	key := tn.dists[dist][index]
 	ip := net.IP{127, byte(dist >> 8), byte(dist), byte(index)}

p2p/discover/v4_udp.go

@@ -19,6 +19,7 @@ package discover
 import (
 	"bytes"
 	"container/list"
+	"context"
 	"crypto/ecdsa"
 	crand "crypto/rand"
 	"errors"
@@ -207,7 +208,8 @@ type UDPv4 struct {
 
 	addReplyMatcher chan *replyMatcher
 	gotreply        chan reply
-	closing         chan struct{}
+	closeCtx        context.Context
+	cancelCloseCtx  func()
 }
 
 // replyMatcher represents a pending reply.
@@ -256,20 +258,23 @@ type reply struct {
 }
 
 func ListenV4(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) {
+	closeCtx, cancel := context.WithCancel(context.Background())
 	t := &UDPv4{
 		conn:            c,
 		priv:            cfg.PrivateKey,
 		netrestrict:     cfg.NetRestrict,
 		localNode:       ln,
 		db:              ln.Database(),
-		closing:         make(chan struct{}),
 		gotreply:        make(chan reply),
 		addReplyMatcher: make(chan *replyMatcher),
+		closeCtx:        closeCtx,
+		cancelCloseCtx:  cancel,
 		log:             cfg.Log,
 	}
 	if t.log == nil {
 		t.log = log.Root()
 	}
+
 	tab, err := newTable(t, ln.Database(), cfg.Bootnodes, t.log)
 	if err != nil {
 		return nil, err
@@ -291,126 +296,13 @@ func (t *UDPv4) Self() *enode.Node {
 // Close shuts down the socket and aborts any running queries.
 func (t *UDPv4) Close() {
 	t.closeOnce.Do(func() {
-		close(t.closing)
+		t.cancelCloseCtx()
 		t.conn.Close()
 		t.wg.Wait()
 		t.tab.close()
 	})
 }
 
-// ReadRandomNodes reads random nodes from the local table.
-func (t *UDPv4) ReadRandomNodes(buf []*enode.Node) int {
-	return t.tab.ReadRandomNodes(buf)
-}
-
-// LookupRandom finds random nodes in the network.
-func (t *UDPv4) LookupRandom() []*enode.Node {
-	if t.tab.len() == 0 {
-		// All nodes were dropped, refresh. The very first query will hit this
-		// case and run the bootstrapping logic.
-		<-t.tab.refresh()
-	}
-	return t.lookupRandom()
-}
-
-func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node {
-	if t.tab.len() == 0 {
-		// All nodes were dropped, refresh. The very first query will hit this
-		// case and run the bootstrapping logic.
-		<-t.tab.refresh()
-	}
-	return unwrapNodes(t.lookup(encodePubkey(key)))
-}
-
-func (t *UDPv4) lookupRandom() []*enode.Node {
-	var target encPubkey
-	crand.Read(target[:])
-	return unwrapNodes(t.lookup(target))
-}
-
-func (t *UDPv4) lookupSelf() []*enode.Node {
-	return unwrapNodes(t.lookup(encodePubkey(&t.priv.PublicKey)))
-}
-
-// lookup performs a network search for nodes close to the given target. It approaches the
-// target by querying nodes that are closer to it on each iteration. The given target does
-// not need to be an actual node identifier.
-func (t *UDPv4) lookup(targetKey encPubkey) []*node {
-	var (
-		target         = enode.ID(crypto.Keccak256Hash(targetKey[:]))
-		asked          = make(map[enode.ID]bool)
-		seen           = make(map[enode.ID]bool)
-		reply          = make(chan []*node, alpha)
-		pendingQueries = 0
-		result         *nodesByDistance
-	)
-	// Don't query further if we hit ourself.
-	// Unlikely to happen often in practice.
-	asked[t.Self().ID()] = true
-
-	// Generate the initial result set.
-	t.tab.mutex.Lock()
-	result = t.tab.closest(target, bucketSize, false)
-	t.tab.mutex.Unlock()
-
-	for {
-		// ask the alpha closest nodes that we haven't asked yet
-		for i := 0; i < len(result.entries) && pendingQueries < alpha; i++ {
-			n := result.entries[i]
-			if !asked[n.ID()] {
-				asked[n.ID()] = true
-				pendingQueries++
-				go t.lookupWorker(n, targetKey, reply)
-			}
-		}
-		if pendingQueries == 0 {
-			// we have asked all closest nodes, stop the search
-			break
-		}
-		select {
-		case nodes := <-reply:
-			for _, n := range nodes {
-				if n != nil && !seen[n.ID()] {
-					seen[n.ID()] = true
-					result.push(n, bucketSize)
-				}
-			}
-		case <-t.tab.closeReq:
-			return nil // shutdown, no need to continue.
-		}
-		pendingQueries--
-	}
-	return result.entries
-}
-
-func (t *UDPv4) lookupWorker(n *node, targetKey encPubkey, reply chan<- []*node) {
-	fails := t.db.FindFails(n.ID(), n.IP())
-	r, err := t.findnode(n.ID(), n.addr(), targetKey)
-	if err == errClosed {
-		// Avoid recording failures on shutdown.
-		reply <- nil
-		return
-	} else if len(r) == 0 {
-		fails++
-		t.db.UpdateFindFails(n.ID(), n.IP(), fails)
-		t.log.Trace("Findnode failed", "id", n.ID(), "failcount", fails, "err", err)
-		if fails >= maxFindnodeFailures {
-			t.log.Trace("Too many findnode failures, dropping", "id", n.ID(), "failcount", fails)
-			t.tab.delete(n)
-		}
-	} else if fails > 0 {
-		// Reset failure counter because it counts _consecutive_ failures.
-		t.db.UpdateFindFails(n.ID(), n.IP(), 0)
-	}
-
-	// Grab as many nodes as possible. Some of them might not be alive anymore, but we'll
-	// just remove those again during revalidation.
-	for _, n := range r {
-		t.tab.addSeenNode(n)
-	}
-	reply <- r
-}
-
 // Resolve searches for a specific node with the given ID and tries to get the most recent
 // version of the node record for it. It returns n if the node could not be resolved.
 func (t *UDPv4) Resolve(n *enode.Node) *enode.Node {
@@ -498,6 +390,45 @@ func (t *UDPv4) makePing(toaddr *net.UDPAddr) *pingV4 {
 	}
 }
 
+// LookupPubkey finds the closest nodes to the given public key.
+func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node {
+	if t.tab.len() == 0 {
+		// All nodes were dropped, refresh. The very first query will hit this
+		// case and run the bootstrapping logic.
+		<-t.tab.refresh()
+	}
+	return t.newLookup(t.closeCtx, encodePubkey(key)).run()
+}
+
+// RandomNodes is an iterator yielding nodes from a random walk of the DHT.
+func (t *UDPv4) RandomNodes() enode.Iterator {
+	return newLookupIterator(t.closeCtx, t.newRandomLookup)
+}
+
+// lookupRandom implements transport.
+func (t *UDPv4) lookupRandom() []*enode.Node {
+	return t.newRandomLookup(t.closeCtx).run()
+}
+
+// lookupSelf implements transport.
+func (t *UDPv4) lookupSelf() []*enode.Node {
+	return t.newLookup(t.closeCtx, encodePubkey(&t.priv.PublicKey)).run()
+}
+
+func (t *UDPv4) newRandomLookup(ctx context.Context) *lookup {
+	var target encPubkey
+	crand.Read(target[:])
+	return t.newLookup(ctx, target)
+}
+
+func (t *UDPv4) newLookup(ctx context.Context, targetKey encPubkey) *lookup {
+	target := enode.ID(crypto.Keccak256Hash(targetKey[:]))
+	it := newLookup(ctx, t.tab, target, func(n *node) ([]*node, error) {
+		return t.findnode(n.ID(), n.addr(), targetKey)
+	})
+	return it
+}
+
 // findnode sends a findnode request to the given node and waits until
 // the node has sent up to k neighbors.
 func (t *UDPv4) findnode(toid enode.ID, toaddr *net.UDPAddr, target encPubkey) ([]*node, error) {
@@ -575,7 +506,7 @@ func (t *UDPv4) pending(id enode.ID, ip net.IP, ptype byte, callback replyMatchF
 	select {
 	case t.addReplyMatcher <- p:
 		// loop will handle it
-	case <-t.closing:
+	case <-t.closeCtx.Done():
 		ch <- errClosed
 	}
 	return p
@@ -589,7 +520,7 @@ func (t *UDPv4) handleReply(from enode.ID, fromIP net.IP, req packetV4) bool {
 	case t.gotreply <- reply{from, fromIP, req, matched}:
 		// loop will handle it
 		return <-matched
-	case <-t.closing:
+	case <-t.closeCtx.Done():
 		return false
 	}
 }
@@ -635,7 +566,7 @@ func (t *UDPv4) loop() {
 		resetTimeout()
 
 		select {
-		case <-t.closing:
+		case <-t.closeCtx.Done():
 			for el := plist.Front(); el != nil; el = el.Next() {
 				el.Value.(*replyMatcher).errc <- errClosed
 			}