From 43c278cdf93d5469702fd1c2f570dbf3c1718ff0 Mon Sep 17 00:00:00 2001
From: gary rong <garyrong0905@gmail.com>
Date: Wed, 28 Oct 2020 20:27:37 +0800
Subject: [PATCH] core/state: disable snapshot iteration if it's not fully
 constructed (#21682)

* core/state/snapshot: add diskRoot function

* core/state/snapshot: disable iteration if the snapshot is generating

* core/state/snapshot: simplify the function

* core/state: panic for undefined layer
---
 core/state/snapshot/snapshot.go | 54 +++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/core/state/snapshot/snapshot.go b/core/state/snapshot/snapshot.go
index f6c5a6a9a8ca..4f496b4ff64f 100644
--- a/core/state/snapshot/snapshot.go
+++ b/core/state/snapshot/snapshot.go
@@ -86,6 +86,10 @@ var (
 	// range of accounts covered.
 	ErrNotCoveredYet = errors.New("not covered yet")
 
+	// ErrNotConstructed is returned if the callers want to iterate the snapshot
+	// while the generation is not finished yet.
+	ErrNotConstructed = errors.New("snapshot is not constructed")
+
 	// errSnapshotCycle is returned if a snapshot is attempted to be inserted
 	// that forms a cycle in the snapshot tree.
 	errSnapshotCycle = errors.New("snapshot cycle")
@@ -609,11 +613,61 @@ func (t *Tree) Rebuild(root common.Hash) {
 // AccountIterator creates a new account iterator for the specified root hash and
 // seeks to a starting account hash.
 func (t *Tree) AccountIterator(root common.Hash, seek common.Hash) (AccountIterator, error) {
+	ok, err := t.generating()
+	if err != nil {
+		return nil, err
+	}
+	if ok {
+		return nil, ErrNotConstructed
+	}
 	return newFastAccountIterator(t, root, seek)
 }
 
 // StorageIterator creates a new storage iterator for the specified root hash and
 // account. The iterator will be move to the specific start position.
 func (t *Tree) StorageIterator(root common.Hash, account common.Hash, seek common.Hash) (StorageIterator, error) {
+	ok, err := t.generating()
+	if err != nil {
+		return nil, err
+	}
+	if ok {
+		return nil, ErrNotConstructed
+	}
 	return newFastStorageIterator(t, root, account, seek)
 }
+
+// disklayer is an internal helper function to return the disk layer.
+// The lock of snapTree is assumed to be held already.
+func (t *Tree) disklayer() *diskLayer {
+	var snap snapshot
+	for _, s := range t.layers {
+		snap = s
+		break
+	}
+	if snap == nil {
+		return nil
+	}
+	switch layer := snap.(type) {
+	case *diskLayer:
+		return layer
+	case *diffLayer:
+		return layer.origin
+	default:
+		panic(fmt.Sprintf("%T: undefined layer", snap))
+	}
+}
+
+// generating is an internal helper function which reports whether the snapshot
+// is still under the construction.
+func (t *Tree) generating() (bool, error) {
+	t.lock.Lock()
+	defer t.lock.Unlock()
+
+	layer := t.disklayer()
+	if layer == nil {
+		return false, errors.New("disk layer is missing")
+	}
+	layer.lock.RLock()
+	defer layer.lock.RUnlock()
+	return layer.genMarker != nil, nil
+}