feat(trace): add deletion proof into storage trace (#245)

* update zktrie, inducde proof with deletion

* Add deletion proof for stateDb

* fill deletionproofs

* lint

* proving more comment for parts prone to be confused

* simplify array appending

Co-authored-by: maskpp <maskpp266@gmail.com>

* bump zktrie version and fix issues

* add delproof into witgen

---------

Co-authored-by: maskpp <maskpp266@gmail.com>

core/state/statedb.go

@@ -351,32 +351,38 @@ func (s *StateDB) GetRootHash() common.Hash {
 }
 
 // StorageTrieProof is not in Db interface and used explictily for reading proof in storage trie (not the dirty value)
-func (s *StateDB) GetStorageTrieProof(a common.Address, key common.Hash) ([][]byte, error) {
+// For zktrie it also provide required data for predict the deletion, else it just fallback to GetStorageProof
+func (s *StateDB) GetStorageTrieProof(a common.Address, key common.Hash) ([][]byte, []byte, error) {
 
 	// try the trie in stateObject first, else we would create one
 	stateObject := s.getStateObject(a)
 	if stateObject == nil {
-		return nil, errors.New("storage trie for requested address does not exist")
+		return nil, nil, errors.New("storage trie for requested address does not exist")
 	}
 
-	trie := stateObject.trie
+	trieS := stateObject.trie
 	var err error
-	if trie == nil {
+	if trieS == nil {
 		// use a new, temporary trie
-		trie, err = s.db.OpenStorageTrie(stateObject.addrHash, stateObject.data.Root)
+		trieS, err = s.db.OpenStorageTrie(stateObject.addrHash, stateObject.data.Root)
 		if err != nil {
-			return nil, fmt.Errorf("can't create storage trie on root %s: %v ", stateObject.data.Root, err)
+			return nil, nil, fmt.Errorf("can't create storage trie on root %s: %v ", stateObject.data.Root, err)
 		}
 	}
 
 	var proof proofList
+	var sibling []byte
 	if s.IsZktrie() {
+		zkTrie := trieS.(*trie.ZkTrie)
+		if zkTrie == nil {
+			panic("unexpected trie type for zktrie")
+		}
 		key_s, _ := zkt.ToSecureKeyBytes(key.Bytes())
-		err = trie.Prove(key_s.Bytes(), 0, &proof)
+		sibling, err = zkTrie.ProveWithDeletion(key_s.Bytes(), 0, &proof)
 	} else {
-		err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
+		err = trieS.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
 	}
-	return proof, err
+	return proof, sibling, err
 }
 
 // GetStorageProof returns the Merkle proof for given storage slot.

core/types/l2trace.go

@@ -48,6 +48,10 @@ type StorageTrace struct {
 
 	// All storage proofs BEFORE execution
 	StorageProofs map[string]map[string][]hexutil.Bytes `json:"storageProofs,omitempty"`
+
+	// Node entries for deletion, no need to distinguish what it is from, just read them
+	// into the partial db
+	DeletionProofs []hexutil.Bytes `json:"deletionProofs,omitempty"`
 }
 
 // ExecutionResult groups all structured logs emitted by the EVM

eth/tracers/api_blocktrace.go

@@ -314,7 +314,7 @@ func (api *API) getTxResult(env *traceEnv, state *state.StateDB, index int, bloc
 			}
 			env.sMu.Unlock()
 
-			proof, err := state.GetStorageTrieProof(addr, key)
+			proof, sibling, err := state.GetStorageTrieProof(addr, key)
 			if err != nil {
 				log.Error("Storage proof not available", "error", err, "address", addrStr, "key", keyStr)
 				// but we still mark the proofs map with nil array
@@ -325,6 +325,9 @@ func (api *API) getTxResult(env *traceEnv, state *state.StateDB, index int, bloc
 			}
 			env.sMu.Lock()
 			m[keyStr] = wrappedProof
+			if sibling != nil {
+				env.DeletionProofs = append(env.DeletionProofs, sibling)
+			}
 			env.sMu.Unlock()
 		}
 	}

go.mod

@@ -49,7 +49,7 @@ require (
 	github.com/prometheus/tsdb v0.7.1
 	github.com/rjeczalik/notify v0.9.1
 	github.com/rs/cors v1.7.0
-	github.com/scroll-tech/zktrie v0.5.0
+	github.com/scroll-tech/zktrie v0.5.2
 	github.com/shirou/gopsutil v3.21.4-0.20210419000835-c7a38de76ee5+incompatible
 	github.com/status-im/keycard-go v0.0.0-20190316090335-8537d3370df4
 	github.com/stretchr/testify v1.7.0

go.sum

@@ -378,8 +378,8 @@ github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFR
 github.com/rs/cors v1.7.0 h1:+88SsELBHx5r+hZ8TCkggzSstaWNbDvThkVK8H6f9ik=
 github.com/rs/cors v1.7.0/go.mod h1:gFx+x8UowdsKA9AchylcLynDq+nNFfI8FkUZdN/jGCU=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
-github.com/scroll-tech/zktrie v0.5.0 h1:dABDR6lMZq6Hs+fWQSiHbX8s3AOX6hY+5nkhSYm5rmU=
-github.com/scroll-tech/zktrie v0.5.0/go.mod h1:XvNo7vAk8yxNyTjBDj5WIiFzYW4bx/gJ78+NK6Zn6Uk=
+github.com/scroll-tech/zktrie v0.5.2 h1:U34jPXMLGOlRHfdvYp5VVgOcC0RuPeJmcS3bWotCWiY=
+github.com/scroll-tech/zktrie v0.5.2/go.mod h1:XvNo7vAk8yxNyTjBDj5WIiFzYW4bx/gJ78+NK6Zn6Uk=
 github.com/segmentio/kafka-go v0.1.0/go.mod h1:X6itGqS9L4jDletMsxZ7Dz+JFWxM6JHfPOCvTvk+EJo=
 github.com/segmentio/kafka-go v0.2.0/go.mod h1:X6itGqS9L4jDletMsxZ7Dz+JFWxM6JHfPOCvTvk+EJo=
 github.com/sergi/go-diff v1.0.0/go.mod h1:0CfEIISq7TuYL3j771MWULgwwjU+GofnZX9QAmXWZgo=

trie/zk_trie.go

@@ -174,29 +174,49 @@ func (t *ZkTrie) NodeIterator(start []byte) NodeIterator {
 // nodes of the longest existing prefix of the key (at least the root node), ending
 // with the node that proves the absence of the key.
 func (t *ZkTrie) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) error {
-	err := t.ZkTrie.Prove(key, fromLevel, func(n *zktrie.Node) error {
-		nodeHash, err := n.NodeHash()
-		if err != nil {
-			return err
-		}
+	// omit sibling, which is not required for proving only
+	_, err := t.ProveWithDeletion(key, fromLevel, proofDb)
+	return err
+}
 
-		if n.Type == zktrie.NodeTypeLeaf {
-			preImage := t.GetKey(n.NodeKey.Bytes())
-			if len(preImage) > 0 {
-				n.KeyPreimage = &zkt.Byte32{}
-				copy(n.KeyPreimage[:], preImage)
-				//return fmt.Errorf("key preimage not found for [%x] ref %x", n.NodeKey.Bytes(), k.Bytes())
+// ProveWithDeletion is the implement of Prove, it also return possible sibling node
+// (if there is, i.e. the node of key exist and is not the only node in trie)
+// so witness generator can predict the final state root after deletion of this key
+// the returned sibling node has no key along with it for witness generator must decode
+// the node for its purpose
+func (t *ZkTrie) ProveWithDeletion(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) (sibling []byte, err error) {
+	err = t.ZkTrie.ProveWithDeletion(key, fromLevel,
+		func(n *zktrie.Node) error {
+			nodeHash, err := n.NodeHash()
+			if err != nil {
+				return err
 			}
-		}
-		return proofDb.Put(nodeHash[:], n.Value())
-	})
+
+			if n.Type == zktrie.NodeTypeLeaf {
+				preImage := t.GetKey(n.NodeKey.Bytes())
+				if len(preImage) > 0 {
+					n.KeyPreimage = &zkt.Byte32{}
+					copy(n.KeyPreimage[:], preImage)
+					//return fmt.Errorf("key preimage not found for [%x] ref %x", n.NodeKey.Bytes(), k.Bytes())
+				}
+			}
+			return proofDb.Put(nodeHash[:], n.Value())
+		},
+		func(_ *zktrie.Node, n *zktrie.Node) {
+			// the sibling for each leaf should be unique except for EmptyNode
+			if n != nil && n.Type != zktrie.NodeTypeEmpty {
+				sibling = n.Value()
+			}
+		},
+	)
 	if err != nil {
-		return err
+		return
 	}
 
 	// we put this special kv pair in db so we can distinguish the type and
 	// make suitable Proof
-	return proofDb.Put(magicHash, zktrie.ProofMagicBytes())
+	err = proofDb.Put(magicHash, zktrie.ProofMagicBytes())
+	return
 }
 
 // VerifyProof checks merkle proofs. The given proof must contain the value for

trie/zk_trie_proof_test.go

@@ -143,8 +143,8 @@ func TestSMTMissingKeyProof(t *testing.T) {
 	tr, _ := NewZkTrie(common.Hash{}, NewZktrieDatabase((memorydb.New())))
 	mt := &zkTrieImplTestWrapper{tr.Tree()}
 	err := mt.UpdateWord(
-		zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("k"), 20)),
-		zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("v"), 20)),
+		zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("k"), 32)),
+		zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("v"), 32)),
 	)
 	assert.Nil(t, err)
 
@@ -195,3 +195,42 @@ func randomZktrie(t *testing.T, n int) (*ZkTrie, map[string]*kv) {
 
 	return tr, vals
 }
+
+// Tests that new "proof with deletion" feature
+func TestProofWithDeletion(t *testing.T) {
+	tr, _ := NewZkTrie(common.Hash{}, NewZktrieDatabase((memorydb.New())))
+	mt := &zkTrieImplTestWrapper{tr.Tree()}
+	key1 := bytes.Repeat([]byte("k"), 32)
+	key2 := bytes.Repeat([]byte("m"), 32)
+	err := mt.UpdateWord(
+		zkt.NewByte32FromBytesPaddingZero(key1),
+		zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("v"), 32)),
+	)
+	assert.NoError(t, err)
+	err = mt.UpdateWord(
+		zkt.NewByte32FromBytesPaddingZero(key2),
+		zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("n"), 32)),
+	)
+	assert.NoError(t, err)
+
+	proof := memorydb.New()
+	s_key1, err := zkt.ToSecureKeyBytes(key1)
+	assert.NoError(t, err)
+
+	sibling1, err := tr.ProveWithDeletion(s_key1.Bytes(), 0, proof)
+	assert.NoError(t, err)
+	nd, err := tr.TryGet(key2)
+	assert.NoError(t, err)
+	l := len(sibling1)
+	// a hacking to grep the value part directly from the encoded leaf node,
+	// notice the sibling of key `k*32`` is just the leaf of key `m*32`
+	assert.Equal(t, sibling1[l-33:l-1], nd)
+
+	s_key2, err := zkt.ToSecureKeyBytes(bytes.Repeat([]byte("x"), 32))
+	assert.NoError(t, err)
+
+	sibling2, err := tr.ProveWithDeletion(s_key2.Bytes(), 0, proof)
+	assert.NoError(t, err)
+	assert.Nil(t, sibling2)
+
+}

trie/zkproof/writer.go

@@ -201,7 +201,7 @@ func NewZkTrieProofWriter(storage *types.StorageTrace) (*zktrieProofWriter, erro
 				var err error
 				storages[addr], err = trie.NewZkTrie(accState.Root, zkDb)
 				if err != nil {
-					return nil, fmt.Errorf("zktrie create failure for storage in addr <%s>: %s", err, addrs)
+					return nil, fmt.Errorf("zktrie create failure for storage in addr <%s>: %s, (root %s)", addrs, err, accState.Root)
 				}
 
 			} else {
@@ -211,6 +211,23 @@ func NewZkTrieProofWriter(storage *types.StorageTrace) (*zktrieProofWriter, erro
 		}
 	}
 
+	for _, delProof := range storage.DeletionProofs {
+
+		n, err := zktrie.DecodeSMTProof(delProof)
+		if err != nil {
+			log.Warn("decode delproof string fail", "error", err, "node", delProof)
+		} else if n != nil {
+			hash, err := n.NodeHash()
+			if err != nil {
+				log.Warn("node has no valid node hash", "error", err)
+			} else {
+				//notice: must consistent with trie/merkletree.go
+				bt := hash[:]
+				underlayerDb.Put(bt, delProof)
+			}
+		}
+	}
+
 	zktrie, err := trie.NewZkTrie(
 		storage.RootBefore,
 		trie.NewZktrieDatabase(underlayerDb),
@@ -381,7 +398,7 @@ func (w *zktrieProofWriter) traceAccountUpdate(addr common.Address, updateAccDat
 	var proof proofList
 	s_key, _ := zkt.ToSecureKeyBytes(addr.Bytes())
 	if err := w.tracingZktrie.Prove(s_key.Bytes(), 0, &proof); err != nil {
-		return nil, fmt.Errorf("prove BEFORE state for <%x> fail: %s", addr.Bytes(), err)
+		return nil, fmt.Errorf("prove BEFORE state fail: %s", err)
 	}
 
 	decodeProofForMPTPath(proof, out.AccountPath[0])