[evm] clear snapshots in RevertToSnapshot()

iotexproject · Oct 19, 2021 · fcb8d7b · fcb8d7b
1 parent 13f7452
commit fcb8d7b
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
diff --git a/action/protocol/execution/evm/evmstatedbadapter.go b/action/protocol/execution/evm/evmstatedbadapter.go
@@ -446,6 +446,16 @@ func (stateDB *StateDBAdapter) RevertToSnapshot(snapshot int) {
 	// restore the suicide accounts
 	stateDB.suicided = nil
 	stateDB.suicided = ds
+	if stateDB.fixSnapshotOrder {
+		delete(stateDB.suicideSnapshot, snapshot)
+		for i := snapshot + 1; ; i++ {
+			if _, ok := stateDB.suicideSnapshot[snapshot]; ok {
+				delete(stateDB.suicideSnapshot, i)
+			} else {
+				break
+			}
+		}
+	}
 	// restore modified contracts
 	stateDB.cachedContract = nil
 	stateDB.cachedContract = stateDB.contractSnapshot[snapshot]
@@ -456,9 +466,29 @@ func (stateDB *StateDBAdapter) RevertToSnapshot(snapshot int) {
 			return
 		}
 	}
+	if stateDB.fixSnapshotOrder {
+		delete(stateDB.contractSnapshot, snapshot)
+		for i := snapshot + 1; ; i++ {
+			if _, ok := stateDB.contractSnapshot[snapshot]; ok {
+				delete(stateDB.contractSnapshot, i)
+			} else {
+				break
+			}
+		}
+	}
 	// restore preimages
 	stateDB.preimages = nil
 	stateDB.preimages = stateDB.preimageSnapshot[snapshot]
+	if stateDB.fixSnapshotOrder {
+		delete(stateDB.preimageSnapshot, snapshot)
+		for i := snapshot + 1; ; i++ {
+			if _, ok := stateDB.preimageSnapshot[snapshot]; ok {
+				delete(stateDB.preimageSnapshot, i)
+			} else {
+				break
+			}
+		}
+	}
 }
 
 func (stateDB *StateDBAdapter) cachedContractAddrs() []hash.Hash160 {
@@ -484,7 +514,11 @@ func (stateDB *StateDBAdapter) Snapshot() int {
 	sn := stateDB.sm.Snapshot()
 	if _, ok := stateDB.suicideSnapshot[sn]; ok {
 		err := errors.New("unexpected error: duplicate snapshot version")
-		log.L().Error("Failed to snapshot.", zap.Error(err))
+		if stateDB.fixSnapshotOrder {
+			log.L().Panic("Failed to snapshot.", zap.Error(err))
+		} else {
+			log.L().Error("Failed to snapshot.", zap.Error(err))
+		}
 		// stateDB.err = err
 		return sn
 	}

diff --git a/action/protocol/execution/evm/evmstatedbadapter_test.go b/action/protocol/execution/evm/evmstatedbadapter_test.go
@@ -378,6 +378,17 @@ func TestSnapshotRevertAndCommit(t *testing.T) {
 			}
 		}
 
+		// snapshot after revert
+		require.Equal(1, stateDB.Snapshot())
+		if fixSnapshot {
+			require.Equal(1, len(stateDB.contractSnapshot))
+			require.Equal(1, len(stateDB.suicideSnapshot))
+			require.Equal(1, len(stateDB.preimageSnapshot))
+		} else {
+			require.Equal(3, len(stateDB.contractSnapshot))
+			require.Equal(3, len(stateDB.suicideSnapshot))
+			require.Equal(3, len(stateDB.preimageSnapshot))
+		}
 		// commit snapshot 0's state
 		require.NoError(stateDB.CommitContracts())
 		stateDB.clear()