p2p/discover: integrate iterator with UDPv4

This commit is contained in:
Felix Lange 2019-07-05 12:04:59 +02:00
parent 91063c264e
commit 14858684c8

View file

@ -202,6 +202,7 @@ type UDPv4 struct {
localNode *enode.LocalNode localNode *enode.LocalNode
db *enode.DB db *enode.DB
tab *Table tab *Table
randomWalk *lookupWalker
closeOnce sync.Once closeOnce sync.Once
wg sync.WaitGroup wg sync.WaitGroup
@ -270,6 +271,7 @@ func ListenV4(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) {
if t.log == nil { if t.log == nil {
t.log = log.Root() t.log = log.Root()
} }
tab, err := newTable(t, ln.Database(), cfg.Bootnodes, t.log) tab, err := newTable(t, ln.Database(), cfg.Bootnodes, t.log)
if err != nil { if err != nil {
return nil, err return nil, err
@ -277,6 +279,8 @@ func ListenV4(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) {
t.tab = tab t.tab = tab
go tab.loop() go tab.loop()
t.randomWalk = newLookupWalker(t.randomLookupWithCallback)
t.wg.Add(2) t.wg.Add(2)
go t.loop() go t.loop()
go t.readLoop(cfg.Unhandled) go t.readLoop(cfg.Unhandled)
@ -295,22 +299,25 @@ func (t *UDPv4) Close() {
t.conn.Close() t.conn.Close()
t.wg.Wait() t.wg.Wait()
t.tab.close() t.tab.close()
t.randomWalk.close()
}) })
} }
// ReadRandomNodes reads random nodes from the local table. // RandomNodes is an iterator yielding nodes from a random walk of the DHT.
func (t *UDPv4) ReadRandomNodes(buf []*enode.Node) int { func (t *UDPv4) RandomNodes() Iterator {
return t.tab.ReadRandomNodes(buf) return t.randomWalk.newIterator()
} }
// LookupRandom finds random nodes in the network. // LookupRandom finds random nodes in the network.
func (t *UDPv4) LookupRandom() []*enode.Node { func (t *UDPv4) randomLookupWithCallback(callback func(*enode.Node)) {
if t.tab.len() == 0 { if t.tab.len() == 0 {
// All nodes were dropped, refresh. The very first query will hit this // All nodes were dropped, refresh. The very first query will hit this
// case and run the bootstrapping logic. // case and run the bootstrapping logic.
<-t.tab.refresh() <-t.tab.refresh()
} }
return t.lookupRandom() var target encPubkey
crand.Read(target[:])
t.lookup(target, callback)
} }
func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node { func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node {
@ -319,23 +326,23 @@ func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node {
// case and run the bootstrapping logic. // case and run the bootstrapping logic.
<-t.tab.refresh() <-t.tab.refresh()
} }
return unwrapNodes(t.lookup(encodePubkey(key))) return unwrapNodes(t.lookup(encodePubkey(key), nil))
} }
func (t *UDPv4) lookupRandom() []*enode.Node { func (t *UDPv4) lookupRandom() []*enode.Node {
var target encPubkey var target encPubkey
crand.Read(target[:]) crand.Read(target[:])
return unwrapNodes(t.lookup(target)) return unwrapNodes(t.lookup(target, nil))
} }
func (t *UDPv4) lookupSelf() []*enode.Node { func (t *UDPv4) lookupSelf() []*enode.Node {
return unwrapNodes(t.lookup(encodePubkey(&t.priv.PublicKey))) return unwrapNodes(t.lookup(encodePubkey(&t.priv.PublicKey), nil))
} }
// lookup performs a network search for nodes close to the given target. It approaches the // 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 // target by querying nodes that are closer to it on each iteration. The given target does
// not need to be an actual node identifier. // not need to be an actual node identifier.
func (t *UDPv4) lookup(targetKey encPubkey) []*node { func (t *UDPv4) lookup(targetKey encPubkey, nodeCallback func(*enode.Node)) []*node {
var ( var (
target = enode.ID(crypto.Keccak256Hash(targetKey[:])) target = enode.ID(crypto.Keccak256Hash(targetKey[:]))
asked = make(map[enode.ID]bool) asked = make(map[enode.ID]bool)
@ -360,7 +367,7 @@ func (t *UDPv4) lookup(targetKey encPubkey) []*node {
if !asked[n.ID()] { if !asked[n.ID()] {
asked[n.ID()] = true asked[n.ID()] = true
pendingQueries++ pendingQueries++
go t.lookupWorker(n, targetKey, reply) go t.lookupWorker(n, targetKey, reply, nodeCallback)
} }
} }
if pendingQueries == 0 { if pendingQueries == 0 {
@ -383,7 +390,7 @@ func (t *UDPv4) lookup(targetKey encPubkey) []*node {
return result.entries return result.entries
} }
func (t *UDPv4) lookupWorker(n *node, targetKey encPubkey, reply chan<- []*node) { func (t *UDPv4) lookupWorker(n *node, targetKey encPubkey, reply chan<- []*node, callback func(*enode.Node)) {
fails := t.db.FindFails(n.ID(), n.IP()) fails := t.db.FindFails(n.ID(), n.IP())
r, err := t.findnode(n.ID(), n.addr(), targetKey) r, err := t.findnode(n.ID(), n.addr(), targetKey)
if err == errClosed { if err == errClosed {
@ -407,7 +414,11 @@ func (t *UDPv4) lookupWorker(n *node, targetKey encPubkey, reply chan<- []*node)
// just remove those again during revalidation. // just remove those again during revalidation.
for _, n := range r { for _, n := range r {
t.tab.addSeenNode(n) t.tab.addSeenNode(n)
if callback != nil {
callback(unwrapNode(n))
}
} }
reply <- r reply <- r
} }