proof: verify keys, pre-state and post-state

Signed-off-by: Ignacio Hagopian <[email protected]>

core/chain_makers.go

@@ -17,8 +17,10 @@
 package core
 
 import (
+	"bytes"
 	"fmt"
 	"math/big"
+	"sort"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/consensus"
@@ -358,32 +360,35 @@ func GenerateChainWithGenesis(genesis *Genesis, engine consensus.Engine, n int,
 		panic(err)
 	}
 	if genesis.Config != nil && genesis.Config.IsPrague(genesis.ToBlock().Number(), genesis.ToBlock().Time()) {
-		blocks, receipts, _, _ := GenerateVerkleChain(genesis.Config, genesis.ToBlock(), engine, db, n, gen)
+		blocks, receipts, _, _, _, _, _ := GenerateVerkleChain(genesis.Config, genesis.ToBlock(), engine, db, n, gen)
 		return db, blocks, receipts
 	}
 	blocks, receipts := GenerateChain(genesis.Config, genesis.ToBlock(), engine, db, n, gen)
 	return db, blocks, receipts
 }
 
-func GenerateVerkleChain(config *params.ChainConfig, parent *types.Block, engine consensus.Engine, db ethdb.Database, n int, gen func(int, *BlockGen)) ([]*types.Block, []types.Receipts, []*verkle.VerkleProof, []verkle.StateDiff) {
+func GenerateVerkleChain(config *params.ChainConfig, parent *types.Block, engine consensus.Engine, db ethdb.Database, n int, gen func(int, *BlockGen)) ([]*types.Block, []types.Receipts, []*verkle.VerkleProof, []verkle.StateDiff, [][][]byte, [][][]byte, [][][]byte) {
 	if config == nil {
 		config = params.TestChainConfig
 	}
 	proofs := make([]*verkle.VerkleProof, 0, n)
 	keyvals := make([]verkle.StateDiff, 0, n)
 	blocks, receipts := make(types.Blocks, n), make([]types.Receipts, n)
+	computedKeys := make([][][]byte, n)
+	computedPreStateValues := make([][][]byte, n)
+	computedPostStateValues := make([][][]byte, n)
 	chainreader := &generatedLinearChainReader{
 		config: config,
 		// GenerateVerkleChain should only be called with the genesis block
 		// as parent.
 		genesis: parent,
 		chain:   blocks,
 	}
-	genblock := func(i int, parent *types.Block, statedb *state.StateDB) (*types.Block, types.Receipts) {
+	genblock := func(i int, parent *types.Block, statedb *state.StateDB) (*types.Block, types.Receipts, [][]byte, [][]byte, [][]byte) {
 		b := &BlockGen{i: i, chain: blocks, parent: parent, statedb: statedb, config: config, engine: engine}
 		b.header = makeHeader(chainreader, parent, statedb, b.engine)
-		preState := statedb.Copy()
-		fmt.Println("prestate", preState.GetTrie().(*trie.VerkleTrie).ToDot())
+		preStateTree := statedb.Copy().GetTrie().(*trie.VerkleTrie)
+		fmt.Println("prestate", preStateTree.ToDot())
 
 		// Mutate the state and block according to any hard-fork specs
 		if daoBlock := config.DAOForkBlock; daoBlock != nil {
@@ -408,6 +413,19 @@ func GenerateVerkleChain(config *params.ChainConfig, parent *types.Block, engine
 				panic(err)
 			}
 
+			// Get the keys collected in the witness in the witness.
+			computedKeys := statedb.Witness().Keys()
+			sort.Slice(computedKeys, func(i, j int) bool { return bytes.Compare(computedKeys[i], computedKeys[j]) < 0 })
+
+			// Get the pre-state values from the database.
+			computedPreStateValues := make([][]byte, len(computedKeys))
+			for i := range computedKeys {
+				computedPreStateValues[i], err = preStateTree.GetWithHashedKey(computedKeys[i])
+				if err != nil {
+					panic(err)
+				}
+			}
+
 			// Write state changes to db
 			root, err := statedb.Commit(b.header.Number.Uint64(), config.IsEIP158(b.header.Number))
 			if err != nil {
@@ -420,9 +438,23 @@ func GenerateVerkleChain(config *params.ChainConfig, parent *types.Block, engine
 			proofs = append(proofs, block.ExecutionWitness().VerkleProof)
 			keyvals = append(keyvals, block.ExecutionWitness().StateDiff)
 
-			return block, b.receipts
+			computedPostStateValues := make([][]byte, len(computedKeys))
+			treeWrites := statedb.GetTrie().(*trie.VerkleTrie).GetTreeWrites()
+			for i := range computedKeys {
+				treeWrittenValue, ok := treeWrites[string(computedKeys[i])]
+				// If we didn't tracked this key, it means it was never written to the post-state tree,
+				// thus the newValue must be `nil`. (i.e: same as currentValue)
+				// Additionally, if we wrote to this key but it has the same pre-state value, then must
+				// also be `nil`.
+				if !ok || bytes.Equal(treeWrittenValue, computedPreStateValues[i]) {
+					continue
+				}
+				computedPostStateValues[i] = treeWrittenValue
+			}
+
+			return block, b.receipts, computedKeys, computedPreStateValues, computedPostStateValues
 		}
-		return nil, nil
+		return nil, nil, nil, nil, nil
 	}
 	var snaps *snapshot.Tree
 	for i := 0; i < n; i++ {
@@ -432,13 +464,16 @@ func GenerateVerkleChain(config *params.ChainConfig, parent *types.Block, engine
 		if err != nil {
 			panic(fmt.Sprintf("could not find state for block %d: err=%v, parent root=%x", i, err, parent.Root()))
 		}
-		block, receipt := genblock(i, parent, statedb)
+		block, receipt, keys, preStateValues, postStateValues := genblock(i, parent, statedb)
 		blocks[i] = block
 		receipts[i] = receipt
+		computedKeys[i] = keys
+		computedPreStateValues[i] = preStateValues
+		computedPostStateValues[i] = postStateValues
 		parent = block
 		snaps = statedb.Snaps()
 	}
-	return blocks, receipts, proofs, keyvals
+	return blocks, receipts, proofs, keyvals, computedKeys, computedPreStateValues, computedPostStateValues
 }
 
 func makeHeader(chain consensus.ChainReader, parent *types.Block, state *state.StateDB, engine consensus.Engine) *types.Header {

core/state_processor_test.go

@@ -18,6 +18,7 @@ package core
 
 import (
 	//"bytes"
+
 	"crypto/ecdsa"
 
 	//"fmt"
@@ -490,7 +491,7 @@ func TestProcessVerkle(t *testing.T) {
 		txCost1*2 + txCost2 + contractCreationCost + codeWithExtCodeCopyGas,
 	}
 	// TODO utiliser GenerateChainWithGenesis pour le rendre plus pratique
-	chain, _, proofs, keyvals := GenerateVerkleChain(gspec.Config, genesis, beacon.New(ethash.NewFaker()), gendb, 2, func(i int, gen *BlockGen) {
+	chain, _, proofs, statediff, computedKeys, computedPreStateValues, computedPostStateValues := GenerateVerkleChain(gspec.Config, genesis, beacon.New(ethash.NewFaker()), gendb, 2, func(i int, gen *BlockGen) {
 		gen.SetPoS()
 		// TODO need to check that the tx cost provided is the exact amount used (no remaining left-over)
 		tx, _ := types.SignTx(types.NewTransaction(uint64(i)*3, common.Address{byte(i), 2, 3}, big.NewInt(999), txCost1, big.NewInt(875000000), nil), signer, testKey)
@@ -518,11 +519,13 @@ func TestProcessVerkle(t *testing.T) {
 	//f.Write(buf.Bytes())
 	//fmt.Printf("root= %x\n", chain[0].Root())
 	// check the proof for the last block
-	err := trie.DeserializeAndVerifyVerkleProof(proofs[1], chain[0].Root().Bytes(), chain[1].Root().Bytes(), keyvals[1])
-	if err != nil {
-		t.Fatal(err)
+
+	for i := 1; i < len(proofs); i++ {
+		if err := trie.DeserializeAndVerifyVerkleProof(proofs[i], statediff[i], chain[i-1].Root().Bytes(), computedKeys[i], computedPreStateValues[i], computedPostStateValues[i]); err != nil {
+			t.Fatal(err)
+		}
+		t.Log("verfied verkle proof")
 	}
-	t.Log("verfied verkle proof")
 
 	endnum, err := blockchain.InsertChain(chain)
 	if err != nil {

miner/worker_test.go

@@ -158,7 +158,7 @@ func (b *testWorkerBackend) newRandomVerkleUncle() *types.Block {
 	} else {
 		parent = b.chain.GetBlockByHash(b.chain.CurrentBlock().ParentHash)
 	}
-	blocks, _, _, _ := core.GenerateVerkleChain(b.chain.Config(), parent, b.chain.Engine(), b.db, 1, func(i int, gen *core.BlockGen) {
+	blocks, _, _, _, _ := core.GenerateVerkleChain(b.chain.Config(), parent, b.chain.Engine(), b.db, 1, func(i int, gen *core.BlockGen) {
 		var addr = make([]byte, common.AddressLength)
 		rand.Read(addr)
 		gen.SetCoinbase(common.BytesToAddress(addr))

trie/verkle.go

@@ -39,6 +39,8 @@ type VerkleTrie struct {
 	db         *Database
 	pointCache *utils.PointCache
 	ended      bool
+
+	treeWrites map[string][]byte
 }
 
 func (vt *VerkleTrie) ToDot() string {
@@ -51,6 +53,7 @@ func NewVerkleTrie(root verkle.VerkleNode, db *Database, pointCache *utils.Point
 		db:         db,
 		pointCache: pointCache,
 		ended:      ended,
+		treeWrites: make(map[string][]byte),
 	}
 }
 
@@ -59,6 +62,7 @@ func (trie *VerkleTrie) FlatdbNodeResolver(path []byte) ([]byte, error) {
 }
 
 func (trie *VerkleTrie) InsertMigratedLeaves(leaves []verkle.LeafNode) error {
+	// Note: these values intentionally not inserted in the postValues map.
 	return trie.root.(*verkle.InternalNode).InsertMigratedLeaves(leaves, trie.FlatdbNodeResolver)
 }
 
@@ -185,16 +189,22 @@ func (t *VerkleTrie) UpdateAccount(addr common.Address, acc *types.StateAccount)
 	}
 	// TODO figure out if the code size needs to be updated, too
 
+	t.trackPostStateValues(stem, values)
+
 	return nil
 }
 
-func (trie *VerkleTrie) UpdateStem(key []byte, values [][]byte) error {
+func (trie *VerkleTrie) UpdateStem(stem []byte, values [][]byte) error {
 	switch root := trie.root.(type) {
 	case *verkle.InternalNode:
-		return root.InsertStem(key, values, trie.FlatdbNodeResolver)
+		if err := root.InsertStem(stem, values, trie.FlatdbNodeResolver); err != nil {
+			return fmt.Errorf("updating stem: %v", err)
+		}
+		trie.trackPostStateValues(stem, values)
 	default:
 		panic("invalid root type")
 	}
+	return nil
 }
 
 // Update associates key with value in the trie. If value has length zero, any
@@ -209,7 +219,13 @@ func (trie *VerkleTrie) UpdateStorage(address common.Address, key, value []byte)
 	} else {
 		copy(v[32-len(value):], value[:])
 	}
-	return trie.root.Insert(k, v[:], trie.FlatdbNodeResolver)
+	if err := trie.root.Insert(k, v[:], trie.FlatdbNodeResolver); err != nil {
+		return fmt.Errorf("inserting key: %s", err)
+	}
+
+	trie.treeWrites[string(k)] = v[:]
+
+	return nil
 }
 
 func (t *VerkleTrie) DeleteAccount(addr common.Address) error {
@@ -234,6 +250,8 @@ func (t *VerkleTrie) DeleteAccount(addr common.Address) error {
 	}
 	// TODO figure out if the code size needs to be updated, too
 
+	t.trackPostStateValues(stem, values)
+
 	return nil
 }
 
@@ -243,7 +261,12 @@ func (trie *VerkleTrie) DeleteStorage(addr common.Address, key []byte) error {
 	pointEval := trie.pointCache.GetTreeKeyHeader(addr[:])
 	k := utils.GetTreeKeyStorageSlotWithEvaluatedAddress(pointEval, key)
 	var zero [32]byte
-	return trie.root.Insert(k, zero[:], trie.FlatdbNodeResolver)
+	if err := trie.root.Insert(k, zero[:], trie.FlatdbNodeResolver); err != nil {
+		return fmt.Errorf("inserting key: %s", err)
+	}
+	trie.treeWrites[string(k)] = zero[:]
+
+	return nil
 }
 
 // Hash returns the root hash of the trie. It does not write to the database and
@@ -306,6 +329,13 @@ func (trie *VerkleTrie) Prove(key []byte, proofDb ethdb.KeyValueWriter) error {
 	panic("not implemented")
 }
 
+// GetTreeWrites returns the a map that contains the addresses and values that were
+// written to the trie. The returned map **is not** a copy, so any mutation to it
+// can affect further calls. It's recommended to treat it as read-only.
+func (trie *VerkleTrie) GetTreeWrites() map[string][]byte {
+	return trie.treeWrites
+}
+
 func (trie *VerkleTrie) Copy() *VerkleTrie {
 	return &VerkleTrie{
 		root:       trie.root.Copy(),
@@ -336,14 +366,48 @@ func ProveAndSerialize(pretrie, posttrie *VerkleTrie, keys [][]byte, resolver ve
 	return p, kvps, nil
 }
 
-func DeserializeAndVerifyVerkleProof(vp *verkle.VerkleProof, preStateRoot []byte, postStateRoot []byte, statediff verkle.StateDiff) error {
-	// TODO: check that `OtherStems` have expected length and values.
-
+func DeserializeAndVerifyVerkleProof(
+	vp *verkle.VerkleProof,
+	statediff verkle.StateDiff,
+	preStateRoot []byte,
+	computedKeys [][]byte,
+	computedPreStateValues [][]byte,
+	computedPostStateValues [][]byte) error {
 	proof, err := verkle.DeserializeProof(vp, statediff)
 	if err != nil {
 		return fmt.Errorf("verkle proof deserialization error: %w", err)
 	}
 
+	// Verify the provided `statediff` by checking that the keys, pre-values and post-values match exactly
+	// with the ones provided from the EVM block execution witness.
+	if len(computedKeys) != len(proof.Keys) {
+		return fmt.Errorf("witness keys length doesn't match proof keys length: expected %d, got %d", len(computedKeys), len(proof.Keys))
+	}
+	for i := range computedKeys {
+		if !bytes.Equal(computedKeys[i], proof.Keys[i]) {
+			return fmt.Errorf("witness keys don't match proof keys: expected %x, got %x", computedKeys[i], proof.Keys[i])
+		}
+	}
+	if len(computedPreStateValues) != len(proof.PreValues) {
+		return fmt.Errorf("witness pre-values length doesn't match proof pre-values length: expected %d, got %d", len(computedPreStateValues), len(proof.PreValues))
+	}
+	for i := range computedPreStateValues {
+		if !bytes.Equal(computedPreStateValues[i], proof.PreValues[i]) {
+			return fmt.Errorf("witness pre-values don't match proof pre-values: expected %x, got %x", computedPreStateValues[i], proof.PreValues[i])
+		}
+	}
+	if len(computedPostStateValues) != len(proof.PostValues) {
+		return fmt.Errorf("witness post-values length doesn't match proof post-values length: expected %d, got %d", len(computedPostStateValues), len(proof.PostValues))
+
+	}
+	for i := range computedPostStateValues {
+		if !bytes.Equal(computedPostStateValues[i], proof.PostValues[i]) {
+			return fmt.Errorf("witness post-values don't match proof post-values: expected %x, got %x", computedPostStateValues[i], proof.PostValues[i])
+		}
+	}
+
+	// At the point we know that the pre and post values are correct, we we proceed with reconstructing
+	// the pre-state tree, and getting the elements to verify the cryptographic proof.
 	rootC := new(verkle.Point)
 	rootC.SetBytes(preStateRoot)
 	pretree, err := verkle.PreStateTreeFromProof(proof, rootC)
@@ -374,19 +438,6 @@ func DeserializeAndVerifyVerkleProof(vp *verkle.VerkleProof, preStateRoot []byte
 		}
 	}
 
-	// TODO: this is necessary to verify that the post-values are the correct ones.
-	// But all this can be avoided with a even faster way. The EVM block execution can
-	// keep track of the written keys, and compare that list with this post-values list.
-	// This can avoid regenerating the post-tree which is somewhat expensive.
-	posttree, err := verkle.PostStateTreeFromStateDiff(pretree, statediff)
-	if err != nil {
-		return fmt.Errorf("error rebuilding the post-tree from proof: %w", err)
-	}
-	regeneratedPostTreeRoot := posttree.Commitment().Bytes()
-	if !bytes.Equal(regeneratedPostTreeRoot[:], postStateRoot) {
-		return fmt.Errorf("post tree root mismatch: %x != %x", regeneratedPostTreeRoot, postStateRoot)
-	}
-
 	return verkle.VerifyVerkleProofWithPreState(proof, pretree)
 }
 
@@ -496,3 +547,16 @@ func (t *VerkleTrie) UpdateContractCode(addr common.Address, codeHash common.Has
 	}
 	return nil
 }
+
+func (trie *VerkleTrie) trackPostStateValues(stem []byte, values [][]byte) {
+	addr := make([]byte, verkle.StemSize+1)
+	copy(addr[:verkle.StemSize], stem)
+	for i := range values {
+		if len(values[i]) == 0 {
+			continue
+		}
+		addr[verkle.StemSize] = byte(i)
+		fmt.Printf("Tracking %x: %x\n", string(addr), values[i])
+		trie.treeWrites[string(addr)] = values[i]
+	}
+}