From 239887e8b849018cdeea9043affdf50873e83753 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Wed, 18 Mar 2020 09:40:01 +0100
Subject: [PATCH] tests: validate snapshot after test

---
 core/blockchain.go       |  9 +++++++++
 tests/block_test_util.go | 14 ++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/core/blockchain.go b/core/blockchain.go
index de0d4f3993..21326de2e5 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -520,6 +520,15 @@ func (bc *BlockChain) CurrentBlock() *types.Block {
 	return bc.currentBlock.Load().(*types.Block)
 }
 
+// Snapshot returns the blockchain snapshot tree. This method is mainly used for
+// testing, to make it possible to verify the snapshot after execution.
+//
+// Warning: There are no guarantees about the safety of using the returned 'snap' if the
+// blockchain is simultaneously importing blocks, so take care.
+func (bc *BlockChain) Snapshot() *snapshot.Tree {
+	return bc.snaps
+}
+
 // CurrentFastBlock retrieves the current fast-sync head block of the canonical
 // chain. The block is retrieved from the blockchain's internal cache.
 func (bc *BlockChain) CurrentFastBlock() *types.Block {
diff --git a/tests/block_test_util.go b/tests/block_test_util.go
index 1ae986e3ca..37f63f538a 100644
--- a/tests/block_test_util.go
+++ b/tests/block_test_util.go
@@ -32,6 +32,7 @@ import (
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/params"
@@ -144,6 +145,19 @@ func (t *BlockTest) Run(snapshotter bool) error {
 	if err = t.validatePostState(newDB); err != nil {
 		return fmt.Errorf("post state validation failed: %v", err)
 	}
+	// Cross-check the snapshot-to-hash against the trie hash
+	if snapshotter {
+		snapTree := chain.Snapshot()
+		root := chain.CurrentBlock().Root()
+		it, err := snapTree.AccountIterator(root, common.Hash{})
+		if err != nil {
+			return fmt.Errorf("Could not create iterator for root %x: %v", root, err)
+		}
+		generatedRoot := snapshot.GenerateTrieRoot(it)
+		if generatedRoot != root {
+			return fmt.Errorf("Snapshot corruption, got %d exp %d", generatedRoot, root)
+		}
+	}
 	return t.validateImportedHeaders(chain, validBlocks)
 }