From 59ac6a482de8605303b0e3f7a959404491f3a5b7 Mon Sep 17 00:00:00 2001
From: KimiWu <leonhartwu@gmail.com>
Date: Mon, 8 Apr 2024 15:46:33 +0800
Subject: [PATCH] fix isExt()

---
 geth-utils/gethutil/mpt/trie/stacktrie.go          | 13 ++++++++++---
 geth-utils/gethutil/mpt/witness/branch.go          |  4 +++-
 geth-utils/gethutil/mpt/witness/prepare_witness.go |  2 +-
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/geth-utils/gethutil/mpt/trie/stacktrie.go b/geth-utils/gethutil/mpt/trie/stacktrie.go
index 108e0cdad9..e78172004c 100644
--- a/geth-utils/gethutil/mpt/trie/stacktrie.go
+++ b/geth-utils/gethutil/mpt/trie/stacktrie.go
@@ -651,15 +651,21 @@ func isTxLeaf(proofEl []byte) bool {
 	c, _ := rlp.CountValues(elems)
 
 	// 9: for tx (Nonce, Gas, GasPrice, Value, To, Data, r, s, v)
-	return c == 9 || c == 2
+	return (c == 9 || c == 2) && !isTxExt(proofEl)
+}
+
+func isTxExt(proofEl []byte) bool {
+	elems, _, _ := rlp.SplitList(proofEl)
+	idx := proofEl[0] - 225
+	return len(proofEl) < 50 && proofEl[0] < 248 && elems[idx] == 160
 }
 
 func printProof(ps [][]byte, idx []byte) {
 
-	enable := byte(150)
+	enable := byte(200)
 	fmt.Print(" [")
 	for _, p := range ps {
-		if p[0] == 226 && p[1]%16 == 0 && p[2] == 160 {
+		if isTxExt(p) {
 			fmt.Print("EXT - ")
 			if idx[0] > enable {
 				fmt.Print(" (", p, ") - ")
@@ -784,6 +790,7 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, []
 	isHashed := false
 
 	for i := 0; i < len(k); i++ {
+		// fmt.Print(k[i], "- ")
 		if c.nodeType == extNode {
 			nodes = append(nodes, c)
 			c = c.children[0]
diff --git a/geth-utils/gethutil/mpt/witness/branch.go b/geth-utils/gethutil/mpt/witness/branch.go
index f7ff15f9ff..c9caff2975 100644
--- a/geth-utils/gethutil/mpt/witness/branch.go
+++ b/geth-utils/gethutil/mpt/witness/branch.go
@@ -38,7 +38,9 @@ func isTxLeaf(proofEl []byte) bool {
 }
 
 func isTxExt(proofEl []byte) bool {
-	return proofEl[0] == 226 && proofEl[1]%16 == 0 && proofEl[2] == 160
+	elems, _, _ := rlp.SplitList(proofEl)
+	idx := proofEl[0] - 225
+	return len(proofEl) < 50 && proofEl[0] < 248 && elems[idx] == 160
 }
 
 // prepareBranchWitness takes the rows that are to be filled with branch data and it takes
diff --git a/geth-utils/gethutil/mpt/witness/prepare_witness.go b/geth-utils/gethutil/mpt/witness/prepare_witness.go
index ed4366d7a7..b1d72bf33e 100644
--- a/geth-utils/gethutil/mpt/witness/prepare_witness.go
+++ b/geth-utils/gethutil/mpt/witness/prepare_witness.go
@@ -346,13 +346,13 @@ func prepareWitnessSpecial(testName string, trieModifications []TrieModification
 //	-[] [EXT] -> [BRANCH - BRANCH - BRANCH - LEAF] --> 144
 //	-[] [BRANCH - LEAF] -> [BRANCH - BRANCH - LEAF] (modified extension)
 //	-[] [BRANCH - BRANCH - (...BRANCH)] -> [BRANCH - BRANCH - (...BRANCH) - LEAF]
-//	-[] [BRANCH - BRANCH - BRANCH - EXT] -> [BRANCH - BRANCH - BRANCH - BRANCH - LEAF]
 //	-[] [BRANCH - BRANCH - LEAF ] -> [BRANCH - BRANCH - LEAF - BRANCH - LEAF]
 //	-[] [LEAF] -> [LEAF]
 //	-[] [EXT] -> [EXT]
 //	-[] [EXT - EXT] -> [EXT - EXT]
 //	-[] [EXT - LEAF] -> [LEAF]
 //	-[] [BRANCH - LEAF] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF]
+//	-[] [BRANCH - BRANCH - EXT] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF]
 //	-[] [BRANCH - BRANCH - EXT - BRANCH - (LEAF)] -> [BRANCH - BRANCH - EXT - BRANCH - EXT - BRANCH - LEAF]
 //	-[] [BRANCH - BRANCH - EXT - BRANCH - (...BRANCH)] -> [BRANCH - BRANCH - EXT - BRANCH - (...BRANCH) - LEAF]
 //	-[] [LEAF] -> [BRANCH - BRANCH - EXT - BRANCH - BRANCH - LEAF]