From bd82e59599bafd76d031dc3f1867bce50177ae12 Mon Sep 17 00:00:00 2001
From: miha-stopar <stopar.miha@gmail.com>
Date: Thu, 22 Feb 2024 11:51:58 +0100
Subject: [PATCH] Preparing witness for wrong extension node

---
 geth-utils/gethutil/mpt/witness/branch.go     | 12 ++++--
 .../gethutil/mpt/witness/prepare_witness.go   | 37 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/geth-utils/gethutil/mpt/witness/branch.go b/geth-utils/gethutil/mpt/witness/branch.go
index dec0854efa..eb59b86fe8 100644
--- a/geth-utils/gethutil/mpt/witness/branch.go
+++ b/geth-utils/gethutil/mpt/witness/branch.go
@@ -149,9 +149,8 @@ func prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []b
 	return node
 }
 
-// getDriftedPosition returns the position in branch to which the leaf drifted because another
-// leaf has been added to the same slot. This information is stored into a branch init row.
-func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
+// getNibbles returns the nibbles of the leaf or extension node.
+func getNibbles(leafKeyRow []byte) []byte {
 	var nibbles []byte
 	if leafKeyRow[0] != 248 {
 		var keyLen int
@@ -193,6 +192,13 @@ func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
 		}
 	}
 
+	return nibbles
+}
+
+// getDriftedPosition returns the position in branch to which the leaf drifted because another
+// leaf has been added to the same slot. This information is stored into a branch init row.
+func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
+	nibbles := getNibbles(leafKeyRow)	
 	return nibbles[numberOfNibbles]
 }
 
diff --git a/geth-utils/gethutil/mpt/witness/prepare_witness.go b/geth-utils/gethutil/mpt/witness/prepare_witness.go
index 9335136485..c5c4705f06 100644
--- a/geth-utils/gethutil/mpt/witness/prepare_witness.go
+++ b/geth-utils/gethutil/mpt/witness/prepare_witness.go
@@ -520,13 +520,46 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 				// to obtain the underlying branch to be able to finally add (besides this branch)
 				// the placeholder leaf. So we need to query getProof again with one of the leaves that is
 				// actually in this extension node.
-				// TODO:
 				if isAccountProof {
+					// TODO
 					node := prepareAccountLeafPlaceholderNode(addr, addrh, key, keyIndex)
 					nodes = append(nodes, node)
 				} else {
-					node := prepareStorageLeafPlaceholderNode(storage_key, key, keyIndex)
+					nibbles := getNibbles(proof2[len(proof2)-1])
+					newKey := make([]byte, len(key))
+					copy(newKey, key)
+
+					for i := 0; i < len(nibbles); i++ {
+						n := nibbles[i]
+						if key[i] != n {
+							newKey[i] = n
+						}
+					}
+
+					// The last nibble should be the one that gets one of the leaves
+					// in the branch (not nil):
+					var proof [][]byte
+					var err error
+					for i := 0; i < 16; i++ {
+						newKey[keyIndex] = byte(i)
+						k := trie.HexToKeybytes(newKey)
+						ky := common.BytesToHash(k)
+						proof, _, _, _, _, err = statedb.GetStorageProof(addr, ky)
+						check(err)
+						if !isBranch(proof[len(proof)-1]) {
+							break
+						}
+					}
+
+					branchRlp := proof[len(proof)-2] // the last element has to be a leaf
+					isExtension := true
+
+					extNode := proof2[len(proof2)-1]
+					bNode := prepareBranchNode(branchRlp, branchRlp, extNode, extNode, extListRlpBytes, extValues,
+						key[keyIndex], key[keyIndex], false, false, isExtension)
+					nodes = append(nodes, bNode)
 
+					node := prepareStorageLeafNode(proof[len(proof)-1], proof[len(proof)-1], nil, storage_key, key, nonExistingStorageProof, false, false, false, false)
 					nodes = append(nodes, node)
 				}
 			}