Fix in account address hashing prevention

Author: miha-stopar

geth-utils/gethutil/mpt/trie/secure_trie.go

@@ -109,9 +109,6 @@ func (t *SecureTrie) Update(key, value []byte) {
 // If a node was not found in the database, a MissingNodeError is returned.
 func (t *SecureTrie) TryUpdate(key, value []byte) error {
 	hk := t.hashKey(key)
-	if oracle.PreventHashingInSecureTrie {
-		hk = append(hk, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
-	}
 	err := t.trie.TryUpdate(hk, value)
 	if err != nil {
 		return err
@@ -214,6 +211,9 @@ func (t *SecureTrie) hashKey(key []byte) []byte {
 		return t.hashKeyBuf[:]
 	} else {
 		// For generating special tests for MPT circuit.
+		if len(key) < 32 { // accounts
+			key = append(key, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		}
 		return key
 	}
 }

geth-utils/gethutil/mpt/witness/gen_witness_from_infura_blockchain_test.go

@@ -2347,6 +2347,7 @@ func TestStorageWrongExtensionNode(t *testing.T) {
 	oracle.PreventHashingInSecureTrie = true // to store the unchanged key
 
 	statedb.CreateAccount(addr)
+	statedb.IntermediateRoot(false)
 
 	key1 := common.HexToHash("0x1230000000000000000000000000000000000000000000000000000000000000")
 	key2 := common.HexToHash("0x1231000000000000000000000000000000000000000000000000000000000000")

geth-utils/gethutil/mpt/witness/prepare_witness.go

@@ -65,14 +65,15 @@ func GetWitness(nodeUrl string, blockNum int, trieModifications []TrieModificati
 	return obtainTwoProofsAndConvertToWitness(trieModifications, statedb, 0)
 }
 
-func obtainAccountProofAndConvertToWitness(i int, tMod TrieModification, tModsLen int, statedb *state.StateDB, specialTest byte) []Node {
+func obtainAccountProofAndConvertToWitness(tMod TrieModification, statedb *state.StateDB, specialTest byte) []Node {
 	statedb.IntermediateRoot(false)
 
 	addr := tMod.Address
 	addrh := crypto.Keccak256(addr.Bytes())
 	if oracle.PreventHashingInSecureTrie {
 		addrh = addr.Bytes()
-		addrh = append(addrh, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		addrh = append(addrh, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		addr = common.BytesToAddress(addrh)
 	}
 	accountAddr := trie.KeybytesToHex(addrh)
 
@@ -82,7 +83,7 @@ func obtainAccountProofAndConvertToWitness(i int, tMod TrieModification, tModsLe
 	// for cases when statedb.loadRemoteAccountsIntoStateObjects = false.
 	statedb.SetStateObjectIfExists(tMod.Address)
 
-	oracle.PrefetchAccount(statedb.Db.BlockNumber, tMod.Address, nil)
+	oracle.PrefetchAccount(statedb.Db.BlockNumber, addr, nil)
 	accountProof, aNeighbourNode1, aExtNibbles1, isLastLeaf1, aIsNeighbourNodeHashed1, err := statedb.GetProof(addr)
 	check(err)
 
@@ -182,10 +183,11 @@ func obtainTwoProofsAndConvertToWitness(trieModifications []TrieModification, st
 			if oracle.PreventHashingInSecureTrie {
 				addrh = addr.Bytes()
 				addrh = append(addrh, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+				addr = common.BytesToAddress(addrh)
 			}
 			accountAddr := trie.KeybytesToHex(addrh)
 
-			oracle.PrefetchAccount(statedb.Db.BlockNumber, tMod.Address, nil)
+			oracle.PrefetchAccount(statedb.Db.BlockNumber, addr, nil)
 			oracle.PrefetchStorage(statedb.Db.BlockNumber, addr, tMod.Key, nil)
 
 			if specialTest == 1 {
@@ -282,7 +284,7 @@ func obtainTwoProofsAndConvertToWitness(trieModifications []TrieModification, st
 			nodes = append(nodes, nodesStorage...)
 			nodes = append(nodes, GetEndNode())
 		} else {
-			accountNodes := obtainAccountProofAndConvertToWitness(i, tMod, len(trieModifications), statedb, specialTest)
+			accountNodes := obtainAccountProofAndConvertToWitness(tMod, statedb, specialTest)
 			nodes = append(nodes, accountNodes...)
 		}
 	}