From 36f81118f6d00713dbf8d1caf168ad937e0f1a90 Mon Sep 17 00:00:00 2001
From: gary rong <garyrong0905@gmail.com>
Date: Fri, 5 Apr 2019 14:44:02 +0800
Subject: [PATCH] core/state: fix state iterator (#19127)

* core/state: fix state iterator

* core: fix state iterator more elegant
---
 core/state/statedb.go | 21 +++++++++++++++++----
 core/vm/interface.go  |  2 +-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/core/state/statedb.go b/core/state/statedb.go
index db2e0d5e0..3bb9862ed 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -515,20 +515,33 @@ func (self *StateDB) CreateAccount(addr common.Address) {
 	}
 }
 
-func (db *StateDB) ForEachStorage(addr common.Address, cb func(key, value common.Hash) bool) {
+func (db *StateDB) ForEachStorage(addr common.Address, cb func(key, value common.Hash) bool) error {
 	so := db.getStateObject(addr)
 	if so == nil {
-		return
+		return nil
 	}
 	it := trie.NewIterator(so.getTrie(db.db).NodeIterator(nil))
+
 	for it.Next() {
 		key := common.BytesToHash(db.trie.GetKey(it.Key))
 		if value, dirty := so.dirtyStorage[key]; dirty {
-			cb(key, value)
+			if !cb(key, value) {
+				return nil
+			}
 			continue
 		}
-		cb(key, common.BytesToHash(it.Value))
+
+		if len(it.Value) > 0 {
+			_, content, _, err := rlp.Split(it.Value)
+			if err != nil {
+				return err
+			}
+			if !cb(key, common.BytesToHash(content)) {
+				return nil
+			}
+		}
 	}
+	return nil
 }
 
 // Copy creates a deep, independent copy of the state.
diff --git a/core/vm/interface.go b/core/vm/interface.go
index fc15082f1..dd401466a 100644
--- a/core/vm/interface.go
+++ b/core/vm/interface.go
@@ -63,7 +63,7 @@ type StateDB interface {
 	AddLog(*types.Log)
 	AddPreimage(common.Hash, []byte)
 
-	ForEachStorage(common.Address, func(common.Hash, common.Hash) bool)
+	ForEachStorage(common.Address, func(common.Hash, common.Hash) bool) error
 }
 
 // CallContext provides a basic interface for the EVM calling conventions. The EVM