Fixing handling the accounts without hashing in the trie

Author: miha-stopar

geth-utils/gethutil/mpt/state/statedb.go

@@ -298,7 +298,12 @@ func (s *StateDB) GetState(addr common.Address, hash common.Hash) common.Hash {
 
 // GetProof returns the Merkle proof for a given account.
 func (s *StateDB) GetProof(addr common.Address) ([][]byte, []byte, [][]byte, bool, bool, error) {
-	return s.GetProofByHash(crypto.Keccak256Hash(addr.Bytes()))
+	var newAddr common.Hash
+	if oracle.PreventHashingInSecureTrie {
+		bytes := append(addr.Bytes(), []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		newAddr = common.BytesToHash(bytes)
+	}
+	return s.GetProofByHash(newAddr)
 }
 
 // GetProofByHash returns the Merkle proof for a given account.
@@ -311,7 +316,13 @@ func (s *StateDB) GetProofByHash(addrHash common.Hash) ([][]byte, []byte, [][]by
 // GetStorageProof returns the Merkle proof for given storage slot.
 func (s *StateDB) GetStorageProof(a common.Address, key common.Hash) ([][]byte, []byte, [][]byte, bool, bool, error) {
 	var proof proofList
-	trie := s.StorageTrie(a)
+	newAddr := a
+	if oracle.PreventHashingInSecureTrie {
+		bytes := append(a.Bytes(), []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		newAddr = common.BytesToAddress(bytes)
+	}
+
+	trie := s.StorageTrie(newAddr)
 	if trie == nil {
 		return proof, nil, nil, false, false, errors.New("storage trie for requested address does not exist")
 	}

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

@@ -109,6 +109,9 @@ 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

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

@@ -2344,11 +2344,10 @@ func TestStorageWrongExtensionNode(t *testing.T) {
 	addr := common.HexToAddress("0x40efbf12580138bc623c95757286df4e24eb81c9")
 
 	statedb.DisableLoadingRemoteAccounts()
+	oracle.PreventHashingInSecureTrie = true // to store the unchanged key
 
 	statedb.CreateAccount(addr)
 
-	oracle.PreventHashingInSecureTrie = true // to store the unchanged key
-
 	key1 := common.HexToHash("0x1230000000000000000000000000000000000000000000000000000000000000")
 	key2 := common.HexToHash("0x1231000000000000000000000000000000000000000000000000000000000000")
 

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

@@ -179,6 +179,10 @@ func obtainTwoProofsAndConvertToWitness(trieModifications []TrieModification, st
 
 			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}...)
+			}
 			accountAddr := trie.KeybytesToHex(addrh)
 
 			oracle.PrefetchAccount(statedb.Db.BlockNumber, tMod.Address, nil)