From 2b9e2c547a639102853b8bd3b3c2f93655c4f59c 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 af787aa7b520..d2c01662ca6b 100644
--- a/trie/iterator.go
+++ b/trie/iterator.go
@@ -151,8 +151,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)
@@ -402,7 +405,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
@@ -480,8 +483,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 46920e9541fa..761d812bdfcc 100644
--- a/trie/iterator_test.go
+++ b/trie/iterator_test.go
@@ -29,6 +29,19 @@ import (
 	"github.com/scroll-tech/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 }{