Skip to content

Commit

Permalink
cleaned tests add proof verification, addressed comments
Browse files Browse the repository at this point in the history
  • Loading branch information
pnowosie committed Oct 11, 2024
1 parent 31182fd commit fcf4395
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 62 deletions.
2 changes: 1 addition & 1 deletion rpc/handlers.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ var (
ErrSubscriptionNotFound = &jsonrpc.Error{Code: 100, Message: "Subscription not found"}

// TODO[pnowosie]: Update the error while specification describe it
ErrBlockNotRecentForProof = &jsonrpc.Error{Code: 1001, Message: "Block is not sufficiently recent for storage proofs"}
ErrBlockNotRecentForProof = &jsonrpc.Error{Code: 1001, Message: "Block is not sufficiently recent for storage proofs. Use 'latest' as block id"}
)

const (
Expand Down
42 changes: 28 additions & 14 deletions rpc/storage.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package rpc

import (
"errors"
"fmt"

"github.com/NethermindEth/juno/core"
"github.com/NethermindEth/juno/core/felt"
Expand Down Expand Up @@ -50,7 +49,7 @@ func (h *Handler) StorageProof(id BlockID, classes, contracts []felt.Felt, stora
return nil, ErrInternal.CloneWithData(err)
}

if !(id.Latest || head.Number == id.Number) {
if !id.Latest {
return nil, ErrBlockNotRecentForProof
}

Expand Down Expand Up @@ -92,22 +91,41 @@ type StorageKeys struct {

// MerkleNode represents a proof node in a trie
// https://github.com/starkware-libs/starknet-specs/blob/647caa00c0223e1daab1b2f3acc4e613ba2138aa/api/starknet_api_openrpc.json#L3632
type MerkleNode interface{}
// Implemented by MerkleBinaryNode, MerkleEdgeNode
type MerkleNode interface {
AsProofNode() trie.ProofNode
}

// https://github.com/starkware-libs/starknet-specs/blob/647caa00c0223e1daab1b2f3acc4e613ba2138aa/api/starknet_api_openrpc.json#L3644
type MerkleBinaryNode struct {
Left *felt.Felt `json:"left"`
Right *felt.Felt `json:"right"`
}

func (mbn *MerkleBinaryNode) AsProofNode() trie.ProofNode {
return &trie.Binary{
LeftHash: mbn.Left,
RightHash: mbn.Right,
}
}

// TODO[pnowosie]: link to specs
type MerkleEdgeNode struct {
Path *felt.Felt `json:"path"`
Length int `json:"length"`
Child *felt.Felt `json:"child"`
}

// TODO[pnowosie]: link to specs
func (men *MerkleEdgeNode) AsProofNode() trie.ProofNode {
pbs := men.Path.Bytes()
path := trie.NewKey(uint8(men.Length), pbs[:])
return &trie.Edge{
Path: &path,
Child: men.Child,
}
}

// TODO[pnowosie]: link to specs, but hoping for removal
type MerkleLeafNode struct {
Value *felt.Felt `json:"value"`
}
Expand Down Expand Up @@ -147,13 +165,13 @@ type StorageProofResult struct {
}

func getClassesProof(reader core.StateReader, classes []felt.Felt) ([]*HashToNode, error) {
ctrie, _, err := reader.ClassTrie()
cTrie, _, err := reader.ClassTrie()
if err != nil {
return nil, err
}
result := []*HashToNode{}
for _, class := range classes {
nodes, err := getProof(ctrie, &class)
nodes, err := getProof(cTrie, &class)
if err != nil {
return nil, err
}
Expand All @@ -163,7 +181,7 @@ func getClassesProof(reader core.StateReader, classes []felt.Felt) ([]*HashToNod
}

func getContractsProof(reader core.StateReader, contracts []felt.Felt) (*ContractProof, error) {
strie, _, err := reader.StorageTrie()
sTrie, _, err := reader.StorageTrie()
if err != nil {
return nil, err
}
Expand All @@ -180,7 +198,7 @@ func getContractsProof(reader core.StateReader, contracts []felt.Felt) (*Contrac
}
result.LeavesData = append(result.LeavesData, leafData)

nodes, err := getProof(strie, &contract)
nodes, err := getProof(sTrie, &contract)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -213,15 +231,15 @@ func getContractsStorageProofs(reader core.StateReader, keys []StorageKeys) ([][
result := make([][]*HashToNode, 0, len(keys))

for _, key := range keys {
cstrie, err := reader.StorageTrieForAddr(&key.Contract)
csTrie, err := reader.StorageTrieForAddr(&key.Contract)
if err != nil {
// Note: if contract does not exist, `StorageTrieForAddr()` returns an empty trie, not an error
return nil, err
}

nodes := []*HashToNode{}
for _, slot := range key.Keys {
proof, err := getProof(cstrie, &slot)
proof, err := getProof(csTrie, &slot)
if err != nil {
return nil, err
}
Expand All @@ -237,10 +255,6 @@ func getProof(t *trie.Trie, elt *felt.Felt) ([]*HashToNode, error) {
feltBytes := elt.Bytes()
key := trie.NewKey(core.ContractStorageTrieHeight, feltBytes[:])
nodes, err := trie.GetProof(&key, t)
for i, n := range nodes {
fmt.Printf("[%d]", i)
n.PrettyPrint()
}
if err != nil {
return nil, err
}
Expand Down
115 changes: 68 additions & 47 deletions rpc/storage_test.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
package rpc_test

import (
"encoding/json"
"errors"
"fmt"
"testing"

"github.com/NethermindEth/juno/core"
Expand Down Expand Up @@ -110,6 +108,7 @@ func TestStorageProof(t *testing.T) {
key = new(felt.Felt).SetUint64(1)
noSuchKey = new(felt.Felt).SetUint64(0)
value = new(felt.Felt).SetUint64(51)
blockLatest = rpc.BlockID{Latest: true}
blockNumber = uint64(1313)
nopCloser = func() error {
return nil
Expand Down Expand Up @@ -160,7 +159,19 @@ func TestStorageProof(t *testing.T) {
log := utils.NewNopZapLogger()
handler := rpc.New(mockReader, nil, nil, "", log)

blockLatest := rpc.BlockID{Latest: true}
verifyIf := func(proof []*rpc.HashToNode, key *felt.Felt, value *felt.Felt) {
root, err := tempTrie.Root()
require.NoError(t, err)

pnodes := []trie.ProofNode{}
for _, hn := range proof {
pnodes = append(pnodes, hn.Node.AsProofNode())
}

kbs := key.Bytes()
kkey := trie.NewKey(251, kbs[:])
require.True(t, trie.VerifyProof(root, &kkey, value, pnodes, tempTrie.HashFunc()))
}

t.Run("Trie proofs sanity check", func(t *testing.T) {
kbs := key.Bytes()
Expand All @@ -170,6 +181,13 @@ func TestStorageProof(t *testing.T) {
root, err := tempTrie.Root()
require.NoError(t, err)
require.True(t, trie.VerifyProof(root, &kKey, value, proof, tempTrie.HashFunc()))

// non-membership test
kbs = noSuchKey.Bytes()
kKey = trie.NewKey(251, kbs[:])
proof, err = trie.GetProof(&kKey, tempTrie)
require.NoError(t, err)
require.True(t, trie.VerifyProof(root, &kKey, nil, proof, tempTrie.HashFunc()))
})
t.Run("global roots are filled", func(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, nil, nil, nil)
Expand All @@ -186,62 +204,60 @@ func TestStorageProof(t *testing.T) {
assert.Equal(t, rpc.ErrBlockNotRecentForProof, rpcErr)
require.Nil(t, proof)
})
t.Run("no error when blknum matches head", func(t *testing.T) {
t.Run("error is returned even when blknum matches head", func(t *testing.T) {
proof, rpcErr := handler.StorageProof(rpc.BlockID{Number: blockNumber}, nil, nil, nil)
assert.Nil(t, rpcErr)
assert.Equal(t, rpc.ErrBlockNotRecentForProof, rpcErr)
require.Nil(t, proof)
})
t.Run("empty request", func(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, nil, nil, nil)
require.Nil(t, rpcErr)
require.NotNil(t, proof)
arityTest(t, proof, 0, 0, 0, 0)
})
t.Run("class trie hash does not exist in a trie", func(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, []felt.Felt{*noSuchKey}, nil, nil)
require.Nil(t, rpcErr)
require.NotNil(t, proof)
require.NotNil(t, proof.ClassesProof)
require.True(t, len(proof.ClassesProof) > 0)
arityTest(t, proof, 3, 0, 0, 0)
verifyIf(proof.ClassesProof, noSuchKey, nil)
})
t.Run("class trie hash exists in a trie", func(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, []felt.Felt{*key}, nil, nil)
require.Nil(t, rpcErr)
require.NotNil(t, proof)
require.True(t, len(proof.ClassesProof) > 0)
require.Len(t, proof.ContractsStorageProofs, 0)
require.NotNil(t, proof.ContractsProof)
require.Len(t, proof.ContractsProof.Nodes, 0)
jsonStr, err := json.Marshal(proof)
require.NoError(t, err)
fmt.Println(string(jsonStr))
arityTest(t, proof, 3, 0, 0, 0)
verifyIf(proof.ClassesProof, key, value)
})
t.Run("storage trie address does not exist in a trie", func(t *testing.T) {
mockState.EXPECT().ContractNonce(noSuchKey).Return(nil, db.ErrKeyNotFound)
mockState.EXPECT().ContractClassHash(noSuchKey).Return(nil, db.ErrKeyNotFound)
mockState.EXPECT().ContractNonce(noSuchKey).Return(nil, db.ErrKeyNotFound).Times(1)
mockState.EXPECT().ContractClassHash(noSuchKey).Return(nil, db.ErrKeyNotFound).Times(0)

proof, rpcErr := handler.StorageProof(blockLatest, nil, []felt.Felt{*noSuchKey}, nil)
require.Nil(t, rpcErr)
require.NotNil(t, proof)
require.Len(t, proof.ClassesProof, 0)
require.Len(t, proof.ContractsStorageProofs, 0)
require.NotNil(t, proof.ContractsProof)
require.True(t, len(proof.ContractsProof.Nodes) > 0)
require.Len(t, proof.ContractsProof.LeavesData, 1)
arityTest(t, proof, 0, 3, 1, 0)
require.Nil(t, proof.ContractsProof.LeavesData[0])

verifyIf(proof.ContractsProof.Nodes, noSuchKey, nil)
})
t.Run("storage trie address exists in a trie", func(t *testing.T) {
nonce := new(felt.Felt).SetUint64(121)
mockState.EXPECT().ContractNonce(key).Return(nonce, nil)
mockState.EXPECT().ContractNonce(key).Return(nonce, nil).Times(1)
classHasah := new(felt.Felt).SetUint64(1234)
mockState.EXPECT().ContractClassHash(key).Return(classHasah, nil)
mockState.EXPECT().ContractClassHash(key).Return(classHasah, nil).Times(1)

proof, rpcErr := handler.StorageProof(blockLatest, nil, []felt.Felt{*key}, nil)
require.Nil(t, rpcErr)
require.NotNil(t, proof)
require.Len(t, proof.ClassesProof, 0)
require.Len(t, proof.ContractsStorageProofs, 0)
require.NotNil(t, proof.ContractsProof)
require.True(t, len(proof.ContractsProof.Nodes) > 0)
require.Len(t, proof.ContractsProof.LeavesData, 1)
arityTest(t, proof, 0, 3, 1, 0)

require.NotNil(t, proof.ContractsProof.LeavesData[0])
ld := proof.ContractsProof.LeavesData[0]
require.Equal(t, nonce, ld.Nonce)
require.Equal(t, classHasah, ld.ClassHash)

verifyIf(proof.ContractsProof.Nodes, key, value)
})
t.Run("contract storage trie address does not exist in a trie", func(t *testing.T) {
contract := utils.HexToFelt(t, "0xdead")
Expand All @@ -251,10 +267,7 @@ func TestStorageProof(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, nil, nil, storageKeys)
require.NotNil(t, proof)
require.Nil(t, rpcErr)
require.Len(t, proof.ClassesProof, 0)
require.NotNil(t, proof.ContractsProof)
require.Len(t, proof.ContractsProof.Nodes, 0)
require.Len(t, proof.ContractsStorageProofs, 1)
arityTest(t, proof, 0, 0, 0, 1)
require.Len(t, proof.ContractsStorageProofs[0], 0)
})
t.Run("contract storage trie key slot does not exist in a trie", func(t *testing.T) {
Expand All @@ -265,11 +278,10 @@ func TestStorageProof(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, nil, nil, storageKeys)
require.NotNil(t, proof)
require.Nil(t, rpcErr)
require.Len(t, proof.ClassesProof, 0)
require.NotNil(t, proof.ContractsProof)
require.Len(t, proof.ContractsProof.Nodes, 0)
require.Len(t, proof.ContractsStorageProofs, 1)
require.True(t, len(proof.ContractsStorageProofs[0]) > 0)
arityTest(t, proof, 0, 0, 0, 1)
require.Len(t, proof.ContractsStorageProofs[0], 3)

verifyIf(proof.ContractsStorageProofs[0], noSuchKey, nil)
})
t.Run("contract storage trie address/key exists in a trie", func(t *testing.T) {
contract := utils.HexToFelt(t, "0xabcd")
Expand All @@ -279,11 +291,10 @@ func TestStorageProof(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, nil, nil, storageKeys)
require.NotNil(t, proof)
require.Nil(t, rpcErr)
require.Len(t, proof.ClassesProof, 0)
require.NotNil(t, proof.ContractsProof)
require.Len(t, proof.ContractsProof.Nodes, 0)
require.Len(t, proof.ContractsStorageProofs, 1)
require.True(t, len(proof.ContractsStorageProofs[0]) > 0)
arityTest(t, proof, 0, 0, 0, 1)
require.Len(t, proof.ContractsStorageProofs[0], 3)

verifyIf(proof.ContractsStorageProofs[0], key, value)
})
t.Run("class & storage tries proofs requested", func(t *testing.T) {
nonce := new(felt.Felt).SetUint64(121)
Expand All @@ -294,14 +305,24 @@ func TestStorageProof(t *testing.T) {
proof, rpcErr := handler.StorageProof(blockLatest, []felt.Felt{*key}, []felt.Felt{*key}, nil)
require.Nil(t, rpcErr)
require.NotNil(t, proof)
require.True(t, len(proof.ClassesProof) > 0)
require.Len(t, proof.ContractsStorageProofs, 0)
require.NotNil(t, proof.ContractsProof)
require.True(t, len(proof.ContractsProof.Nodes) > 0)
require.Len(t, proof.ContractsProof.LeavesData, 1)
arityTest(t, proof, 3, 3, 1, 0)
})
}

func arityTest(t *testing.T,
proof *rpc.StorageProofResult,
classesProofArity int,
contractsProofNodesArity int,
contractsProofLeavesArity int,
contractStorageArity int,
) {
require.Len(t, proof.ClassesProof, classesProofArity)
require.Len(t, proof.ContractsStorageProofs, contractStorageArity)
require.NotNil(t, proof.ContractsProof)
require.Len(t, proof.ContractsProof.Nodes, contractsProofNodesArity)
require.Len(t, proof.ContractsProof.LeavesData, contractsProofLeavesArity)
}

func emptyTrie(t *testing.T) *trie.Trie {
memdb := pebble.NewMemTest(t)
txn, err := memdb.NewTransaction(true)
Expand Down

0 comments on commit fcf4395

Please sign in to comment.