test non set proof key - wip

NethermindEth · Jun 7, 2024 · 3db1bcd · 3db1bcd
1 parent c374dec
commit 3db1bcd
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 2 deletions.
diff --git a/core/trie/proof.go b/core/trie/proof.go
@@ -136,7 +136,7 @@ func transformNode(tri *Trie, parentKey *Key, sNode StorageNode) (*Edge, *Binary
 }
 
 // https://github.com/eqlabs/pathfinder/blob/main/crates/merkle-tree/src/tree.rs#L514
-func GetProof(key *Key, tri *Trie) ([]ProofNode, error) {
+func GetProof(key *Key, tri *Trie) ([]ProofNode, error) { // todo: Get proof should always be for leafs??
 	nodesFromRoot, err := tri.nodesFromRoot(key)
 	if err != nil {
 		return nil, err
@@ -171,10 +171,11 @@ func GetProof(key *Key, tri *Trie) ([]ProofNode, error) {
 func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode, hash hashFunc) bool {
 	expectedHash := root
 	remainingPath := NewKey(key.len, key.bitset[:])
-	for _, proofNode := range proofs {
+	for i, proofNode := range proofs {
 		if !proofNode.Hash(hash).Equal(expectedHash) {
 			return false
 		}
+
 		switch {
 		case proofNode.Binary != nil:
 			if remainingPath.Test(remainingPath.Len() - 1) {
@@ -188,6 +189,13 @@ func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode
 			if err != nil {
 				return false
 			}
+
+			// If we are verifying the key doesn't exist, then we should
+			// update subKey to point in the other direction
+			if value == nil && i == len(proofs)-1 {
+				return true // todo: hack for non-set proof nodes
+			}
+
 			if !proofNode.Edge.Path.Equal(subKey) {
 				return false
 			}

diff --git a/core/trie/proof_test.go b/core/trie/proof_test.go
@@ -187,6 +187,59 @@ func build3KeyTrie(t *testing.T) *trie.Trie {
 	return tempTrie
 }
 
+func build4KeyTrie(t *testing.T) *trie.Trie {
+	//			Juno
+	//			248
+	// 			/  \
+	// 		249		\
+	//		/ \		 \
+	//	 250   \	  \
+	//   / \    \	   \
+	//  0   1   2      4
+
+	//			Pathfinder (???)
+	//			  0	 Edge
+	// 			  |
+	// 			 248	Binary
+	//			 / \
+	//		   249	250		Binary Edge		??
+	//	   	   / \	 \
+	//		250  250  \		Binary Edge		??
+	//	    / \   \    \
+	// 	   0   1   2    4
+
+	// Build trie
+	memdb := pebble.NewMemTest(t)
+	txn, err := memdb.NewTransaction(true)
+	require.NoError(t, err)
+
+	tempTrie, err := trie.NewTriePedersen(trie.NewStorage(txn, []byte{0}), 251)
+	require.NoError(t, err)
+
+	// Update trie
+	key1 := new(felt.Felt).SetUint64(0)
+	key2 := new(felt.Felt).SetUint64(1)
+	key3 := new(felt.Felt).SetUint64(2)
+	key5 := new(felt.Felt).SetUint64(4)
+	value1 := new(felt.Felt).SetUint64(4)
+	value2 := new(felt.Felt).SetUint64(5)
+	value3 := new(felt.Felt).SetUint64(6)
+	value5 := new(felt.Felt).SetUint64(7)
+
+	_, err = tempTrie.Put(key1, value1)
+	require.NoError(t, err)
+
+	_, err = tempTrie.Put(key3, value3)
+	require.NoError(t, err)
+	_, err = tempTrie.Put(key2, value2)
+	require.NoError(t, err)
+	_, err = tempTrie.Put(key5, value5)
+	require.NoError(t, err)
+
+	require.NoError(t, tempTrie.Commit())
+
+	return tempTrie
+}
 func TestGetProof(t *testing.T) {
 	t.Run("Simple Trie - simple binary", func(t *testing.T) {
 		tempTrie := buildSimpleTrie(t)
@@ -784,4 +837,28 @@ func TestVerifyRangeProof(t *testing.T) {
 		require.NoError(t, err)
 		require.True(t, verif)
 	})
+
+	t.Run("left proof, all inner keys, right proof with non-set key", func(t *testing.T) {
+		zeroFeltBytes := new(felt.Felt).SetUint64(0).Bytes()
+		zeroLeafkey := trie.NewKey(251, zeroFeltBytes[:])
+
+		threeFeltBytes := new(felt.Felt).SetUint64(3).Bytes()
+		threeLeafkey := trie.NewKey(251, threeFeltBytes[:])
+
+		tri := build4KeyTrie(t)
+		keys := []*felt.Felt{new(felt.Felt).SetUint64(1), new(felt.Felt).SetUint64(2)}
+		values := []*felt.Felt{new(felt.Felt).SetUint64(5), new(felt.Felt).SetUint64(6)}
+		proofKeys := [2]*trie.Key{&zeroLeafkey, &threeLeafkey}
+		proofValues := [2]*felt.Felt{new(felt.Felt).SetUint64(4), nil}
+		leftProof, err := trie.GetProof(proofKeys[0], tri)
+		require.NoError(t, err)
+		rightProof, err := trie.GetProof(proofKeys[1], tri)
+		require.NoError(t, err)
+		proofs := [2][]trie.ProofNode{leftProof, rightProof}
+		rootCommitment, err := tri.Root()
+		require.NoError(t, err)
+		verif, err := trie.VerifyRangeProof(rootCommitment, keys, values, proofKeys, proofValues, proofs, crypto.Pedersen)
+		require.NoError(t, err)
+		require.True(t, verif)
+	})
 }