Skip to content

Commit

Permalink
getStorageProof rpc method
Browse files Browse the repository at this point in the history
  • Loading branch information
weiihann authored and pnowosie committed Dec 17, 2024
1 parent 71c7ae9 commit 4b4b538
Show file tree
Hide file tree
Showing 18 changed files with 1,595 additions and 136 deletions.
12 changes: 12 additions & 0 deletions blockchain/blockchain.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ type Reader interface {
L1HandlerTxnHash(msgHash *common.Hash) (l1HandlerTxnHash *felt.Felt, err error)

HeadState() (core.StateReader, StateCloser, error)
HeadTrie() (core.TrieReader, StateCloser, error)
StateAtBlockHash(blockHash *felt.Felt) (core.StateReader, StateCloser, error)
StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)

Expand Down Expand Up @@ -768,6 +769,17 @@ func (b *Blockchain) HeadState() (core.StateReader, StateCloser, error) {
return core.NewState(txn), txn.Discard, nil
}

func (b *Blockchain) HeadTrie() (core.TrieReader, StateCloser, error) {
// Note: I'm not sure I should open a new db txn since the TrieReader is a State
// so the same instance of the state we create in HeadState will do job.
txn, err := b.database.NewTransaction(false)
if err != nil {
return nil, nil, err
}

Check warning on line 778 in blockchain/blockchain.go

View check run for this annotation

Codecov / codecov/patch

blockchain/blockchain.go#L777-L778

Added lines #L777 - L778 were not covered by tests

return core.NewState(txn), txn.Discard, nil
}

// StateAtBlockNumber returns a StateReader that provides a stable view to the state at the given block number
func (b *Blockchain) StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error) {
b.listener.OnRead("StateAtBlockNumber")
Expand Down
58 changes: 57 additions & 1 deletion core/state.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,17 @@ type StateReader interface {
Class(classHash *felt.Felt) (*DeclaredClass, error)
}

// TrieReader used for storage proofs, can only be supported by current state implementation (for now, we plan to add db snapshots)
var _ TrieReader = (*State)(nil)

//go:generate mockgen -destination=../mocks/mock_trie.go -package=mocks github.com/NethermindEth/juno/core TrieReader
type TrieReader interface {
ClassTrie() (*trie.Trie, func() error, error)
StorageTrie() (*trie.Trie, func() error, error)
StorageTrieForAddr(addr *felt.Felt) (*trie.Trie, error)
StateAndClassRoot() (*felt.Felt, *felt.Felt, error)
}

type State struct {
*history
txn db.Transaction
Expand Down Expand Up @@ -129,6 +140,18 @@ func (s *State) storage() (*trie.Trie, func() error, error) {
return s.globalTrie(db.StateTrie, trie.NewTriePedersen)
}

func (s *State) StorageTrie() (*trie.Trie, func() error, error) {
return s.storage()
}

func (s *State) ClassTrie() (*trie.Trie, func() error, error) {
return s.classesTrie()
}

func (s *State) StorageTrieForAddr(addr *felt.Felt) (*trie.Trie, error) {
return storage(addr, s.txn)

Check warning on line 152 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L151-L152

Added lines #L151 - L152 were not covered by tests
}

func (s *State) classesTrie() (*trie.Trie, func() error, error) {
return s.globalTrie(db.ClassesTrie, trie.NewTriePoseidon)
}
Expand Down Expand Up @@ -547,7 +570,7 @@ func (s *State) Revert(blockNumber uint64, update *StateUpdate) error {

err = s.performStateDeletions(blockNumber, update.StateDiff)
if err != nil {
return fmt.Errorf("error performing state deletions: %v", err)
return fmt.Errorf("build reverse diff: %v", err)

Check warning on line 573 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L573

Added line #L573 was not covered by tests
}

stateTrie, storageCloser, err := s.storage()
Expand Down Expand Up @@ -581,6 +604,7 @@ func (s *State) purgeNoClassContracts() error {
// As noClassContracts are not in StateDiff.DeployedContracts we can only purge them if their storage no longer exists.
// Updating contracts with reverse diff will eventually lead to the deletion of noClassContract's storage key from db. Thus,
// we can use the lack of key's existence as reason for purging noClassContracts.

for addr := range noClassContracts {
noClassC, err := NewContractUpdater(&addr, s.txn)
if err != nil {
Expand Down Expand Up @@ -743,3 +767,35 @@ func (s *State) performStateDeletions(blockNumber uint64, diff *StateDiff) error

return nil
}

func (s *State) StateAndClassRoot() (*felt.Felt, *felt.Felt, error) {
var storageRoot, classesRoot *felt.Felt

sStorage, closer, err := s.storage()
if err != nil {
return nil, nil, err
}

Check warning on line 777 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L776-L777

Added lines #L776 - L777 were not covered by tests

if storageRoot, err = sStorage.Root(); err != nil {
return nil, nil, err
}

Check warning on line 781 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L780-L781

Added lines #L780 - L781 were not covered by tests

if err = closer(); err != nil {
return nil, nil, err
}

Check warning on line 785 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L784-L785

Added lines #L784 - L785 were not covered by tests

classes, closer, err := s.classesTrie()
if err != nil {
return nil, nil, err
}

Check warning on line 790 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L789-L790

Added lines #L789 - L790 were not covered by tests

if classesRoot, err = classes.Root(); err != nil {
return nil, nil, err
}

Check warning on line 794 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L793-L794

Added lines #L793 - L794 were not covered by tests

if err = closer(); err != nil {
return nil, nil, err
}

Check warning on line 798 in core/state.go

View check run for this annotation

Codecov / codecov/patch

core/state.go#L797-L798

Added lines #L797 - L798 were not covered by tests

return storageRoot, classesRoot, nil
}
6 changes: 3 additions & 3 deletions core/trie/key.go
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
package trie

import (
"bytes"
"encoding/hex"
"errors"
"fmt"
"io"
"math/big"

"github.com/NethermindEth/juno/core/felt"
Expand Down Expand Up @@ -39,8 +39,8 @@ func (k *Key) unusedBytes() []byte {
return k.bitset[:len(k.bitset)-int(k.bytesNeeded())]
}

func (k *Key) WriteTo(buf *bytes.Buffer) (int64, error) {
if err := buf.WriteByte(k.len); err != nil {
func (k *Key) WriteTo(buf io.Writer) (int64, error) {
if _, err := buf.Write([]byte{k.len}); err != nil {
return 0, err
}

Expand Down
37 changes: 37 additions & 0 deletions core/trie/key_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package trie_test

import (
"bytes"
"errors"
"testing"

"github.com/NethermindEth/juno/core/felt"
Expand Down Expand Up @@ -227,3 +228,39 @@ func TestMostSignificantBits(t *testing.T) {
})
}
}

func TestKeyErrorHandling(t *testing.T) {
t.Run("passed too long key bytes panics", func(t *testing.T) {
defer func() {
r := recover()
require.NotNil(t, r)
require.Contains(t, r.(string), "bytes does not fit in bitset")
}()
tooLongKeyB := make([]byte, 33)
trie.NewKey(8, tooLongKeyB)
})
t.Run("MostSignificantBits n greater than key length", func(t *testing.T) {
key := trie.NewKey(8, []byte{0x01})
_, err := key.MostSignificantBits(9)
require.Error(t, err)
require.Contains(t, err.Error(), "cannot get more bits than the key length")
})
t.Run("MostSignificantBits equals key length return copy of key", func(t *testing.T) {
key := trie.NewKey(8, []byte{0x01})
kCopy, err := key.MostSignificantBits(8)
require.NoError(t, err)
require.Equal(t, key, *kCopy)
})
t.Run("WriteTo returns error", func(t *testing.T) {
key := trie.NewKey(8, []byte{0x01})
wrote, err := key.WriteTo(&errorBuffer{})
require.Error(t, err)
require.Equal(t, int64(0), wrote)
})
}

type errorBuffer struct{}

func (*errorBuffer) Write([]byte) (int, error) {
return 0, errors.New("expected to fail")
}
8 changes: 4 additions & 4 deletions core/trie/node.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
package trie

import (
"bytes"
"errors"
"fmt"
"io"

"github.com/NethermindEth/juno/core/felt"
)
Expand All @@ -19,7 +19,7 @@ type Node struct {
}

// Hash calculates the hash of a [Node]
func (n *Node) Hash(path *Key, hashFunc hashFunc) *felt.Felt {
func (n *Node) Hash(path *Key, hashFunc HashFunc) *felt.Felt {
if path.Len() == 0 {
// we have to deference the Value, since the Node can released back
// to the NodePool and be reused anytime
Expand All @@ -34,12 +34,12 @@ func (n *Node) Hash(path *Key, hashFunc hashFunc) *felt.Felt {
}

// Hash calculates the hash of a [Node]
func (n *Node) HashFromParent(parentKey, nodeKey *Key, hashFunc hashFunc) *felt.Felt {
func (n *Node) HashFromParent(parentKey, nodeKey *Key, hashFunc HashFunc) *felt.Felt {

Check warning on line 37 in core/trie/node.go

View check run for this annotation

Codecov / codecov/patch

core/trie/node.go#L37

Added line #L37 was not covered by tests
path := path(nodeKey, parentKey)
return n.Hash(&path, hashFunc)
}

func (n *Node) WriteTo(buf *bytes.Buffer) (int64, error) {
func (n *Node) WriteTo(buf io.Writer) (int64, error) {
if n.Value == nil {
return 0, errors.New("cannot marshal node with nil value")
}
Expand Down
32 changes: 32 additions & 0 deletions core/trie/node_test.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
package trie_test

import (
"bytes"
"encoding/hex"
"errors"
"testing"

"github.com/NethermindEth/juno/core/crypto"
Expand All @@ -26,3 +28,33 @@ func TestNodeHash(t *testing.T) {

assert.Equal(t, expected, node.Hash(&path, crypto.Pedersen), "TestTrieNode_Hash failed")
}

func TestNodeErrorHandling(t *testing.T) {
t.Run("WriteTo node value is nil", func(t *testing.T) {
node := trie.Node{}
var buffer bytes.Buffer
_, err := node.WriteTo(&buffer)
require.Error(t, err)
})
t.Run("WriteTo returns error", func(t *testing.T) {
node := trie.Node{
Value: new(felt.Felt).SetUint64(42),
Left: &trie.Key{},
Right: &trie.Key{},
}

wrote, err := node.WriteTo(&errorBuffer{})
require.Error(t, err)
require.Equal(t, int64(0), wrote)
})
t.Run("UnmarshalBinary returns error", func(t *testing.T) {
node := trie.Node{}

err := node.UnmarshalBinary([]byte{42})
require.Equal(t, errors.New("size of input data is less than felt size"), err)

bs := new(felt.Felt).Bytes()
err = node.UnmarshalBinary(append(bs[:], 0, 0, 42))
require.Equal(t, errors.New("the node does not contain both left and right hash"), err)
})
}
8 changes: 4 additions & 4 deletions core/trie/proof.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ func NewProofNodeSet() *ProofNodeSet {
}

type ProofNode interface {
Hash(hash hashFunc) *felt.Felt
Hash(hash HashFunc) *felt.Felt
Len() uint8
String() string
}
Expand All @@ -26,7 +26,7 @@ type Binary struct {
RightHash *felt.Felt
}

func (b *Binary) Hash(hash hashFunc) *felt.Felt {
func (b *Binary) Hash(hash HashFunc) *felt.Felt {
return hash(b.LeftHash, b.RightHash)
}

Expand All @@ -43,7 +43,7 @@ type Edge struct {
Path *Key // path from parent to child
}

func (e *Edge) Hash(hash hashFunc) *felt.Felt {
func (e *Edge) Hash(hash HashFunc) *felt.Felt {
length := make([]byte, len(e.Path.bitset))
length[len(e.Path.bitset)-1] = e.Path.len
pathFelt := e.Path.Felt()
Expand Down Expand Up @@ -137,7 +137,7 @@ func (t *Trie) GetRangeProof(leftKey, rightKey *felt.Felt, proofSet *ProofNodeSe
// - Any node's computed hash doesn't match its expected hash
// - The path bits don't match the key bits
// - The proof ends before processing all key bits
func VerifyProof(root, keyFelt *felt.Felt, proof *ProofNodeSet, hash hashFunc) (*felt.Felt, error) {
func VerifyProof(root, keyFelt *felt.Felt, proof *ProofNodeSet, hash HashFunc) (*felt.Felt, error) {
key := FeltToKey(globalTrieHeight, keyFelt)
expectedHash := root
keyLen := key.Len()
Expand Down
49 changes: 49 additions & 0 deletions core/trie/proofset.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
package trie

import (
"sync"

"github.com/NethermindEth/juno/core/felt"
)

// ProofSet represents a set of trie nodes used in a Merkle proof verification process.
// Rather than relying on only either map of list, ProofSet provides both for the following reasons:
// - map allows for unique node insertion
// - list allows for ordered iteration over the proof nodes
// It also supports concurrent read and write operations.
type ProofSet struct {
nodeSet map[felt.Felt]ProofNode
nodeList []ProofNode
size int
lock sync.RWMutex
}

func NewProofSet() *ProofSet {
return &ProofSet{
nodeSet: make(map[felt.Felt]ProofNode),
}

Check warning on line 24 in core/trie/proofset.go

View check run for this annotation

Codecov / codecov/patch

core/trie/proofset.go#L21-L24

Added lines #L21 - L24 were not covered by tests
}

func (ps *ProofSet) Put(key felt.Felt, node ProofNode) {
ps.lock.Lock()
defer ps.lock.Unlock()

ps.nodeSet[key] = node
ps.nodeList = append(ps.nodeList, node)
ps.size++

Check warning on line 33 in core/trie/proofset.go

View check run for this annotation

Codecov / codecov/patch

core/trie/proofset.go#L27-L33

Added lines #L27 - L33 were not covered by tests
}

func (ps *ProofSet) Get(key felt.Felt) (ProofNode, bool) {
ps.lock.RLock()
defer ps.lock.RUnlock()

node, ok := ps.nodeSet[key]
return node, ok

Check warning on line 41 in core/trie/proofset.go

View check run for this annotation

Codecov / codecov/patch

core/trie/proofset.go#L36-L41

Added lines #L36 - L41 were not covered by tests
}

func (ps *ProofSet) Size() int {
ps.lock.RLock()
defer ps.lock.RUnlock()

return ps.size

Check warning on line 48 in core/trie/proofset.go

View check run for this annotation

Codecov / codecov/patch

core/trie/proofset.go#L44-L48

Added lines #L44 - L48 were not covered by tests
}
13 changes: 9 additions & 4 deletions core/trie/trie.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ import (

const globalTrieHeight = 251 // TODO(weiihann): this is declared in core also, should be moved to a common place

type hashFunc func(*felt.Felt, *felt.Felt) *felt.Felt
type HashFunc func(*felt.Felt, *felt.Felt) *felt.Felt

// Trie is a dense Merkle Patricia Trie (i.e., all internal nodes have two children).
//
Expand All @@ -40,7 +40,7 @@ type Trie struct {
rootKey *Key
maxKey *felt.Felt
storage *Storage
hash hashFunc
hash HashFunc

dirtyNodes []*Key
rootKeyIsDirty bool
Expand All @@ -56,7 +56,7 @@ func NewTriePoseidon(storage *Storage, height uint8) (*Trie, error) {
return newTrie(storage, height, crypto.Poseidon)
}

func newTrie(storage *Storage, height uint8, hash hashFunc) (*Trie, error) {
func newTrie(storage *Storage, height uint8, hash HashFunc) (*Trie, error) {
if height > felt.Bits {
return nil, fmt.Errorf("max trie height is %d, got: %d", felt.Bits, height)
}
Expand Down Expand Up @@ -96,12 +96,17 @@ func RunOnTempTriePoseidon(height uint8, do func(*Trie) error) error {
return do(trie)
}

// feltToKey Converts a key, given in felt, to a trie.Key which when followed on a [Trie],
// FeltToKey Converts a key, given in felt, to a trie.Key which when followed on a [Trie],
// leads to the corresponding [Node]
func (t *Trie) FeltToKey(k *felt.Felt) Key {
return FeltToKey(t.height, k)
}

// HashFunc returns the hash function used by the trie
func (t *Trie) HashFunc() HashFunc {
return t.hash
}

// path returns the path as mentioned in the [specification] for commitment calculations.
// path is suffix of key that diverges from parentKey. For example,
// for a key 0b1011 and parentKey 0b10, this function would return the path object of 0b0.
Expand Down
Loading

0 comments on commit 4b4b538

Please sign in to comment.