From 16bc57438bd3d28e947d12f6f295da62e4ca9e26 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Thu, 20 May 2021 09:24:41 +0200
Subject: [PATCH] p2p/dnsdisc: fix crash when iterator closed before first call
 to Next (#22906)

---
 p2p/dnsdisc/client.go      |  6 ++++++
 p2p/dnsdisc/client_test.go | 15 +++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/p2p/dnsdisc/client.go b/p2p/dnsdisc/client.go
index f2a4bed4c6e4..096df06a54f9 100644
--- a/p2p/dnsdisc/client.go
+++ b/p2p/dnsdisc/client.go
@@ -298,6 +298,12 @@ func (it *randomIterator) pickTree() *clientTree {
 	it.mu.Lock()
 	defer it.mu.Unlock()
 
+	// First check if iterator was closed.
+	// Need to do this here to avoid nil map access in rebuildTrees.
+	if it.trees == nil {
+		return nil
+	}
+
 	// Rebuild the trees map if any links have changed.
 	if it.lc.changed {
 		it.rebuildTrees()
diff --git a/p2p/dnsdisc/client_test.go b/p2p/dnsdisc/client_test.go
index 741bee4230be..9320dd667a94 100644
--- a/p2p/dnsdisc/client_test.go
+++ b/p2p/dnsdisc/client_test.go
@@ -115,6 +115,21 @@ func TestIterator(t *testing.T) {
 	checkIterator(t, it, nodes)
 }
 
+func TestIteratorCloseWithoutNext(t *testing.T) {
+	tree1, url1 := makeTestTree("t1", nil, nil)
+	c := NewClient(Config{Resolver: newMapResolver(tree1.ToTXT("t1"))})
+	it, err := c.NewIterator(url1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	it.Close()
+	ok := it.Next()
+	if ok {
+		t.Fatal("Next returned true after Close")
+	}
+}
+
 // This test checks if closing randomIterator races.
 func TestIteratorClose(t *testing.T) {
 	nodes := testNodes(nodesSeed1, 500)