p2p/dnsdisc: add support for enode.Iterator

This changes the dnsdisc.Client API to support the enode.Iterator
interface.

p2p/dnsdisc/client.go

@@ -23,6 +23,7 @@ import (
 	"math/rand"
 	"net"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common/mclock"
@@ -37,9 +38,6 @@ import (
 type Client struct {
 	cfg     Config
 	clock   mclock.Clock
-	linkCache linkCache
-	trees     map[string]*clientTree
-
 	entries *lru.Cache
 }
 
@@ -86,32 +84,25 @@ func (cfg Config) withDefaults() Config {
 }
 
 // NewClient creates a client.
-func NewClient(cfg Config, urls ...string) (*Client, error) {
+func NewClient(cfg Config) *Client {
 	c := &Client{
 		cfg:   cfg.withDefaults(),
 		clock: mclock.System{},
-		trees: make(map[string]*clientTree),
 	}
 	var err error
 	if c.entries, err = lru.New(c.cfg.CacheLimit); err != nil {
-		return nil, err
+		panic(err)
 	}
-	for _, url := range urls {
+	return c
-		if err := c.AddTree(url); err != nil {
-			return nil, err
-		}
-	}
-	return c, nil
 }
 
-// SyncTree downloads the entire node tree at the given URL. This doesn't add the tree for
+// SyncTree downloads the entire node tree at the given URL.
-// later use, but any previously-synced entries are reused.
 func (c *Client) SyncTree(url string) (*Tree, error) {
 	le, err := parseLink(url)
 	if err != nil {
 		return nil, fmt.Errorf("invalid enrtree URL: %v", err)
 	}
-	ct := newClientTree(c, le)
+	ct := newClientTree(c, new(linkCache), le)
 	t := &Tree{entries: make(map[string]entry)}
 	if err := ct.syncAll(t.entries); err != nil {
 		return nil, err
@@ -120,75 +111,16 @@ func (c *Client) SyncTree(url string) (*Tree, error) {
 	return t, nil
 }
 
-// AddTree adds a enrtree:// URL to crawl.
+// NewIterator creates an iterator that visits all nodes at the
-func (c *Client) AddTree(url string) error {
+// given tree URLs.
-	le, err := parseLink(url)
+func (c *Client) NewIterator(urls ...string) (enode.Iterator, error) {
-	if err != nil {
+	it := c.newRandomIterator()
-		return fmt.Errorf("invalid enrtree URL: %v", err)
+	for _, url := range urls {
-	}
+		if err := it.addTree(url); err != nil {
-	ct, err := c.ensureTree(le)
+			return nil, err
-	if err != nil {
-		return err
-	}
-	c.linkCache.add(ct)
-	return nil
-}
-
-func (c *Client) ensureTree(le *linkEntry) (*clientTree, error) {
-	if tree, ok := c.trees[le.domain]; ok {
-		if !tree.matchPubkey(le.pubkey) {
-			return nil, fmt.Errorf("conflicting public keys for domain %q", le.domain)
-		}
-		return tree, nil
-	}
-	ct := newClientTree(c, le)
-	c.trees[le.domain] = ct
-	return ct, nil
-}
-
-// RandomNode retrieves the next random node.
-func (c *Client) RandomNode(ctx context.Context) *enode.Node {
-	for {
-		ct := c.randomTree()
-		if ct == nil {
-			return nil
-		}
-		n, err := ct.syncRandom(ctx)
-		if err != nil {
-			if err == ctx.Err() {
-				return nil // context canceled.
-			}
-			c.cfg.Logger.Debug("Error in DNS random node sync", "tree", ct.loc.domain, "err", err)
-			continue
-		}
-		if n != nil {
-			return n
 		}
 	}
-}
+	return it, nil
-
-// randomTree returns a random tree.
-func (c *Client) randomTree() *clientTree {
-	if !c.linkCache.valid() {
-		c.gcTrees()
-	}
-	limit := rand.Intn(len(c.trees))
-	for _, ct := range c.trees {
-		if limit == 0 {
-			return ct
-		}
-		limit--
-	}
-	return nil
-}
-
-// gcTrees rebuilds the 'trees' map.
-func (c *Client) gcTrees() {
-	trees := make(map[string]*clientTree)
-	for t := range c.linkCache.all() {
-		trees[t.loc.domain] = t
-	}
-	c.trees = trees
 }
 
 // resolveRoot retrieves a root entry via DNS.
@@ -258,3 +190,115 @@ func (c *Client) doResolveEntry(ctx context.Context, domain, hash string) (entry
 	}
 	return nil, nameError{name, errNoEntry}
 }
+
+// randomIterator traverses a set of trees and returns nodes found in them.
+type randomIterator struct {
+	cur      *enode.Node
+	ctx      context.Context
+	cancelFn context.CancelFunc
+	c        *Client
+
+	mu    sync.Mutex
+	trees map[string]*clientTree // all trees
+	lc    linkCache              // tracks tree dependencies
+}
+
+func (c *Client) newRandomIterator() *randomIterator {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &randomIterator{
+		c:        c,
+		ctx:      ctx,
+		cancelFn: cancel,
+		trees:    make(map[string]*clientTree),
+	}
+}
+
+// Node returns the current node.
+func (it *randomIterator) Node() *enode.Node {
+	return it.cur
+}
+
+// Close closes the iterator.
+func (it *randomIterator) Close() {
+	it.mu.Lock()
+	defer it.mu.Unlock()
+
+	it.cancelFn()
+	it.trees = nil
+}
+
+// Next moves the iterator to the next node.
+func (it *randomIterator) Next() bool {
+	it.cur = it.nextNode()
+	return it.cur != nil
+}
+
+// addTree adds a enrtree:// URL to the iterator.
+func (it *randomIterator) addTree(url string) error {
+	le, err := parseLink(url)
+	if err != nil {
+		return fmt.Errorf("invalid enrtree URL: %v", err)
+	}
+	it.lc.addLink("", le.str)
+	return nil
+}
+
+// nextNode syncs random tree entries until it finds a node.
+func (it *randomIterator) nextNode() *enode.Node {
+	for {
+		ct := it.nextTree()
+		if ct == nil {
+			return nil
+		}
+		n, err := ct.syncRandom(it.ctx)
+		if err != nil {
+			if err == it.ctx.Err() {
+				return nil // context canceled.
+			}
+			it.c.cfg.Logger.Debug("Error in DNS random node sync", "tree", ct.loc.domain, "err", err)
+			continue
+		}
+		if n != nil {
+			return n
+		}
+	}
+}
+
+// nextTree returns a random tree.
+func (it *randomIterator) nextTree() *clientTree {
+	it.mu.Lock()
+	defer it.mu.Unlock()
+
+	if it.lc.changed {
+		it.rebuildTrees()
+		it.lc.changed = false
+	}
+	if len(it.trees) == 0 {
+		return nil
+	}
+	limit := rand.Intn(len(it.trees))
+	for _, ct := range it.trees {
+		if limit == 0 {
+			return ct
+		}
+		limit--
+	}
+	return nil
+}
+
+// rebuildTrees rebuilds the 'trees' map.
+func (it *randomIterator) rebuildTrees() {
+	// Delete removed trees.
+	for loc := range it.trees {
+		if !it.lc.isReferenced(loc) {
+			delete(it.trees, loc)
+		}
+	}
+	// Add new trees.
+	for loc := range it.lc.backrefs {
+		if it.trees[loc] == nil {
+			link, _ := parseLink(linkPrefix + loc)
+			it.trees[loc] = newClientTree(it.c, &it.lc, link)
+		}
+	}
+}

p2p/dnsdisc/client_test.go

@@ -54,7 +54,7 @@ func TestClientSyncTree(t *testing.T) {
 		wantSeq   = uint(1)
 	)
 
-	c, _ := NewClient(Config{Resolver: r, Logger: testlog.Logger(t, log.LvlTrace)})
+	c := NewClient(Config{Resolver: r, Logger: testlog.Logger(t, log.LvlTrace)})
 	stree, err := c.SyncTree("enrtree://AKPYQIUQIL7PSIACI32J7FGZW56E5FKHEFCCOFHILBIMW3M6LWXS2@n")
 	if err != nil {
 		t.Fatal("sync error:", err)
@@ -68,9 +68,6 @@ func TestClientSyncTree(t *testing.T) {
 	if stree.Seq() != wantSeq {
 		t.Errorf("synced tree has wrong seq: %d", stree.Seq())
 	}
-	if len(c.trees) > 0 {
-		t.Errorf("tree from SyncTree added to client")
-	}
 }
 
 // In this test, syncing the tree fails because it contains an invalid ENR entry.
@@ -91,7 +88,7 @@ func TestClientSyncTreeBadNode(t *testing.T) {
 		"C7HRFPF3BLGF3YR4DY5KX3SMBE.n": "enrtree://AM5FCQLWIZX2QFPNJAP7VUERCCRNGRHWZG3YYHIUV7BVDQ5FDPRT2@morenodes.example.org",
 		"INDMVBZEEQ4ESVYAKGIYU74EAA.n": "enr:-----",
 	}
-	c, _ := NewClient(Config{Resolver: r, Logger: testlog.Logger(t, log.LvlTrace)})
+	c := NewClient(Config{Resolver: r, Logger: testlog.Logger(t, log.LvlTrace)})
 	_, err := c.SyncTree("enrtree://AKPYQIUQIL7PSIACI32J7FGZW56E5FKHEFCCOFHILBIMW3M6LWXS2@n")
 	wantErr := nameError{name: "INDMVBZEEQ4ESVYAKGIYU74EAA.n", err: entryError{typ: "enr", err: errInvalidENR}}
 	if err != wantErr {
@@ -99,57 +96,83 @@ func TestClientSyncTreeBadNode(t *testing.T) {
 	}
 }
 
-// This test checks that RandomNode hits all entries.
+// This test checks that randomIterator finds all entries.
-func TestClientRandomNode(t *testing.T) {
+func TestIterator(t *testing.T) {
 	nodes := testNodes(nodesSeed1, 30)
 	tree, url := makeTestTree("n", nodes, nil)
 	r := mapResolver(tree.ToTXT("n"))
-	c, _ := NewClient(Config{Resolver: r, Logger: testlog.Logger(t, log.LvlTrace)})
+	c := NewClient(Config{Resolver: r, Logger: testlog.Logger(t, log.LvlTrace)})
-	if err := c.AddTree(url); err != nil {
+	it, err := c.NewIterator(url)
+	if err != nil {
 		t.Fatal(err)
 	}
 
-	checkRandomNode(t, c, nodes)
+	checkIterator(t, it, nodes)
 }
 
-// This test checks that RandomNode traverses linked trees as well as explicitly added trees.
+// This test checks if closing randomIterator races.
-func TestClientRandomNodeLinks(t *testing.T) {
+func TestIteratorClose(t *testing.T) {
+	nodes := testNodes(nodesSeed1, 500)
+	tree1, url1 := makeTestTree("t1", nodes, nil)
+	c := NewClient(Config{Resolver: newMapResolver(tree1.ToTXT("t1"))})
+	it, err := c.NewIterator(url1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	done := make(chan struct{})
+	go func() {
+		for it.Next() {
+			_ = it.Node()
+		}
+		close(done)
+	}()
+
+	time.Sleep(50 * time.Millisecond)
+	it.Close()
+	<-done
+}
+
+// This test checks that randomIterator traverses linked trees as well as explicitly added trees.
+func TestIteratorLinks(t *testing.T) {
 	nodes := testNodes(nodesSeed1, 40)
 	tree1, url1 := makeTestTree("t1", nodes[:10], nil)
 	tree2, url2 := makeTestTree("t2", nodes[10:], []string{url1})
-	cfg := Config{
+	c := NewClient(Config{
 		Resolver: newMapResolver(tree1.ToTXT("t1"), tree2.ToTXT("t2")),
 		Logger:   testlog.Logger(t, log.LvlTrace),
-	}
+	})
-	c, _ := NewClient(cfg)
+	it, err := c.NewIterator(url2)
-	if err := c.AddTree(url2); err != nil {
+	if err != nil {
 		t.Fatal(err)
 	}
 
-	checkRandomNode(t, c, nodes)
+	checkIterator(t, it, nodes)
 }
 
-// This test verifies that RandomNode re-checks the root of the tree to catch
+// This test verifies that randomIterator re-checks the root of the tree to catch
 // updates to nodes.
-func TestClientRandomNodeUpdates(t *testing.T) {
+func TestIteratorNodeUpdates(t *testing.T) {
 	var (
 		clock    = new(mclock.Simulated)
 		nodes    = testNodes(nodesSeed1, 30)
 		resolver = newMapResolver()
-		cfg      = Config{
+		c        = NewClient(Config{
 			Resolver:        resolver,
 			Logger:          testlog.Logger(t, log.LvlTrace),
 			RecheckInterval: 20 * time.Minute,
-		}
+		})
-		c, _ = NewClient(cfg)
 	)
 	c.clock = clock
 	tree1, url := makeTestTree("n", nodes[:25], nil)
+	it, err := c.NewIterator(url)
+	if err != nil {
+		t.Fatal(err)
+	}
 
-	// Sync the original tree.
+	// sync the original tree.
 	resolver.add(tree1.ToTXT("n"))
-	c.AddTree(url)
+	checkIterator(t, it, nodes[:25])
-	checkRandomNode(t, c, nodes[:25])
 
 	// Update some nodes and ensure RandomNode returns the new nodes as well.
 	keys := testKeys(nodesSeed1, len(nodes))
@@ -162,25 +185,24 @@ func TestClientRandomNodeUpdates(t *testing.T) {
 		nodes[i] = n2
 	}
 	tree2, _ := makeTestTree("n", nodes, nil)
-	clock.Run(cfg.RecheckInterval + 1*time.Second)
+	clock.Run(c.cfg.RecheckInterval + 1*time.Second)
 	resolver.clear()
 	resolver.add(tree2.ToTXT("n"))
-	checkRandomNode(t, c, nodes)
+	checkIterator(t, it, nodes)
 }
 
-// This test verifies that RandomNode re-checks the root of the tree to catch
+// This test verifies that randomIterator re-checks the root of the tree to catch
 // updates to links.
-func TestClientRandomNodeLinkUpdates(t *testing.T) {
+func TestIteratorLinkUpdates(t *testing.T) {
 	var (
 		clock    = new(mclock.Simulated)
 		nodes    = testNodes(nodesSeed1, 30)
 		resolver = newMapResolver()
-		cfg      = Config{
+		c        = NewClient(Config{
 			Resolver:        resolver,
 			Logger:          testlog.Logger(t, log.LvlTrace),
 			RecheckInterval: 20 * time.Minute,
-		}
+		})
-		c, _ = NewClient(cfg)
 	)
 	c.clock = clock
 	tree3, url3 := makeTestTree("t3", nodes[20:30], nil)
@@ -190,49 +212,53 @@ func TestClientRandomNodeLinkUpdates(t *testing.T) {
 	resolver.add(tree2.ToTXT("t2"))
 	resolver.add(tree3.ToTXT("t3"))
 
+	it, err := c.NewIterator(url1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	// Sync tree1 using RandomNode.
-	c.AddTree(url1)
+	checkIterator(t, it, nodes[:20])
-	checkRandomNode(t, c, nodes[:20])
 
 	// Add link to tree3, remove link to tree2.
 	tree1, _ = makeTestTree("t1", nodes[:10], []string{url3})
 	resolver.add(tree1.ToTXT("t1"))
-	clock.Run(cfg.RecheckInterval + 1*time.Second)
+	clock.Run(c.cfg.RecheckInterval + 1*time.Second)
 	t.Log("tree1 updated")
 
 	var wantNodes []*enode.Node
 	wantNodes = append(wantNodes, tree1.Nodes()...)
 	wantNodes = append(wantNodes, tree3.Nodes()...)
-	checkRandomNode(t, c, wantNodes)
+	checkIterator(t, it, wantNodes)
 
 	// Check that linked trees are GCed when they're no longer referenced.
-	if len(c.trees) != 2 {
+	knownTrees := it.(*randomIterator).trees
-		t.Errorf("client knows %d trees, want 2", len(c.trees))
+	if len(knownTrees) != 2 {
+		t.Errorf("client knows %d trees, want 2", len(knownTrees))
 	}
 }
 
-func checkRandomNode(t *testing.T, c *Client, wantNodes []*enode.Node) {
+func checkIterator(t *testing.T, it enode.Iterator, wantNodes []*enode.Node) {
 	t.Helper()
 
 	var (
 		want     = make(map[enode.ID]*enode.Node)
 		maxCalls = len(wantNodes) * 2
 		calls    = 0
-		ctx      = context.Background()
 	)
 	for _, n := range wantNodes {
 		want[n.ID()] = n
 	}
 	for ; len(want) > 0 && calls < maxCalls; calls++ {
-		n := c.RandomNode(ctx)
+		if !it.Next() {
-		if n == nil {
+			t.Fatalf("Next returned false (call %d)", calls)
-			t.Fatalf("RandomNode returned nil (call %d)", calls)
 		}
+		n := it.Node()
 		delete(want, n.ID())
 	}
-	t.Logf("checkRandomNode called RandomNode %d times to find %d nodes", calls, len(wantNodes))
+	t.Logf("checkIterator called Next %d times to find %d nodes", calls, len(wantNodes))
 	for _, n := range want {
-		t.Errorf("RandomNode didn't discover node %v", n.ID())
+		t.Errorf("iterator didn't discover node %v", n.ID())
 	}
 }
 

p2p/dnsdisc/sync.go

@@ -29,22 +29,16 @@ import (
 // clientTree is a full tree being synced.
 type clientTree struct {
 	c             *Client
+	lc            *linkCache
 	loc           *linkEntry
 	root          *rootEntry
 	lastRootCheck mclock.AbsTime // last revalidation of root
 	enrs          *subtreeSync
 	links         *subtreeSync
-	linkCache     linkCache
 }
 
-func newClientTree(c *Client, loc *linkEntry) *clientTree {
+func newClientTree(c *Client, lc *linkCache, loc *linkEntry) *clientTree {
-	ct := &clientTree{c: c, loc: loc}
+	return &clientTree{c: c, lc: lc, loc: loc}
-	ct.linkCache.self = ct
-	return ct
-}
-
-func (ct *clientTree) matchPubkey(key *ecdsa.PublicKey) bool {
-	return keysEqual(ct.loc.pubkey, key)
 }
 
 func keysEqual(k1, k2 *ecdsa.PublicKey) bool {
@@ -95,12 +89,8 @@ func (ct *clientTree) syncNextLink(ctx context.Context) error {
 	}
 	ct.links.missing = ct.links.missing[1:]
 
-	if le, ok := e.(*linkEntry); ok {
+	if dest, ok := e.(*linkEntry); ok {
-		lt, err := ct.c.ensureTree(le)
+		ct.lc.addLink(ct.loc.str, dest.str)
-		if err != nil {
-			return err
-		}
-		ct.linkCache.add(lt)
 	}
 	return nil
 }
@@ -150,7 +140,7 @@ func (ct *clientTree) updateRoot() error {
 	// Invalidate subtrees if changed.
 	if ct.links == nil || root.lroot != ct.links.root {
 		ct.links = newSubtreeSync(ct.c, ct.loc, root.lroot, true)
-		ct.linkCache.reset()
+		ct.lc.resetLinks(ct.loc.str)
 	}
 	if ct.enrs == nil || root.eroot != ct.enrs.root {
 		ct.enrs = newSubtreeSync(ct.c, ct.loc, root.eroot, false)
@@ -215,63 +205,48 @@ func (ts *subtreeSync) resolveNext(ctx context.Context, hash string) (entry, err
 	return e, nil
 }
 
-// linkCache tracks the links of a tree.
+// linkCache tracks links between trees.
 type linkCache struct {
-	self    *clientTree
+	backrefs map[string]map[string]struct{}
-	directM map[*clientTree]struct{} // direct links
+	changed  bool
-	allM    map[*clientTree]struct{} // direct & transitive links
 }
 
-// reset clears the cache.
+func (lc *linkCache) isReferenced(r string) bool {
-func (lc *linkCache) reset() {
+	return len(lc.backrefs[r]) != 0
-	lc.directM = nil
-	lc.allM = nil
 }
 
-// add adds a direct link to the cache.
+func (lc *linkCache) addLink(from, to string) {
-func (lc *linkCache) add(ct *clientTree) {
+	if _, ok := lc.backrefs[to][from]; ok {
-	if lc.directM == nil {
+		return
-		lc.directM = make(map[*clientTree]struct{})
 	}
-	if _, ok := lc.directM[ct]; !ok {
+
-		lc.invalidate()
+	if lc.backrefs == nil {
+		lc.backrefs = make(map[string]map[string]struct{})
 	}
-	lc.directM[ct] = struct{}{}
+	if _, ok := lc.backrefs[to]; !ok {
+		lc.backrefs[to] = make(map[string]struct{})
+	}
+	lc.backrefs[to][from] = struct{}{}
+	lc.changed = true
 }
 
-// invalidate resets the cache of transitive links.
+// resetLinks clears all links of the given tree.
-func (lc *linkCache) invalidate() {
+func (lc *linkCache) resetLinks(from string) {
-	lc.allM = nil
+	stk := []string{from}
-}
+	for len(stk) > 0 {
+		item := stk[len(stk)-1]
+		stk = stk[:len(stk)-1]
 
-// valid returns true when the cache of transitive links is up-to-date.
+		for r, refs := range lc.backrefs {
-func (lc *linkCache) valid() bool {
+			if _, ok := refs[item]; !ok {
-	// Re-check validity of child caches to catch updates.
+				continue
-	for ct := range lc.allM {
+			}
-		if ct != lc.self && !ct.linkCache.valid() {
+			lc.changed = true
-			lc.allM = nil
+			delete(refs, item)
-			break
+			if len(refs) == 0 {
+				delete(lc.backrefs, r)
+				stk = append(stk, r)
+			}
 		}
 	}
-	return lc.allM != nil
-}
-
-// all returns all trees reachable through the cache.
-func (lc *linkCache) all() map[*clientTree]struct{} {
-	if lc.valid() {
-		return lc.allM
-	}
-	// Remake lc.allM it by taking the union of all() across children.
-	m := make(map[*clientTree]struct{})
-	if lc.self != nil {
-		m[lc.self] = struct{}{}
-	}
-	for ct := range lc.directM {
-		m[ct] = struct{}{}
-		for lt := range ct.linkCache.all() {
-			m[lt] = struct{}{}
-		}
-	}
-	lc.allM = m
-	return m
 }

p2p/dnsdisc/sync_test.go

@@ -0,0 +1,83 @@
+// Copyright 2018 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 dnsdisc
+
+import (
+	"math/rand"
+	"strconv"
+	"testing"
+)
+
+func TestLinkCache(t *testing.T) {
+	var lc linkCache
+
+	// Check adding links.
+	lc.addLink("1", "2")
+	if !lc.changed {
+		t.Error("changed flag not set")
+	}
+	lc.changed = false
+	lc.addLink("1", "2")
+	if lc.changed {
+		t.Error("changed flag set after adding link that's already present")
+	}
+	lc.addLink("2", "3")
+	lc.addLink("3", "1")
+	lc.addLink("2", "4")
+	lc.changed = false
+
+	if !lc.isReferenced("3") {
+		t.Error("3 not referenced")
+	}
+	if lc.isReferenced("6") {
+		t.Error("6 is referenced")
+	}
+
+	lc.resetLinks("1")
+	if !lc.changed {
+		t.Error("changed flag not set")
+	}
+	if len(lc.backrefs) != 0 {
+		t.Logf("%+v", lc)
+		t.Error("reference maps should be empty")
+	}
+}
+
+func TestLinkCacheRandom(t *testing.T) {
+	tags := make([]string, 1000)
+	for i := range tags {
+		tags[i] = strconv.Itoa(i)
+	}
+
+	// Create random links.
+	var lc linkCache
+	var remove []string
+	for i := 0; i < 100; i++ {
+		a, b := tags[rand.Intn(len(tags))], tags[rand.Intn(len(tags))]
+		lc.addLink(a, b)
+		remove = append(remove, a)
+	}
+
+	// Remove all the links.
+	for _, s := range remove {
+		lc.resetLinks(s)
+	}
+	if len(lc.backrefs) != 0 {
+		t.Logf("%+v", lc)
+		t.Error("reference maps should be empty")
+	}
+}

p2p/dnsdisc/tree.go

@@ -48,7 +48,7 @@ func (t *Tree) Sign(key *ecdsa.PrivateKey, domain string) (url string, err error
 	}
 	root.sig = sig
 	t.root = &root
-	link := &linkEntry{domain, &key.PublicKey}
+	link := newLinkEntry(domain, &key.PublicKey)
 	return link.String(), nil
 }
 
@@ -209,6 +209,7 @@ type (
 		node *enode.Node
 	}
 	linkEntry struct {
+		str    string
 		domain string
 		pubkey *ecdsa.PublicKey
 	}
@@ -246,7 +247,8 @@ func (e *rootEntry) sigHash() []byte {
 
 func (e *rootEntry) verifySignature(pubkey *ecdsa.PublicKey) bool {
 	sig := e.sig[:crypto.RecoveryIDOffset] // remove recovery id
-	return crypto.VerifySignature(crypto.FromECDSAPub(pubkey), e.sigHash(), sig)
+	enckey := crypto.FromECDSAPub(pubkey)
+	return crypto.VerifySignature(enckey, e.sigHash(), sig)
 }
 
 func (e *branchEntry) String() string {
@@ -258,8 +260,13 @@ func (e *enrEntry) String() string {
 }
 
 func (e *linkEntry) String() string {
-	pubkey := b32format.EncodeToString(crypto.CompressPubkey(e.pubkey))
+	return linkPrefix + e.str
-	return fmt.Sprintf("%s%s@%s", linkPrefix, pubkey, e.domain)
+}
+
+func newLinkEntry(domain string, pubkey *ecdsa.PublicKey) *linkEntry {
+	key := b32format.EncodeToString(crypto.CompressPubkey(pubkey))
+	str := key + "@" + domain
+	return &linkEntry{str, domain, pubkey}
 }
 
 // Entry Parsing
@@ -319,7 +326,7 @@ func parseLink(e string) (*linkEntry, error) {
 	if err != nil {
 		return nil, entryError{"link", errBadPubkey}
 	}
-	return &linkEntry{domain, key}, nil
+	return &linkEntry{e, domain, key}, nil
 }
 
 func parseBranch(e string) (entry, error) {

p2p/dnsdisc/tree_test.go

@@ -91,7 +91,7 @@ func TestParseEntry(t *testing.T) {
 		// Links
 		{
 			input: "enrtree://AKPYQIUQIL7PSIACI32J7FGZW56E5FKHEFCCOFHILBIMW3M6LWXS2@nodes.example.org",
-			e:     &linkEntry{"nodes.example.org", &testkey.PublicKey},
+			e:     &linkEntry{"AKPYQIUQIL7PSIACI32J7FGZW56E5FKHEFCCOFHILBIMW3M6LWXS2@nodes.example.org", "nodes.example.org", &testkey.PublicKey},
 		},
 		{
 			input: "enrtree://nodes.example.org",