Initialize Merkle Layers and Recompute Dirty Fields in BeaconState Proofs (#10032)

* handle trie recomputes

* recomp

* state trie

* args

* deep source

* amend test

* proofs test more fixes

* tests pass

beacon-chain/state/phase0.go

@@ -24,9 +24,9 @@ type BeaconState interface {
 
 // StateProver defines the ability to create Merkle proofs for beacon state fields.
 type StateProver interface {
-	FinalizedRootProof() ([][]byte, error)
-	CurrentSyncCommitteeProof() ([][]byte, error)
-	NextSyncCommitteeProof() ([][]byte, error)
+	FinalizedRootProof(ctx context.Context) ([][]byte, error)
+	CurrentSyncCommitteeProof(ctx context.Context) ([][]byte, error)
+	NextSyncCommitteeProof(ctx context.Context) ([][]byte, error)
 }
 
 // ReadOnlyBeaconState defines a struct which only has read access to beacon state methods.

beacon-chain/state/v1/proofs.go

@@ -1,6 +1,7 @@
 package v1
 
 import (
+	"context"
 	"encoding/binary"
 
 	"github.com/pkg/errors"
@@ -18,20 +19,26 @@ func FinalizedRootGeneralizedIndex() uint64 {
 }
 
 // CurrentSyncCommitteeProof from the state's Merkle trie representation.
-func (*BeaconState) CurrentSyncCommitteeProof() ([][]byte, error) {
+func (*BeaconState) CurrentSyncCommitteeProof(_ context.Context) ([][]byte, error) {
 	return nil, errors.New("CurrentSyncCommitteeProof() unsupported for v1 beacon state")
 }
 
 // NextSyncCommitteeProof from the state's Merkle trie representation.
-func (*BeaconState) NextSyncCommitteeProof() ([][]byte, error) {
+func (*BeaconState) NextSyncCommitteeProof(_ context.Context) ([][]byte, error) {
 	return nil, errors.New("NextSyncCommitteeProof() unsupported for v1 beacon state")
 }
 
 // FinalizedRootProof crafts a Merkle proof for the finalized root
 // contained within the finalized checkpoint of a beacon state.
-func (b *BeaconState) FinalizedRootProof() ([][]byte, error) {
-	b.lock.RLock()
-	defer b.lock.RUnlock()
+func (b *BeaconState) FinalizedRootProof(ctx context.Context) ([][]byte, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return nil, err
+	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return nil, err
+	}
 	cpt := b.state.FinalizedCheckpoint
 	// The epoch field of a finalized checkpoint is the neighbor
 	// index of the finalized root field in its Merkle tree representation

beacon-chain/state/v1/proofs_test.go

@@ -16,19 +16,49 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 	htr, err := st.HashTreeRoot(ctx)
 	require.NoError(t, err)
 	t.Run("current sync committee", func(t *testing.T) {
-		_, err := st.CurrentSyncCommitteeProof()
+		_, err := st.CurrentSyncCommitteeProof(ctx)
 		require.ErrorContains(t, "unsupported", err)
 	})
 	t.Run("next sync committee", func(t *testing.T) {
-		_, err := st.NextSyncCommitteeProof()
+		_, err := st.NextSyncCommitteeProof(ctx)
 		require.ErrorContains(t, "unsupported", err)
 	})
 	t.Run("finalized root", func(t *testing.T) {
 		finalizedRoot := st.FinalizedCheckpoint().Root
-		proof, err := st.FinalizedRootProof()
+		proof, err := st.FinalizedRootProof(ctx)
 		require.NoError(t, err)
 		gIndex := v1.FinalizedRootGeneralizedIndex()
 		valid := trie.VerifyMerkleProof(htr[:], finalizedRoot, gIndex, proof)
 		require.Equal(t, true, valid)
 	})
+	t.Run("recomputes root on dirty fields", func(t *testing.T) {
+		currentRoot, err := st.HashTreeRoot(ctx)
+		require.NoError(t, err)
+		cpt := st.FinalizedCheckpoint()
+		require.NoError(t, err)
+
+		// Edit the checkpoint.
+		cpt.Epoch = 100
+		require.NoError(t, st.SetFinalizedCheckpoint(cpt))
+
+		// Produce a proof for the finalized root.
+		proof, err := st.FinalizedRootProof(ctx)
+		require.NoError(t, err)
+
+		// We expect the previous step to have triggered
+		// a recomputation of dirty fields in the beacon state, resulting
+		// in a new hash tree root as the finalized checkpoint had previously
+		// changed and should have been marked as a dirty state field.
+		// The proof validity should be false for the old root, but true for the new.
+		finalizedRoot := st.FinalizedCheckpoint().Root
+		gIndex := v1.FinalizedRootGeneralizedIndex()
+		valid := trie.VerifyMerkleProof(currentRoot[:], finalizedRoot, gIndex, proof)
+		require.Equal(t, false, valid)
+
+		newRoot, err := st.HashTreeRoot(ctx)
+		require.NoError(t, err)
+
+		valid = trie.VerifyMerkleProof(newRoot[:], finalizedRoot, gIndex, proof)
+		require.Equal(t, true, valid)
+	})
 }

beacon-chain/state/v1/state_trie.go

@@ -204,27 +204,44 @@ func (b *BeaconState) HashTreeRoot(ctx context.Context) ([32]byte, error) {
 
 	b.lock.Lock()
 	defer b.lock.Unlock()
-
-	if b.merkleLayers == nil || len(b.merkleLayers) == 0 {
-		fieldRoots, err := computeFieldRoots(ctx, b.state)
-		if err != nil {
-			return [32]byte{}, err
-		}
-		layers := stateutil.Merkleize(fieldRoots)
-		b.merkleLayers = layers
-		b.dirtyFields = make(map[types.FieldIndex]bool, params.BeaconConfig().BeaconStateFieldCount)
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return [32]byte{}, err
 	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return [32]byte{}, err
+	}
+	return bytesutil.ToBytes32(b.merkleLayers[len(b.merkleLayers)-1][0]), nil
+}
 
+// Initializes the Merkle layers for the beacon state if they are empty.
+// WARNING: Caller must acquire the mutex before using.
+func (b *BeaconState) initializeMerkleLayers(ctx context.Context) error {
+	if len(b.merkleLayers) > 0 {
+		return nil
+	}
+	fieldRoots, err := computeFieldRoots(ctx, b.state)
+	if err != nil {
+		return err
+	}
+	layers := stateutil.Merkleize(fieldRoots)
+	b.merkleLayers = layers
+	b.dirtyFields = make(map[types.FieldIndex]bool, params.BeaconConfig().BeaconStateFieldCount)
+	return nil
+}
+
+// Recomputes the Merkle layers for the dirty fields in the state.
+// WARNING: Caller must acquire the mutex before using.
+func (b *BeaconState) recomputeDirtyFields(ctx context.Context) error {
 	for field := range b.dirtyFields {
 		root, err := b.rootSelector(ctx, field)
 		if err != nil {
-			return [32]byte{}, err
+			return err
 		}
 		b.merkleLayers[0][field] = root[:]
 		b.recomputeRoot(int(field))
 		delete(b.dirtyFields, field)
 	}
-	return bytesutil.ToBytes32(b.merkleLayers[len(b.merkleLayers)-1][0]), nil
+	return nil
 }
 
 // FieldReferencesCount returns the reference count held by each field. This

beacon-chain/state/v2/proofs.go

@@ -1,6 +1,7 @@
 package v2
 
 import (
+	"context"
 	"encoding/binary"
 
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/fieldtrie"
@@ -27,24 +28,46 @@ func NextSyncCommitteeGeneralizedIndex() uint64 {
 }
 
 // CurrentSyncCommitteeProof from the state's Merkle trie representation.
-func (b *BeaconState) CurrentSyncCommitteeProof() ([][]byte, error) {
-	b.lock.RLock()
-	defer b.lock.RUnlock()
+func (b *BeaconState) CurrentSyncCommitteeProof(ctx context.Context) ([][]byte, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	// In case the Merkle layers of the trie are not populated, we need
+	// to perform some initialization.
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return nil, err
+	}
+	// Our beacon state uses a "dirty" fields pattern which requires us to
+	// recompute branches of the Merkle layers that are marked as dirty.
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return nil, err
+	}
 	return fieldtrie.ProofFromMerkleLayers(b.merkleLayers, currentSyncCommittee), nil
 }
 
 // NextSyncCommitteeProof from the state's Merkle trie representation.
-func (b *BeaconState) NextSyncCommitteeProof() ([][]byte, error) {
-	b.lock.RLock()
-	defer b.lock.RUnlock()
+func (b *BeaconState) NextSyncCommitteeProof(ctx context.Context) ([][]byte, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return nil, err
+	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return nil, err
+	}
 	return fieldtrie.ProofFromMerkleLayers(b.merkleLayers, nextSyncCommittee), nil
 }
 
 // FinalizedRootProof crafts a Merkle proof for the finalized root
 // contained within the finalized checkpoint of a beacon state.
-func (b *BeaconState) FinalizedRootProof() ([][]byte, error) {
-	b.lock.RLock()
-	defer b.lock.RUnlock()
+func (b *BeaconState) FinalizedRootProof(ctx context.Context) ([][]byte, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return nil, err
+	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return nil, err
+	}
 	cpt := b.state.FinalizedCheckpoint
 	// The epoch field of a finalized checkpoint is the neighbor
 	// index of the finalized root field in its Merkle tree representation

beacon-chain/state/v2/proofs_test.go

@@ -23,7 +23,7 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 		// Verify the Merkle proof.
 		scRoot, err := sc.HashTreeRoot()
 		require.NoError(t, err)
-		proof, err := st.CurrentSyncCommitteeProof()
+		proof, err := st.CurrentSyncCommitteeProof(ctx)
 		require.NoError(t, err)
 		valid := trie.VerifyMerkleProof(htr[:], scRoot[:], v2.CurrentSyncCommitteeGeneralizedIndex(), proof)
 		require.Equal(t, true, valid)
@@ -31,7 +31,7 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 	t.Run("next sync committee", func(t *testing.T) {
 		nextSC, err := st.NextSyncCommittee()
 		require.NoError(t, err)
-		proof, err := st.NextSyncCommitteeProof()
+		proof, err := st.NextSyncCommitteeProof(ctx)
 		require.NoError(t, err)
 
 		// Verify the Merkle proof.
@@ -53,7 +53,7 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 		require.Equal(t, false, valid)
 
 		// Generating a new, valid proof should pass.
-		proof, err = st.NextSyncCommitteeProof()
+		proof, err = st.NextSyncCommitteeProof(ctx)
 		require.NoError(t, err)
 		htr, err = st.HashTreeRoot(ctx)
 		require.NoError(t, err)
@@ -66,10 +66,40 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 		// Verify the Merkle proof.
 		htr, err = st.HashTreeRoot(ctx)
 		require.NoError(t, err)
-		proof, err := st.FinalizedRootProof()
+		proof, err := st.FinalizedRootProof(ctx)
 		require.NoError(t, err)
 		gIndex := v2.FinalizedRootGeneralizedIndex()
 		valid := trie.VerifyMerkleProof(htr[:], finalizedRoot, gIndex, proof)
 		require.Equal(t, true, valid)
 	})
+	t.Run("recomputes root on dirty fields", func(t *testing.T) {
+		currentRoot, err := st.HashTreeRoot(ctx)
+		require.NoError(t, err)
+		cpt := st.FinalizedCheckpoint()
+		require.NoError(t, err)
+
+		// Edit the checkpoint.
+		cpt.Epoch = 100
+		require.NoError(t, st.SetFinalizedCheckpoint(cpt))
+
+		// Produce a proof for the finalized root.
+		proof, err := st.FinalizedRootProof(ctx)
+		require.NoError(t, err)
+
+		// We expect the previous step to have triggered
+		// a recomputation of dirty fields in the beacon state, resulting
+		// in a new hash tree root as the finalized checkpoint had previously
+		// changed and should have been marked as a dirty state field.
+		// The proof validity should be false for the old root, but true for the new.
+		finalizedRoot := st.FinalizedCheckpoint().Root
+		gIndex := v2.FinalizedRootGeneralizedIndex()
+		valid := trie.VerifyMerkleProof(currentRoot[:], finalizedRoot, gIndex, proof)
+		require.Equal(t, false, valid)
+
+		newRoot, err := st.HashTreeRoot(ctx)
+		require.NoError(t, err)
+
+		valid = trie.VerifyMerkleProof(newRoot[:], finalizedRoot, gIndex, proof)
+		require.Equal(t, true, valid)
+	})
 }

beacon-chain/state/v2/state_trie.go

@@ -232,27 +232,44 @@ func (b *BeaconState) HashTreeRoot(ctx context.Context) ([32]byte, error) {
 
 	b.lock.Lock()
 	defer b.lock.Unlock()
-
-	if b.merkleLayers == nil || len(b.merkleLayers) == 0 {
-		fieldRoots, err := computeFieldRoots(ctx, b.state)
-		if err != nil {
-			return [32]byte{}, err
-		}
-		layers := stateutil.Merkleize(fieldRoots)
-		b.merkleLayers = layers
-		b.dirtyFields = make(map[types.FieldIndex]bool, params.BeaconConfig().BeaconStateAltairFieldCount)
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return [32]byte{}, err
 	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return [32]byte{}, err
+	}
+	return bytesutil.ToBytes32(b.merkleLayers[len(b.merkleLayers)-1][0]), nil
+}
 
+// Initializes the Merkle layers for the beacon state if they are empty.
+// WARNING: Caller must acquire the mutex before using.
+func (b *BeaconState) initializeMerkleLayers(ctx context.Context) error {
+	if len(b.merkleLayers) > 0 {
+		return nil
+	}
+	fieldRoots, err := computeFieldRoots(ctx, b.state)
+	if err != nil {
+		return err
+	}
+	layers := stateutil.Merkleize(fieldRoots)
+	b.merkleLayers = layers
+	b.dirtyFields = make(map[types.FieldIndex]bool, params.BeaconConfig().BeaconStateAltairFieldCount)
+	return nil
+}
+
+// Recomputes the Merkle layers for the dirty fields in the state.
+// WARNING: Caller must acquire the mutex before using.
+func (b *BeaconState) recomputeDirtyFields(ctx context.Context) error {
 	for field := range b.dirtyFields {
 		root, err := b.rootSelector(ctx, field)
 		if err != nil {
-			return [32]byte{}, err
+			return err
 		}
 		b.merkleLayers[0][field] = root[:]
 		b.recomputeRoot(int(field))
 		delete(b.dirtyFields, field)
 	}
-	return bytesutil.ToBytes32(b.merkleLayers[len(b.merkleLayers)-1][0]), nil
+	return nil
 }
 
 // FieldReferencesCount returns the reference count held by each field. This

beacon-chain/state/v3/proofs.go

@@ -1,6 +1,7 @@
 package v3
 
 import (
+	"context"
 	"encoding/binary"
 
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/fieldtrie"
@@ -27,24 +28,46 @@ func NextSyncCommitteeGeneralizedIndex() uint64 {
 }
 
 // CurrentSyncCommitteeProof from the state's Merkle trie representation.
-func (b *BeaconState) CurrentSyncCommitteeProof() ([][]byte, error) {
-	b.lock.RLock()
-	defer b.lock.RUnlock()
+func (b *BeaconState) CurrentSyncCommitteeProof(ctx context.Context) ([][]byte, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	// In case the Merkle layers of the trie are not populated, we need
+	// to perform some initialization.
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return nil, err
+	}
+	// Our beacon state uses a "dirty" fields pattern which requires us to
+	// recompute branches of the Merkle layers that are marked as dirty.
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return nil, err
+	}
 	return fieldtrie.ProofFromMerkleLayers(b.merkleLayers, currentSyncCommittee), nil
 }
 
 // NextSyncCommitteeProof from the state's Merkle trie representation.
-func (b *BeaconState) NextSyncCommitteeProof() ([][]byte, error) {
-	b.lock.RLock()
-	defer b.lock.RUnlock()
+func (b *BeaconState) NextSyncCommitteeProof(ctx context.Context) ([][]byte, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return nil, err
+	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return nil, err
+	}
 	return fieldtrie.ProofFromMerkleLayers(b.merkleLayers, nextSyncCommittee), nil
 }
 
 // FinalizedRootProof crafts a Merkle proof for the finalized root
 // contained within the finalized checkpoint of a beacon state.
-func (b *BeaconState) FinalizedRootProof() ([][]byte, error) {
-	b.lock.RLock()
-	defer b.lock.RUnlock()
+func (b *BeaconState) FinalizedRootProof(ctx context.Context) ([][]byte, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return nil, err
+	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return nil, err
+	}
 	cpt := b.state.FinalizedCheckpoint
 	// The epoch field of a finalized checkpoint is the neighbor
 	// index of the finalized root field in its Merkle tree representation

beacon-chain/state/v3/proofs_test.go

@@ -23,7 +23,7 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 		// Verify the Merkle proof.
 		scRoot, err := sc.HashTreeRoot()
 		require.NoError(t, err)
-		proof, err := st.CurrentSyncCommitteeProof()
+		proof, err := st.CurrentSyncCommitteeProof(ctx)
 		require.NoError(t, err)
 		valid := trie.VerifyMerkleProof(htr[:], scRoot[:], v3.CurrentSyncCommitteeGeneralizedIndex(), proof)
 		require.Equal(t, true, valid)
@@ -31,7 +31,7 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 	t.Run("next sync committee", func(t *testing.T) {
 		nextSC, err := st.NextSyncCommittee()
 		require.NoError(t, err)
-		proof, err := st.NextSyncCommitteeProof()
+		proof, err := st.NextSyncCommitteeProof(ctx)
 		require.NoError(t, err)
 
 		// Verify the Merkle proof.
@@ -53,7 +53,7 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 		require.Equal(t, false, valid)
 
 		// Generating a new, valid proof should pass.
-		proof, err = st.NextSyncCommitteeProof()
+		proof, err = st.NextSyncCommitteeProof(ctx)
 		require.NoError(t, err)
 		htr, err = st.HashTreeRoot(ctx)
 		require.NoError(t, err)
@@ -66,10 +66,40 @@ func TestBeaconStateMerkleProofs(t *testing.T) {
 		// Verify the Merkle proof.
 		htr, err = st.HashTreeRoot(ctx)
 		require.NoError(t, err)
-		proof, err := st.FinalizedRootProof()
+		proof, err := st.FinalizedRootProof(ctx)
 		require.NoError(t, err)
 		gIndex := v3.FinalizedRootGeneralizedIndex()
 		valid := trie.VerifyMerkleProof(htr[:], finalizedRoot, gIndex, proof)
 		require.Equal(t, true, valid)
 	})
+	t.Run("recomputes root on dirty fields", func(t *testing.T) {
+		currentRoot, err := st.HashTreeRoot(ctx)
+		require.NoError(t, err)
+		cpt := st.FinalizedCheckpoint()
+		require.NoError(t, err)
+
+		// Edit the checkpoint.
+		cpt.Epoch = 100
+		require.NoError(t, st.SetFinalizedCheckpoint(cpt))
+
+		// Produce a proof for the finalized root.
+		proof, err := st.FinalizedRootProof(ctx)
+		require.NoError(t, err)
+
+		// We expect the previous step to have triggered
+		// a recomputation of dirty fields in the beacon state, resulting
+		// in a new hash tree root as the finalized checkpoint had previously
+		// changed and should have been marked as a dirty state field.
+		// The proof validity should be false for the old root, but true for the new.
+		finalizedRoot := st.FinalizedCheckpoint().Root
+		gIndex := v3.FinalizedRootGeneralizedIndex()
+		valid := trie.VerifyMerkleProof(currentRoot[:], finalizedRoot, gIndex, proof)
+		require.Equal(t, false, valid)
+
+		newRoot, err := st.HashTreeRoot(ctx)
+		require.NoError(t, err)
+
+		valid = trie.VerifyMerkleProof(newRoot[:], finalizedRoot, gIndex, proof)
+		require.Equal(t, true, valid)
+	})
 }

beacon-chain/state/v3/state_trie.go

@@ -208,27 +208,44 @@ func (b *BeaconState) HashTreeRoot(ctx context.Context) ([32]byte, error) {
 
 	b.lock.Lock()
 	defer b.lock.Unlock()
-
-	if b.merkleLayers == nil || len(b.merkleLayers) == 0 {
-		fieldRoots, err := computeFieldRoots(ctx, b.state)
-		if err != nil {
-			return [32]byte{}, err
-		}
-		layers := stateutil.Merkleize(fieldRoots)
-		b.merkleLayers = layers
-		b.dirtyFields = make(map[types.FieldIndex]bool, params.BeaconConfig().BeaconStateBellatrixFieldCount)
+	if err := b.initializeMerkleLayers(ctx); err != nil {
+		return [32]byte{}, err
 	}
+	if err := b.recomputeDirtyFields(ctx); err != nil {
+		return [32]byte{}, err
+	}
+	return bytesutil.ToBytes32(b.merkleLayers[len(b.merkleLayers)-1][0]), nil
+}
 
+// Initializes the Merkle layers for the beacon state if they are empty.
+// WARNING: Caller must acquire the mutex before using.
+func (b *BeaconState) initializeMerkleLayers(ctx context.Context) error {
+	if len(b.merkleLayers) > 0 {
+		return nil
+	}
+	fieldRoots, err := computeFieldRoots(ctx, b.state)
+	if err != nil {
+		return err
+	}
+	layers := stateutil.Merkleize(fieldRoots)
+	b.merkleLayers = layers
+	b.dirtyFields = make(map[types.FieldIndex]bool, params.BeaconConfig().BeaconStateBellatrixFieldCount)
+	return nil
+}
+
+// Recomputes the Merkle layers for the dirty fields in the state.
+// WARNING: Caller must acquire the mutex before using.
+func (b *BeaconState) recomputeDirtyFields(_ context.Context) error {
 	for field := range b.dirtyFields {
 		root, err := b.rootSelector(field)
 		if err != nil {
-			return [32]byte{}, err
+			return err
 		}
 		b.merkleLayers[0][field] = root[:]
 		b.recomputeRoot(int(field))
 		delete(b.dirtyFields, field)
 	}
-	return bytesutil.ToBytes32(b.merkleLayers[len(b.merkleLayers)-1][0]), nil
+	return nil
 }
 
 // FieldReferencesCount returns the reference count held by each field. This