fix: use tracer when generating storage proofs (#1093)

* fix: possible deadlock in proof tracer * fix: use correct tracer instance * fix(ProofTracer): add magicSMTBytes (#1094) add magic ending --------- Co-authored-by: colin <[email protected]>
scroll-tech · Nov 26, 2024 · f8e7694 · f8e7694
1 parent 56864e6
commit f8e7694
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git a/rollup/tracing/proof.go b/rollup/tracing/proof.go
@@ -8,6 +8,11 @@ import (
 	"github.com/scroll-tech/go-ethereum/trie"
 )
 
+var (
+	magicHash     = []byte("THIS IS THE MAGIC INDEX FOR ZKTRIE")
+	magicSMTBytes = []byte("THIS IS SOME MAGIC BYTES FOR SMT m1rRXgP2xpDI")
+)
+
 type ProofTracer struct {
 	trie           *trie.ZkTrie
 	deletionTracer map[trie.Hash]struct{}
@@ -124,7 +129,7 @@ func (t *ProofTracer) MarkDeletion(key []byte) error {
 func (t *ProofTracer) Prove(key []byte, proofDb ethdb.KeyValueWriter) error {
 	fromLevel := uint(0)
 	var mptPath []*trie.Node
-	return t.trie.ProveWithDeletion(key, fromLevel,
+	if err := t.trie.ProveWithDeletion(key, fromLevel,
 		func(n *trie.Node) error {
 			nodeHash, err := n.NodeHash()
 			if err != nil {
@@ -133,11 +138,6 @@ func (t *ProofTracer) Prove(key []byte, proofDb ethdb.KeyValueWriter) error {
 
 			switch n.Type {
 			case trie.NodeTypeLeaf_New:
-				preImage := t.trie.GetKey(n.NodeKey.Bytes())
-				if len(preImage) > 0 {
-					n.KeyPreimage = &trie.Byte32{}
-					copy(n.KeyPreimage[:], preImage)
-				}
 			case trie.NodeTypeBranch_0, trie.NodeTypeBranch_1,
 				trie.NodeTypeBranch_2, trie.NodeTypeBranch_3:
 				mptPath = append(mptPath, n)
@@ -158,5 +158,11 @@ func (t *ProofTracer) Prove(key []byte, proofDb ethdb.KeyValueWriter) error {
 			mptPath = append(mptPath, n)
 			t.rawPaths[string(key)] = mptPath
 		},
-	)
+	); err != nil {
+		return err
+	}
+
+	// we put this special kv pair in db so we can distinguish the type and
+	// make suitable Proof
+	return proofDb.Put(magicHash, magicSMTBytes)
 }
diff --git a/rollup/tracing/tracing.go b/rollup/tracing/tracing.go
@@ -477,7 +477,7 @@ func (env *TraceEnv) getTxResult(state *state.StateDB, index int, block *types.B
 			env.sMu.Unlock()
 
 			var proof zkproof.ProofList
-			if err = env.ZkTrieTracer[addrStr].Prove(key.Bytes(), &proof); err != nil {
+			if err = zktrieTracer.Prove(key.Bytes(), &proof); err != nil {
 				log.Error("Storage proof not available", "error", err, "address", addrStr, "key", keyStr)
 				// but we still mark the proofs map with nil array
 			}

diff --git a/trie/zk_trie.go b/trie/zk_trie.go
@@ -1198,6 +1198,11 @@ func (mt *ZkTrie) prove(kHash *Hash, fromLevel uint, writeNode func(*Node) error
 		case NodeTypeLeaf_New:
 			// notice even we found a leaf whose entry didn't match the expected k,
 			// we still include it as the proof of absence
+			preImage := mt.getKey(n.NodeKey.Bytes())
+			if len(preImage) > 0 {
+				n.KeyPreimage = &Byte32{}
+				copy(n.KeyPreimage[:], preImage)
+			}
 		case NodeTypeBranch_0, NodeTypeBranch_1, NodeTypeBranch_2, NodeTypeBranch_3:
 			finished = false
 			if path[i] {