fix(rpc): avoid race when reading storage proof state

blockchain/blockchain.go

@@ -85,6 +85,7 @@ type Reader interface {
 	HeadState() (core.StateReader, StateCloser, error)
 	StateAtBlockHash(blockHash *felt.Felt) (core.StateReader, StateCloser, error)
 	StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)
+	TrieStateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)
 
 	BlockCommitmentsByNumber(blockNumber uint64) (*core.BlockCommitments, error)
 
@@ -464,6 +465,14 @@ func (b *Blockchain) StateAtBlockHash(
 	return b.stateBackend.StateAtBlockHash(blockHash)
 }
 
+// TrieStateAtBlockNumber returns a StateReader with trie access at the given block number.
+func (b *Blockchain) TrieStateAtBlockNumber(
+	blockNumber uint64,
+) (core.StateReader, StateCloser, error) {
+	b.listener.OnRead("TrieStateAtBlockNumber")
+	return b.stateBackend.TrieStateAtBlockNumber(blockNumber)
+}
+
 // EventFilter returns an EventFilter object that is tied to a snapshot of the blockchain
 func (b *Blockchain) EventFilter(
 	addresses []felt.Address,

blockchain/statebackend/deprecated.go

@@ -63,6 +63,30 @@
 	), NoopStateCloser, nil
 }
 
+func (b *deprecatedStateBackend) TrieStateAtBlockNumber(
+	blockNumber uint64,
+) (core.StateReader, StateCloser, error) {
+	if err := pruner.RequireStateRetainedByBlockNumber(
+		b.database,
+		b.retentionFloor,
+		blockNumber,
+	); err != nil {
+		return nil, nil, err
+	}
+
+	//nolint:staticcheck,nolintlint // used by old state
+	txn := b.database.NewIndexedBatch()
+	height, err := core.GetChainHeight(txn)
+	if err != nil {
+		return nil, nil, err
+	}
+	if height != blockNumber {
+		return nil, nil, db.ErrKeyNotFound
+	}
+
+	return deprecatedstate.New(txn), NoopStateCloser, nil
+}
+
 func (b *deprecatedStateBackend) Store(
 	block *core.Block,
 	blockCommitments *core.BlockCommitments,

blockchain/statebackend/statebackend.go

@@ -77,6 +77,25 @@
 	return &history, NoopStateCloser, nil
 }
 
+func (b *stateBackend) TrieStateAtBlockNumber(
+	blockNumber uint64,
+) (core.StateReader, StateCloser, error) {
+	stateRoot, err := pruner.StateRootIfStateRetainedByBlockNumber(
+		b.database,
+		b.retentionFloor,
+		blockNumber,
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	st, err := state.NewStateReader(stateRoot, b.stateDB)
+	if err != nil {
+		return nil, nil, err
+	}
+	return st, NoopStateCloser, nil
+}
+
 func (b *stateBackend) Store(
 	block *core.Block,
 	blockCommitments *core.BlockCommitments,

blockchain/statebackend/types.go

@@ -15,6 +15,7 @@ type StateBackend interface {
 	HeadState() (core.StateReader, StateCloser, error)
 	StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)
 	StateAtBlockHash(blockHash *felt.Felt) (core.StateReader, StateCloser, error)
+	TrieStateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)
 	Store(
 		block *core.Block,
 		commitments *core.BlockCommitments,

rpc/v10/storage.go

@@ -206,7 +206,7 @@ func (h *Handler) StorageProof(
 		return nil, rpccore.ErrInternal.CloneWithData(err)
 	}
 
-	state, closer, err := h.bcReader.HeadState()
+	state, closer, err := h.bcReader.TrieStateAtBlockNumber(chainHeight)
 	if err != nil {
 		return nil, rpccore.ErrInternal.CloneWithData(err)
 	}

rpc/v10/storage_bench_test.go

@@ -150,7 +150,10 @@ func (f *storageProofBenchmarkFixture) setup(
 
 	mockReader.EXPECT().Height().Return(f.blockNumber, nil).AnyTimes()
 	mockReader.EXPECT().BlockHeaderHashByNumber(f.blockNumber).Return(f.blkHash, nil).AnyTimes()
-	mockReader.EXPECT().HeadState().Return(mockState, func() error { return nil }, nil).AnyTimes()
+	mockReader.EXPECT().
+		TrieStateAtBlockNumber(f.blockNumber).
+		Return(mockState, func() error { return nil }, nil).
+		AnyTimes()
 
 	mockState.EXPECT().ClassTrie().Return(classTrie, nil).AnyTimes()
 	mockState.EXPECT().ContractTrie().Return(contractTrie, nil).AnyTimes()

rpc/v10/storage_test.go

@@ -694,7 +694,10 @@ func TestStorageProof(t *testing.T) {
 
 	mockReader := mocks.NewMockReader(mockCtrl)
 	mockState := mocks.NewMockStateReader(mockCtrl)
-	mockReader.EXPECT().HeadState().Return(mockState, func() error { return nil }, nil).AnyTimes()
+	mockReader.EXPECT().
+		TrieStateAtBlockNumber(blockNumber).
+		Return(mockState, func() error { return nil }, nil).
+		AnyTimes()
 	mockReader.EXPECT().Head().Return(headBlock, nil).AnyTimes()
 	mockReader.EXPECT().BlockByNumber(blockNumber).Return(headBlock, nil).AnyTimes()
 	mockReader.EXPECT().Height().Return(blockNumber, nil).AnyTimes()

rpc/v9/storage.go

@@ -93,7 +93,7 @@ func (h *Handler) StorageProof(
 		return nil, rpccore.ErrInternal.CloneWithData(err)
 	}
 
-	state, closer, err := h.bcReader.HeadState()
+	state, closer, err := h.bcReader.TrieStateAtBlockNumber(chainHeight)
 	if err != nil {
 		return nil, rpccore.ErrInternal.CloneWithData(err)
 	}

rpc/v9/storage_test.go

@@ -351,7 +351,10 @@ func TestStorageProof(t *testing.T) {
 
 	mockReader := mocks.NewMockReader(mockCtrl)
 	mockState := mocks.NewMockStateReader(mockCtrl)
-	mockReader.EXPECT().HeadState().Return(mockState, func() error { return nil }, nil).AnyTimes()
+	mockReader.EXPECT().
+		TrieStateAtBlockNumber(blockNumber).
+		Return(mockState, func() error { return nil }, nil).
+		AnyTimes()
 	mockReader.EXPECT().Head().Return(headBlock, nil).AnyTimes()
 	mockReader.EXPECT().BlockByNumber(blockNumber).Return(headBlock, nil).AnyTimes()
 	mockReader.EXPECT().Height().Return(blockNumber, nil).AnyTimes()