Starknet's getStorageProof rpc method

blockchain/blockchain.go

@@ -36,6 +36,7 @@
 	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)
 
@@ -768,6 +769,17 @@
 	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
+	}
+
+	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")

core/state.go

@@ -44,6 +44,17 @@
 	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
@@ -129,6 +140,18 @@
 	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)
+}
+
 func (s *State) classesTrie() (*trie.Trie, func() error, error) {
 	return s.globalTrie(db.ClassesTrie, trie.NewTriePoseidon)
 }
@@ -547,7 +570,7 @@
 
 	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)
 	}
 
 	stateTrie, storageCloser, err := s.storage()
@@ -581,6 +604,7 @@
 	// 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 {
@@ -743,3 +767,35 @@
 
 	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
+	}
+
+	if storageRoot, err = sStorage.Root(); err != nil {
+		return nil, nil, err
+	}
+
+	if err = closer(); err != nil {
+		return nil, nil, err
+	}
+
+	classes, closer, err := s.classesTrie()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if classesRoot, err = classes.Root(); err != nil {
+		return nil, nil, err
+	}
+
+	if err = closer(); err != nil {
+		return nil, nil, err
+	}
+
+	return storageRoot, classesRoot, nil
+}

core/trie/key.go

@@ -1,10 +1,10 @@
 package trie
 
 import (
-	"bytes"
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"io"
 	"math/big"
 
 	"github.com/NethermindEth/juno/core/felt"
@@ -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
 	}
 

core/trie/key_test.go

@@ -2,6 +2,7 @@ package trie_test
 
 import (
 	"bytes"
+	"errors"
 	"testing"
 
 	"github.com/NethermindEth/juno/core/felt"
@@ -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")
+}

core/trie/node.go

@@ -1,9 +1,9 @@
 package trie
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
+	"io"
 
 	"github.com/NethermindEth/juno/core/felt"
 )
@@ -19,7 +19,7 @@
 }
 
 // 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
@@ -34,12 +34,12 @@
 }
 
 // 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 {
 	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")
 	}

core/trie/node_test.go

@@ -1,7 +1,9 @@
 package trie_test
 
 import (
+	"bytes"
 	"encoding/hex"
+	"errors"
 	"testing"
 
 	"github.com/NethermindEth/juno/core/crypto"
@@ -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)
+	})
+}

core/trie/proof.go

@@ -16,7 +16,7 @@ func NewProofNodeSet() *ProofNodeSet {
 }
 
 type ProofNode interface {
-	Hash(hash hashFunc) *felt.Felt
+	Hash(hash HashFunc) *felt.Felt
 	Len() uint8
 	String() string
 }
@@ -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)
 }
 
@@ -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()
@@ -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()

core/trie/proofset.go

@@ -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),
+	}
+}
+
+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++
+}
+
+func (ps *ProofSet) Get(key felt.Felt) (ProofNode, bool) {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
+
+	node, ok := ps.nodeSet[key]
+	return node, ok
+}
+
+func (ps *ProofSet) Size() int {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
+
+	return ps.size
+}

core/trie/trie.go

@@ -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).
 //
@@ -40,7 +40,7 @@ type Trie struct {
 	rootKey *Key
 	maxKey  *felt.Felt
 	storage *Storage
-	hash    hashFunc
+	hash    HashFunc
 
 	dirtyNodes     []*Key
 	rootKeyIsDirty bool
@@ -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)
 	}
@@ -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.