From e2f3e87842d005ef141a87454743c0cf7355fd39 Mon Sep 17 00:00:00 2001
From: antonis19 <antonis19@users.noreply.github.com>
Date: Mon, 13 Jan 2025 17:43:31 +0100
Subject: [PATCH] add state consistency check during witness generation

---
 erigon-lib/commitment/hex_patricia_hashed.go |  4 ++
 turbo/jsonrpc/eth_call.go                    | 62 ++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/erigon-lib/commitment/hex_patricia_hashed.go b/erigon-lib/commitment/hex_patricia_hashed.go
index 3d692dca755..03b969c65b3 100644
--- a/erigon-lib/commitment/hex_patricia_hashed.go
+++ b/erigon-lib/commitment/hex_patricia_hashed.go
@@ -1131,6 +1131,10 @@ func (c *cell) String() string {
 	return s
 }
 
+func (hph *HexPatriciaHashed) GetPatriciaContext() PatriciaContext {
+	return hph.ctx
+}
+
 func (hph *HexPatriciaHashed) PrintGrid() {
 	fmt.Printf("GRID:\n")
 	for row := 0; row < hph.activeRows; row++ {
diff --git a/turbo/jsonrpc/eth_call.go b/turbo/jsonrpc/eth_call.go
index 1b8f9a7f46d..0c4f6bc9718 100644
--- a/turbo/jsonrpc/eth_call.go
+++ b/turbo/jsonrpc/eth_call.go
@@ -464,6 +464,55 @@ func verifyExecResult(execResult *core.EphemeralExecResult, block *types.Block)
 	return nil
 }
 
+// verify the state consistency between the history state reader and the hph.Ctx reader
+func verifyStateConsistency(touchedPlainKeys [][]byte, stateReader state.StateReader, hph *commitment.HexPatriciaHashed) (bool, error) {
+	ctx := hph.GetPatriciaContext()
+	for _, key := range touchedPlainKeys {
+		if len(key) == 20 { //account address
+			address := libcommon.Address(key)
+			accountInStateReader, err := stateReader.ReadAccountData(address)
+			if err != nil {
+				return false, err
+			}
+			accountDataInTrie, err := ctx.Account(key)
+			if err != nil {
+				return false, err
+			}
+			accountInTrie := accounts.Account{
+				Nonce:           accountDataInTrie.Nonce,
+				Balance:         accountDataInTrie.Balance,
+				Root:            accountDataInTrie.Storage,
+				CodeHash:        accountDataInTrie.CodeHash,
+				Incarnation:     accountInStateReader.Incarnation, // these 2 are copied over since they don't affect the RLP encoding
+				PrevIncarnation: accountInStateReader.PrevIncarnation,
+			}
+			if !accountInStateReader.Equals(&accountInTrie) {
+				err := fmt.Errorf("accountInStateReader(%+v) != accountInPatriciaContext(%+v)", accountInStateReader, accountInTrie)
+				return false, err
+			}
+
+		} else {
+			address := libcommon.Address(key[:20])
+			storageKey := libcommon.Hash(key[20:])
+			storageValInStateReader, err := stateReader.ReadAccountStorage(address, 0, &storageKey)
+			if err != nil {
+				return false, err
+			}
+
+			storageDataInTrie, err := ctx.Storage(storageKey[:])
+			if err != nil {
+				return false, err
+			}
+			storageValInPatriciaContext := storageDataInTrie.Storage[:]
+			if !bytes.Equal(storageValInStateReader, storageDataInTrie.Storage[:]) {
+				return false, fmt.Errorf("storageValInStateReader(%x) != storageValInPatriciaContext(%x)", storageValInStateReader, storageValInPatriciaContext)
+			}
+
+		}
+	}
+	return true, nil
+}
+
 func (api *BaseAPI) getWitness(ctx context.Context, db kv.RoDB, blockNrOrHash rpc.BlockNumberOrHash, txIndex hexutil.Uint, fullBlock bool, maxGetProofRewindBlockCount int, logger log.Logger) (hexutility.Bytes, error) {
 	roTx, err := db.BeginRo(ctx)
 	if err != nil {
@@ -574,6 +623,19 @@ func (api *BaseAPI) getWitness(ctx context.Context, db kv.RoDB, blockNrOrHash rp
 	touchedPlainKeys, touchedHashedKeys := store.Tds.GetTouchedPlainKeys()
 	codeReads := store.Tds.BuildCodeTouches()
 
+	// Verify that consistency of account and storage data for touched keys between the state reader
+	// and the patricia context (hph.ctx)
+	// If there is any inconsistency there will inevitably be a root hash mismatch later on when
+	// we try to construct the witness trie
+	isStateConsistent, err := verifyStateConsistency(touchedPlainKeys, store.Tds.StateReader, hph)
+	if err != nil {
+		return nil, err
+	}
+
+	if !isStateConsistent {
+		return nil, fmt.Errorf("state is inconsistent: %w", err)
+	}
+
 	// define these keys as "updates", but we are not really updating anything, we just want to load them into the grid,
 	// so this is just to satisfy the current hex patricia trie api.
 	updates := commitment.NewUpdates(commitment.ModeDirect, sdCtx.TempDir(), hph.HashAndNibblizeKey)