From fb2ae8e9955d2cbb65383e6ff0f3ba73e88c0b09 Mon Sep 17 00:00:00 2001
From: rjl493456442 <garyrong0905@gmail.com>
Date: Tue, 15 Mar 2022 17:23:37 +0800
Subject: [PATCH] trie: fix two issues in trie iterator (#24539)

* trie: fix memory leak in trie iterator

In the trie iterator, live nodes are tracked in a stack while iterating.
Popped node states should be explictly set to nil in order to get
garbage-collected.

* trie: fix empty trie iterator
---
 trie/iterator.go      | 14 +++++++++-----
 trie/iterator_test.go | 13 +++++++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/trie/iterator.go b/trie/iterator.go
index 1fbf991af..e0006ee05 100644
--- a/trie/iterator.go
+++ b/trie/iterator.go
@@ -154,8 +154,11 @@ func (e seekError) Error() string {
 }
 
 func newNodeIterator(trie *Trie, start []byte) NodeIterator {
-	if trie.Hash() == emptyState {
-		return new(nodeIterator)
+	if trie.Hash() == emptyRoot {
+		return &nodeIterator{
+			trie: trie,
+			err:  errIteratorEnd,
+		}
 	}
 	it := &nodeIterator{trie: trie}
 	it.err = it.seek(start)
@@ -425,7 +428,7 @@ func findChild(n *fullNode, index int, path []byte, ancestor common.Hash) (node,
 func (it *nodeIterator) nextChild(parent *nodeIteratorState, ancestor common.Hash) (*nodeIteratorState, []byte, bool) {
 	switch node := parent.node.(type) {
 	case *fullNode:
-		//Full node, move to the first non-nil child.
+		// Full node, move to the first non-nil child.
 		if child, state, path, index := findChild(node, parent.index+1, it.path, ancestor); child != nil {
 			parent.index = index - 1
 			return state, path, true
@@ -503,8 +506,9 @@ func (it *nodeIterator) push(state *nodeIteratorState, parentIndex *int, path []
 }
 
 func (it *nodeIterator) pop() {
-	parent := it.stack[len(it.stack)-1]
-	it.path = it.path[:parent.pathlen]
+	last := it.stack[len(it.stack)-1]
+	it.path = it.path[:last.pathlen]
+	it.stack[len(it.stack)-1] = nil
 	it.stack = it.stack[:len(it.stack)-1]
 }
 
diff --git a/trie/iterator_test.go b/trie/iterator_test.go
index 1a9a5e80e..9a46e9b99 100644
--- a/trie/iterator_test.go
+++ b/trie/iterator_test.go
@@ -29,6 +29,19 @@ import (
 	"github.com/ethereum/go-ethereum/ethdb/memorydb"
 )
 
+func TestEmptyIterator(t *testing.T) {
+	trie := newEmpty()
+	iter := trie.NodeIterator(nil)
+
+	seen := make(map[string]struct{})
+	for iter.Next(true) {
+		seen[string(iter.Path())] = struct{}{}
+	}
+	if len(seen) != 0 {
+		t.Fatal("Unexpected trie node iterated")
+	}
+}
+
 func TestIterator(t *testing.T) {
 	trie := newEmpty()
 	vals := []struct{ k, v string }{