Skip to content

Commit

Permalink
Update GetProof and VerifyProof to use Keys (#1893)
Browse files Browse the repository at this point in the history 
* Update proof logic to support inner nodes and keys


---------

Co-authored-by: rian <[email protected]>
  • Loading branch information
@rianhughes
rianhughes and rian authored Jun 7, 2024
1 parent 03c1985 commit 99afb51
Show file tree
Hide file tree
Showing 3 changed files with 148 additions and 76 deletions.
7 changes: 4 additions & 3 deletions core/trie/key.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package trie
import (
"bytes"
"encoding/hex"
"errors"
"fmt"
"math/big"

Expand All @@ -23,9 +24,9 @@ func NewKey(length uint8, keyBytes []byte) Key {
return k
}

func (k *Key) SubKey(n uint8) *Key {
func (k *Key) SubKey(n uint8) (*Key, error) {
if n > k.len {
panic("n is greater than the length of the key")
return nil, errors.New(fmt.Sprint("cannot subtract key of length %i from key of length %i", n, k.len))
}

newKey := &Key{len: n}
Expand All @@ -40,7 +41,7 @@ func (k *Key) SubKey(n uint8) *Key {
}
}

return newKey
return newKey, nil
}

func (k *Key) bytesNeeded() uint {
Expand Down
23 changes: 10 additions & 13 deletions core/trie/proof.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,18 +113,16 @@ 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(leaf *felt.Felt, tri *Trie) ([]ProofNode, error) {
leafKey := tri.feltToKey(leaf)
nodesToLeaf, err := tri.nodesFromRoot(&leafKey)
func GetProof(key *Key, tri *Trie) ([]ProofNode, error) {
nodesFromRoot, err := tri.nodesFromRoot(key)
if err != nil {
return nil, err
}
proofNodes := []ProofNode{}

var parentKey *Key

for i := 0; i < len(nodesToLeaf); i++ {
sNode := nodesToLeaf[i]
for i, sNode := range nodesFromRoot {
sNodeEdge, sNodeBinary, err := transformNode(tri, parentKey, sNode)
if err != nil {
return nil, err
Expand All @@ -140,21 +138,16 @@ func GetProof(leaf *felt.Felt, tri *Trie) ([]ProofNode, error) {
} else if sNodeEdge == nil && sNodeBinary == nil { // sNode is a binary leaf
break
}
parentKey = nodesToLeaf[i].key
parentKey = nodesFromRoot[i].key
}
return proofNodes, nil
}

// verifyProof checks if `leafPath` leads from `root` to `leafHash` along the `proofNodes`
// https://github.com/eqlabs/pathfinder/blob/main/crates/merkle-tree/src/tree.rs#L2006
func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode, hash hashFunc) bool {
if key.Len() != 251 { //nolint:gomnd
return false
}

expectedHash := root
remainingPath := key

for _, proofNode := range proofs {
if !proofNode.Hash(hash).Equal(expectedHash) {
return false
Expand All @@ -168,11 +161,15 @@ func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode
}
remainingPath.RemoveLastBit()
case proofNode.Edge != nil:
if !proofNode.Edge.Path.Equal(remainingPath.SubKey(proofNode.Edge.Path.Len())) {
subKey, err := remainingPath.SubKey(proofNode.Edge.Path.Len())
if err != nil {
return false
}
if !proofNode.Edge.Path.Equal(subKey) {
return false
}
expectedHash = proofNode.Edge.Child
remainingPath.Truncate(proofNode.Edge.Path.Len())
remainingPath.Truncate(251 - proofNode.Edge.Path.Len()) //nolint:gomnd
}
}
return expectedHash.Equal(value)
Expand Down
194 changes: 134 additions & 60 deletions core/trie/proof_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,15 @@ func buildSimpleBinaryRootTrie(t *testing.T) *trie.Trie {
return tempTrie
}

func buildSimpleDoubleBinaryTrie(t *testing.T) *trie.Trie {
func buildSimpleDoubleBinaryTrie(t *testing.T) (*trie.Trie, []trie.ProofNode) {
// (249,0,x3) // Edge
// |
// (0, 0, x3) // Binary
// / \
// (0,0,x1) // B (1, 1, 5) // Edge leaf
// / \ |
// (2) (3) (5)

// Build trie
memdb := pebble.NewMemTest(t)
txn, err := memdb.NewTransaction(true)
Expand All @@ -99,10 +107,35 @@ func buildSimpleDoubleBinaryTrie(t *testing.T) *trie.Trie {
require.NoError(t, err)

require.NoError(t, tempTrie.Commit())
return tempTrie

zero := trie.NewKey(249, []byte{0})
key3Bytes := new(felt.Felt).SetUint64(1).Bytes()
path3 := trie.NewKey(1, key3Bytes[:])
expectedProofNodes := []trie.ProofNode{
{
Edge: &trie.Edge{
Path: &zero,
Child: utils.HexToFelt(t, "0x055C81F6A791FD06FC2E2CCAD922397EC76C3E35F2E06C0C0D43D551005A8DEA"),
},
},
{
Binary: &trie.Binary{
LeftHash: utils.HexToFelt(t, "0x05774FA77B3D843AE9167ABD61CF80365A9B2B02218FC2F628494B5BDC9B33B8"),
RightHash: utils.HexToFelt(t, "0x07C5BC1CC68B7BC8CA2F632DE98297E6DA9594FA23EDE872DD2ABEAFDE353B43"),
},
},
{
Edge: &trie.Edge{
Path: &path3,
Child: value3,
},
},
}

return tempTrie, expectedProofNodes
}

func TestGetProofs(t *testing.T) {
func TestGetProof(t *testing.T) {
t.Run("Simple Trie - simple binary", func(t *testing.T) {
tempTrie := buildSimpleTrie(t)

Expand All @@ -121,8 +154,9 @@ func TestGetProofs(t *testing.T) {
},
},
}

proofNodes, err := trie.GetProof(new(felt.Felt).SetUint64(0), tempTrie)
leafFelt := new(felt.Felt).SetUint64(0).Bytes()
leafKey := trie.NewKey(251, leafFelt[:])
proofNodes, err := trie.GetProof(&leafKey, tempTrie)
require.NoError(t, err)

// Better inspection
Expand All @@ -133,31 +167,18 @@ func TestGetProofs(t *testing.T) {
})

t.Run("Simple Trie - simple double binary", func(t *testing.T) {
tempTrie := buildSimpleDoubleBinaryTrie(t)
tempTrie, expectedProofNodes := buildSimpleDoubleBinaryTrie(t)

zero := trie.NewKey(249, []byte{0})
expectedProofNodes := []trie.ProofNode{
{
Edge: &trie.Edge{
Path: &zero,
Child: utils.HexToFelt(t, "0x055C81F6A791FD06FC2E2CCAD922397EC76C3E35F2E06C0C0D43D551005A8DEA"),
},
},
{
Binary: &trie.Binary{
LeftHash: utils.HexToFelt(t, "0x05774FA77B3D843AE9167ABD61CF80365A9B2B02218FC2F628494B5BDC9B33B8"),
RightHash: utils.HexToFelt(t, "0x07C5BC1CC68B7BC8CA2F632DE98297E6DA9594FA23EDE872DD2ABEAFDE353B43"),
},
},
{
Binary: &trie.Binary{
LeftHash: utils.HexToFelt(t, "0x0000000000000000000000000000000000000000000000000000000000000002"),
RightHash: utils.HexToFelt(t, "0x0000000000000000000000000000000000000000000000000000000000000003"),
},
expectedProofNodes[2] = trie.ProofNode{
Binary: &trie.Binary{
LeftHash: utils.HexToFelt(t, "0x0000000000000000000000000000000000000000000000000000000000000002"),
RightHash: utils.HexToFelt(t, "0x0000000000000000000000000000000000000000000000000000000000000003"),
},
}

proofNodes, err := trie.GetProof(new(felt.Felt).SetUint64(0), tempTrie)
leafFelt := new(felt.Felt).SetUint64(0).Bytes()
leafKey := trie.NewKey(251, leafFelt[:])
proofNodes, err := trie.GetProof(&leafKey, tempTrie)
require.NoError(t, err)

// Better inspection
Expand All @@ -168,34 +189,10 @@ func TestGetProofs(t *testing.T) {
})

t.Run("Simple Trie - simple double binary edge", func(t *testing.T) {
tempTrie := buildSimpleDoubleBinaryTrie(t)

zero := trie.NewKey(249, []byte{0})
value3 := new(felt.Felt).SetUint64(5)
key3Bytes := new(felt.Felt).SetUint64(1).Bytes()
path3 := trie.NewKey(1, key3Bytes[:])
expectedProofNodes := []trie.ProofNode{
{
Edge: &trie.Edge{
Path: &zero,
Child: utils.HexToFelt(t, "0x055C81F6A791FD06FC2E2CCAD922397EC76C3E35F2E06C0C0D43D551005A8DEA"),
},
},
{
Binary: &trie.Binary{
LeftHash: utils.HexToFelt(t, "0x05774FA77B3D843AE9167ABD61CF80365A9B2B02218FC2F628494B5BDC9B33B8"),
RightHash: utils.HexToFelt(t, "0x07C5BC1CC68B7BC8CA2F632DE98297E6DA9594FA23EDE872DD2ABEAFDE353B43"),
},
},
{
Edge: &trie.Edge{
Path: &path3,
Child: value3,
},
},
}

proofNodes, err := trie.GetProof(new(felt.Felt).SetUint64(3), tempTrie)
tempTrie, expectedProofNodes := buildSimpleDoubleBinaryTrie(t)
leafFelt := new(felt.Felt).SetUint64(3).Bytes()
leafKey := trie.NewKey(251, leafFelt[:])
proofNodes, err := trie.GetProof(&leafKey, tempTrie)
require.NoError(t, err)

// Better inspection
Expand Down Expand Up @@ -224,8 +221,10 @@ func TestGetProofs(t *testing.T) {
},
},
}
leafFelt := new(felt.Felt).SetUint64(0).Bytes()
leafKey := trie.NewKey(251, leafFelt[:])

proofNodes, err := trie.GetProof(new(felt.Felt).SetUint64(0), tempTrie)
proofNodes, err := trie.GetProof(&leafKey, tempTrie)
require.NoError(t, err)

// Better inspection
Expand Down Expand Up @@ -267,8 +266,9 @@ func TestGetProofs(t *testing.T) {
},
},
}

proofNodes, err := trie.GetProof(new(felt.Felt).SetUint64(0), tempTrie)
leafFelt := new(felt.Felt).SetUint64(0).Bytes()
leafKey := trie.NewKey(251, leafFelt[:])
proofNodes, err := trie.GetProof(&leafKey, tempTrie)
require.NoError(t, err)

// Better inspection
Expand All @@ -278,9 +278,54 @@ func TestGetProofs(t *testing.T) {
}
require.Equal(t, expectedProofNodes, proofNodes)
})

t.Run("Simple Trie - proof for non-set key", func(t *testing.T) {
tempTrie, expectedProofNodes := buildSimpleDoubleBinaryTrie(t)

leafFelt := new(felt.Felt).SetUint64(123).Bytes() // The (root) edge node would have a shorter len if this key was set
leafKey := trie.NewKey(251, leafFelt[:])
proofNodes, err := trie.GetProof(&leafKey, tempTrie)
require.NoError(t, err)

// Better inspection
for _, pNode := range proofNodes {
pNode.PrettyPrint()
}
require.Equal(t, expectedProofNodes[0:2], proofNodes)
})

t.Run("Simple Trie - proof for inner key", func(t *testing.T) {
tempTrie, expectedProofNodes := buildSimpleDoubleBinaryTrie(t)

innerFelt := new(felt.Felt).SetUint64(2).Bytes()
innerKey := trie.NewKey(123, innerFelt[:]) // The (root) edge node has len 249 which shows this doesn't exist
proofNodes, err := trie.GetProof(&innerKey, tempTrie)
require.NoError(t, err)

// Better inspection
for _, pNode := range proofNodes {
pNode.PrettyPrint()
}
require.Equal(t, expectedProofNodes[0:2], proofNodes)
})

t.Run("Simple Trie - proof for non-set key, with leafs set to right and left", func(t *testing.T) {
tempTrie, expectedProofNodes := buildSimpleDoubleBinaryTrie(t)

leafFelt := new(felt.Felt).SetUint64(2).Bytes()
leafKey := trie.NewKey(251, leafFelt[:])
proofNodes, err := trie.GetProof(&leafKey, tempTrie)
require.NoError(t, err)

// Better inspection
for _, pNode := range proofNodes {
pNode.PrettyPrint()
}
require.Equal(t, expectedProofNodes, proofNodes)
})
}

func TestVerifyProofs(t *testing.T) {
func TestVerifyProof(t *testing.T) {
// https://github.com/eqlabs/pathfinder/blob/main/crates/merkle-tree/src/tree.rs#L2137
t.Run("Simple binary trie", func(t *testing.T) {
tempTrie := buildSimpleTrie(t)
Expand Down Expand Up @@ -310,7 +355,7 @@ func TestVerifyProofs(t *testing.T) {

// https://github.com/eqlabs/pathfinder/blob/main/crates/merkle-tree/src/tree.rs#L2167
t.Run("Simple double binary trie", func(t *testing.T) {
tempTrie := buildSimpleDoubleBinaryTrie(t)
tempTrie, _ := buildSimpleDoubleBinaryTrie(t)
zero := trie.NewKey(249, []byte{0})
expectedProofNodes := []trie.ProofNode{
{
Expand Down Expand Up @@ -338,6 +383,35 @@ func TestVerifyProofs(t *testing.T) {
key1Bytes := new(felt.Felt).SetUint64(0).Bytes()
key1 := trie.NewKey(251, key1Bytes[:])
val1 := new(felt.Felt).SetUint64(2)
assert.True(t, trie.VerifyProof(root, &key1, val1, expectedProofNodes, crypto.Pedersen))
require.True(t, trie.VerifyProof(root, &key1, val1, expectedProofNodes, crypto.Pedersen))
})

t.Run("non existent key - less than root edge", func(t *testing.T) {
tempTrie, _ := buildSimpleDoubleBinaryTrie(t)

nonExistentKey := trie.NewKey(123, []byte{0}) // Diverges before the root node (len root node = 249)
nonExistentKeyValue := new(felt.Felt).SetUint64(2)
proofNodes, err := trie.GetProof(&nonExistentKey, tempTrie)
require.NoError(t, err)

root, err := tempTrie.Root()
require.NoError(t, err)

require.False(t, trie.VerifyProof(root, &nonExistentKey, nonExistentKeyValue, proofNodes, crypto.Pedersen))
})

t.Run("non existent leaf key", func(t *testing.T) {
tempTrie, _ := buildSimpleDoubleBinaryTrie(t)

nonExistentKeyByte := new(felt.Felt).SetUint64(2).Bytes() // Key not set
nonExistentKey := trie.NewKey(251, nonExistentKeyByte[:])
nonExistentKeyValue := new(felt.Felt).SetUint64(2)
proofNodes, err := trie.GetProof(&nonExistentKey, tempTrie)
require.NoError(t, err)

root, err := tempTrie.Root()
require.NoError(t, err)

require.False(t, trie.VerifyProof(root, &nonExistentKey, nonExistentKeyValue, proofNodes, crypto.Pedersen))
})
}

0 comments on commit 99afb51

Please sign in to comment.