From 1f9bca2e92d3fe86f4431484a2ec437df597af3e Mon Sep 17 00:00:00 2001
From: Gav Wood <i@gavwood.com>
Date: Thu, 22 May 2014 20:45:43 +0200
Subject: [PATCH] Fix and paranoia for trie.

---
 alethzero/MainWin.cpp     | 17 ++++++++++++-----
 libethcore/FileSystem.cpp |  1 +
 libethereum/State.cpp     | 33 +++++++++++++++++++++++++++++++--
 libethereum/State.h       |  2 ++
 4 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/alethzero/MainWin.cpp b/alethzero/MainWin.cpp
index 96ef45e2d..75de6f76f 100644
--- a/alethzero/MainWin.cpp
+++ b/alethzero/MainWin.cpp
@@ -767,11 +767,18 @@ void Main::on_contracts_currentItemChanged()
 		auto h = h160((byte const*)hba.data(), h160::ConstructFromPointer);
 
 		stringstream s;
-		auto storage = state().storage(h);
-		for (auto const& i: storage)
-			s << "@" << showbase << hex << i.first << "&nbsp;&nbsp;&nbsp;&nbsp;" << showbase << hex << i.second << "<br/>";
-		s << "<h4>Body Code</h4>" << disassemble(state().code(h));
-		ui->contractInfo->appendHtml(QString::fromStdString(s.str()));
+		try
+		{
+			auto storage = state().storage(h);
+			for (auto const& i: storage)
+				s << "@" << showbase << hex << i.first << "&nbsp;&nbsp;&nbsp;&nbsp;" << showbase << hex << i.second << "<br/>";
+			s << "<h4>Body Code</h4>" << disassemble(state().code(h));
+			ui->contractInfo->appendHtml(QString::fromStdString(s.str()));
+		}
+		catch (eth::InvalidTrie)
+		{
+			ui->contractInfo->appendHtml("Corrupted trie.");
+		}
 	}
 }
 
diff --git a/libethcore/FileSystem.cpp b/libethcore/FileSystem.cpp
index 559a58997..8c55937ed 100644
--- a/libethcore/FileSystem.cpp
+++ b/libethcore/FileSystem.cpp
@@ -23,6 +23,7 @@
 
 #include "FileSystem.h"
 #include "Common.h"
+#include "Log.h"
 
 #ifdef _WIN32
 #include <shlobj.h>
diff --git a/libethereum/State.cpp b/libethereum/State.cpp
index ac605644f..13e9da421 100644
--- a/libethereum/State.cpp
+++ b/libethereum/State.cpp
@@ -848,13 +848,37 @@ bytes const& State::code(Address _contract) const
 	return m_cache[_contract].code();
 }
 
+bool State::isTrieGood()
+{
+	for (auto const& i: m_state)
+	{
+		RLP r(i.second);
+		TrieDB<h256, Overlay> storageDB(&m_db, r[2].toHash<h256>());
+		try
+		{
+			for (auto const& j: storageDB) {}
+		}
+		catch (InvalidTrie)
+		{
+			return false;
+		}
+		if (r[3].toHash<h256>() != EmptySHA3 &&  m_db.lookup(r[3].toHash<h256>()).empty())
+			return false;
+	}
+	return true;
+}
+
 u256 State::execute(bytesConstRef _rlp)
 {
 #ifndef RELEASE
 	commit();	// get an updated hash
 #endif
 
-	// TODO: CHECK TRIE
+	if (!isTrieGood())
+	{
+		cwarn << "BAD TRIE before execution begins.";
+		throw InvalidTrie();
+	}
 
 	State old(*this);
 	auto h = rootHash();
@@ -880,7 +904,12 @@ u256 State::execute(bytesConstRef _rlp)
 	cnote << "Executed; now" << rootHash();
 	cnote << old.diff(*this);
 
-	// TODO: CHECK TRIE
+	if (!isTrieGood())
+	{
+		cwarn << "BAD TRIE immediately after execution.";
+		throw InvalidTrie();
+	}
+
 	// TODO: CHECK TRIE after level DB flush to make sure exactly the same.
 
 	// Add to the user-originated transactions that we've executed.
diff --git a/libethereum/State.h b/libethereum/State.h
index 59e2c1a74..13cfb146e 100644
--- a/libethereum/State.h
+++ b/libethereum/State.h
@@ -299,6 +299,8 @@ private:
 	/// @returns gas used by transactions thus far executed.
 	u256 gasUsed() const { return m_transactions.size() ? m_transactions.back().gasUsed : 0; }
 
+	bool isTrieGood();
+
 	Overlay m_db;								///< Our overlay for the state tree.
 	TrieDB<Address, Overlay> m_state;			///< Our state tree, as an Overlay DB.
 	std::vector<TransactionReceipt> m_transactions;	///< The current list of transactions that we've included in the state.