Only use head if it's compatible with target (#15965)

* Only use head if it's compatible with target

* Allow blocks from the previous epoch to be viable for checkpoints

* Add feature flag to make it configurable

* fix tests

* @satushh's review

* Manu's nit

* Use fields in logs

beacon-chain/blockchain/process_attestation_helpers.go

@@ -1,6 +1,7 @@
 package blockchain
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"strconv"
@@ -16,6 +17,7 @@ import (
 	ethpb "github.com/OffchainLabs/prysm/v6/proto/prysm/v1alpha1"
 	"github.com/OffchainLabs/prysm/v6/time/slots"
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
 )
 
 // The caller of this function must have a lock on forkchoice.
@@ -27,6 +29,20 @@ func (s *Service) getRecentPreState(ctx context.Context, c *ethpb.Checkpoint) st
 	if !s.cfg.ForkChoiceStore.IsCanonical([32]byte(c.Root)) {
 		return nil
 	}
+	// Only use head state if the head state is compatible with the target checkpoint.
+	headRoot, err := s.HeadRoot(ctx)
+	if err != nil {
+		return nil
+	}
+	headTarget, err := s.cfg.ForkChoiceStore.TargetRootForEpoch([32]byte(headRoot), c.Epoch)
+	if err != nil {
+		return nil
+	}
+	if !bytes.Equal(c.Root, headTarget[:]) {
+		return nil
+	}
+
+	// If the head state alone is enough, we can return it directly read only.
 	if c.Epoch == headEpoch {
 		st, err := s.HeadStateReadOnly(ctx)
 		if err != nil {
@@ -34,11 +50,13 @@ func (s *Service) getRecentPreState(ctx context.Context, c *ethpb.Checkpoint) st
 		}
 		return st
 	}
+	// Otherwise we need to advance the head state to the start of the target epoch.
+	// This point can only be reached if c.Root == headRoot and c.Epoch > headEpoch.
 	slot, err := slots.EpochStart(c.Epoch)
 	if err != nil {
 		return nil
 	}
-	// Try if we have already set the checkpoint cache
+	// Try if we have already set the checkpoint cache. This will be tried again if we fail here but the check is cheap anyway.
 	epochKey := strconv.FormatUint(uint64(c.Epoch), 10 /* base 10 */)
 	lock := async.NewMultilock(string(c.Root) + epochKey)
 	lock.Lock()
@@ -50,6 +68,7 @@ func (s *Service) getRecentPreState(ctx context.Context, c *ethpb.Checkpoint) st
 	if cachedState != nil && !cachedState.IsNil() {
 		return cachedState
 	}
+	// If we haven't advanced yet then process the slots from head state.
 	st, err := s.HeadState(ctx)
 	if err != nil {
 		return nil
@@ -114,6 +133,7 @@ func (s *Service) getAttPreState(ctx context.Context, c *ethpb.Checkpoint) (stat
 	}
 
 	// Fallback to state regeneration.
+	log.WithFields(logrus.Fields{"epoch": c.Epoch, "root": fmt.Sprintf("%#x", c.Root)}).Debug("Regenerating attestation pre-state")
 	baseState, err := s.cfg.StateGen.StateByRoot(ctx, bytesutil.ToBytes32(c.Root))
 	if err != nil {
 		return nil, errors.Wrapf(err, "could not get pre state for epoch %d", c.Epoch)

beacon-chain/forkchoice/doubly-linked-tree/BUILD.bazel

@@ -27,6 +27,7 @@ go_library(
         "//beacon-chain/forkchoice:go_default_library",
         "//beacon-chain/forkchoice/types:go_default_library",
         "//beacon-chain/state:go_default_library",
+        "//config/features:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
         "//consensus-types/blocks:go_default_library",

beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go

@@ -8,6 +8,7 @@ import (
 	"github.com/OffchainLabs/prysm/v6/beacon-chain/forkchoice"
 	forkchoicetypes "github.com/OffchainLabs/prysm/v6/beacon-chain/forkchoice/types"
 	"github.com/OffchainLabs/prysm/v6/beacon-chain/state"
+	"github.com/OffchainLabs/prysm/v6/config/features"
 	fieldparams "github.com/OffchainLabs/prysm/v6/config/fieldparams"
 	"github.com/OffchainLabs/prysm/v6/config/params"
 	consensus_blocks "github.com/OffchainLabs/prysm/v6/consensus-types/blocks"
@@ -239,9 +240,12 @@ func (f *ForkChoice) IsViableForCheckpoint(cp *forkchoicetypes.Checkpoint) (bool
 	if node.slot == epochStart {
 		return true, nil
 	}
-	nodeEpoch := slots.ToEpoch(node.slot)
-	if nodeEpoch >= cp.Epoch {
-		return false, nil
+	if !features.Get().DisableLastEpochTargets {
+		// Allow any node from the checkpoint epoch - 1 to be viable.
+		nodeEpoch := slots.ToEpoch(node.slot)
+		if nodeEpoch+1 == cp.Epoch {
+			return true, nil
+		}
 	}
 	for _, child := range node.children {
 		if child.slot > epochStart {

beacon-chain/forkchoice/doubly-linked-tree/forkchoice_test.go

@@ -813,9 +813,10 @@ func TestForkChoiceIsViableForCheckpoint(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, true, viable)
 
+	// Last epoch blocks are still viable
 	viable, err = f.IsViableForCheckpoint(&forkchoicetypes.Checkpoint{Root: blk.Root(), Epoch: 1})
 	require.NoError(t, err)
-	require.Equal(t, false, viable)
+	require.Equal(t, true, viable)
 
 	// No Children but impossible checkpoint
 	viable, err = f.IsViableForCheckpoint(&forkchoicetypes.Checkpoint{Root: blk2.Root()})
@@ -835,9 +836,10 @@ func TestForkChoiceIsViableForCheckpoint(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, false, viable)
 
+	// Last epoch blocks are still viable
 	viable, err = f.IsViableForCheckpoint(&forkchoicetypes.Checkpoint{Root: blk2.Root(), Epoch: 1})
 	require.NoError(t, err)
-	require.Equal(t, false, viable)
+	require.Equal(t, true, viable)
 
 	st, blk4, err := prepareForkchoiceState(ctx, params.BeaconConfig().SlotsPerEpoch, [32]byte{'d'}, blk2.Root(), [32]byte{'D'}, 0, 0)
 	require.NoError(t, err)
@@ -848,9 +850,10 @@ func TestForkChoiceIsViableForCheckpoint(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, false, viable)
 
+	// Last epoch blocks are still viable
 	viable, err = f.IsViableForCheckpoint(&forkchoicetypes.Checkpoint{Root: blk2.Root(), Epoch: 1})
 	require.NoError(t, err)
-	require.Equal(t, false, viable)
+	require.Equal(t, true, viable)
 
 	// Boundary block
 	viable, err = f.IsViableForCheckpoint(&forkchoicetypes.Checkpoint{Root: blk4.Root(), Epoch: 1})

changelog/potuz_head_target_compat.md

@@ -0,0 +1,3 @@
+### Fixed
+
+- Use head only if its compatible with target for attestation validation.

config/features/config.go

@@ -69,6 +69,7 @@ type Flags struct {
 
 	DisableResourceManager     bool // Disables running the node with libp2p's resource manager.
 	DisableStakinContractCheck bool // Disables check for deposit contract when proposing blocks
+	DisableLastEpochTargets    bool // Disables processing of states for attestations to old blocks.
 
 	EnableVerboseSigVerification bool // EnableVerboseSigVerification specifies whether to verify individual signature if batch verification fails
 
@@ -274,11 +275,14 @@ func ConfigureBeaconChain(ctx *cli.Context) error {
 		logEnabled(forceHeadFlag)
 		cfg.ForceHead = ctx.String(forceHeadFlag.Name)
 	}
-
 	if ctx.IsSet(blacklistRoots.Name) {
 		logEnabled(blacklistRoots)
 		cfg.BlacklistedRoots = parseBlacklistedRoots(ctx.StringSlice(blacklistRoots.Name))
 	}
+	if ctx.IsSet(disableLastEpochTargets.Name) {
+		logEnabled(disableLastEpochTargets)
+		cfg.DisableLastEpochTargets = true
+	}
 
 	cfg.AggregateIntervals = [3]time.Duration{aggregateFirstInterval.Value, aggregateSecondInterval.Value, aggregateThirdInterval.Value}
 	Init(cfg)

config/features/flags.go

@@ -197,6 +197,11 @@ var (
 		Usage: "(Work in progress): Enables the web portal for the validator client.",
 		Value: false,
 	}
+	// disableLastEpochTargets is a flag to disable processing of attestations for old blocks.
+	disableLastEpochTargets = &cli.BoolFlag{
+		Name:  "disable-last-epoch-targets",
+		Usage: "Disables processing of last epoch targets.",
+	}
 )
 
 // devModeFlags holds list of flags that are set when development mode is on.
@@ -257,6 +262,7 @@ var BeaconChainFlags = combinedFlags([]cli.Flag{
 	enableExperimentalAttestationPool,
 	forceHeadFlag,
 	blacklistRoots,
+	disableLastEpochTargets,
 }, deprecatedBeaconFlags, deprecatedFlags, upcomingDeprecation)
 
 func combinedFlags(flags ...[]cli.Flag) []cli.Flag {