Prevent Reinsertion of Orphaned Attestations Into Pool (#9442)

* Can re-save orphaned attestation

* Go fmt

* Fix tests

* Go fmt

* Bug fix flag

* Bug fix flag

* Go fmt

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
(cherry picked from commit 5c96a27)

beacon-chain/blockchain/head.go

@@ -15,6 +15,7 @@ import (
 	ethpbv1 "github.com/prysmaticlabs/prysm/proto/eth/v1"
 	"github.com/prysmaticlabs/prysm/proto/interfaces"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/slotutil"
 	"github.com/sirupsen/logrus"
@@ -143,6 +144,10 @@ func (s *Service) saveHead(ctx context.Context, headRoot [32]byte) error {
 			},
 		})
 
+		if err := s.saveOrphanedAtts(ctx, bytesutil.ToBytes32(r)); err != nil {
+			return err
+		}
+
 		reorgCount.Inc()
 	}
 
@@ -360,3 +365,40 @@ func (s *Service) notifyNewHeadEvent(
 	})
 	return nil
 }
+
+// This saves the attestations inside the beacon block with respect to root `orphanedRoot` back into the
+// attestation pool. It also filters out the attestations that is one epoch older as a
+// defense so invalid attestations don't flow into the attestation pool.
+func (s *Service) saveOrphanedAtts(ctx context.Context, orphanedRoot [32]byte) error {
+	if !featureconfig.Get().CorrectlyInsertOrphanedAtts {
+		return nil
+	}
+
+	orphanedBlk, err := s.cfg.BeaconDB.Block(ctx, orphanedRoot)
+	if err != nil {
+		return err
+	}
+
+	if orphanedBlk == nil || orphanedBlk.IsNil() {
+		return errors.New("orphaned block can't be nil")
+	}
+
+	for _, a := range orphanedBlk.Block().Body().Attestations() {
+		// Is the attestation one epoch older.
+		if a.Data.Slot+params.BeaconConfig().SlotsPerEpoch < s.CurrentSlot() {
+			continue
+		}
+		if helpers.IsAggregated(a) {
+			if err := s.cfg.AttPool.SaveAggregatedAttestation(a); err != nil {
+				return err
+			}
+		} else {
+			if err := s.cfg.AttPool.SaveUnaggregatedAttestation(a); err != nil {
+				return err
+			}
+		}
+		saveOrphanedAttCount.Inc()
+	}
+
+	return nil
+}

beacon-chain/blockchain/head_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"testing"
+	"time"
 
 	types "github.com/prysmaticlabs/eth2-types"
 	mock "github.com/prysmaticlabs/prysm/beacon-chain/blockchain/testing"
@@ -13,6 +14,8 @@ import (
 	ethpbv1 "github.com/prysmaticlabs/prysm/proto/eth/v1"
 	ethpb "github.com/prysmaticlabs/prysm/proto/eth/v1alpha1"
 	"github.com/prysmaticlabs/prysm/proto/eth/v1alpha1/wrapper"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
+	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
@@ -36,18 +39,17 @@ func TestSaveHead_Different(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	service := setupBeaconChain(t, beaconDB)
 
-	oldRoot := [32]byte{'A'}
+	testutil.NewBeaconBlock()
+	oldBlock := wrapper.WrappedPhase0SignedBeaconBlock(
+		testutil.NewBeaconBlock(),
+	)
+	require.NoError(t, service.cfg.BeaconDB.SaveBlock(context.Background(), oldBlock))
+	oldRoot, err := oldBlock.Block().HashTreeRoot()
+	require.NoError(t, err)
 	service.head = &head{
-		slot: 0,
-		root: oldRoot,
-		block: wrapper.WrappedPhase0SignedBeaconBlock(
-			&ethpb.SignedBeaconBlock{
-				Block: &ethpb.BeaconBlock{
-					Slot:      0,
-					StateRoot: make([]byte, 32),
-				},
-			},
-		),
+		slot:  0,
+		root:  oldRoot,
+		block: oldBlock,
 	}
 
 	newHeadSignedBlock := testutil.NewBeaconBlock()
@@ -79,18 +81,16 @@ func TestSaveHead_Different_Reorg(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	service := setupBeaconChain(t, beaconDB)
 
-	oldRoot := [32]byte{'A'}
+	oldBlock := wrapper.WrappedPhase0SignedBeaconBlock(
+		testutil.NewBeaconBlock(),
+	)
+	require.NoError(t, service.cfg.BeaconDB.SaveBlock(context.Background(), oldBlock))
+	oldRoot, err := oldBlock.Block().HashTreeRoot()
+	require.NoError(t, err)
 	service.head = &head{
-		slot: 0,
-		root: oldRoot,
-		block: wrapper.WrappedPhase0SignedBeaconBlock(
-			&ethpb.SignedBeaconBlock{
-				Block: &ethpb.BeaconBlock{
-					Slot:      0,
-					StateRoot: make([]byte, 32),
-				},
-			},
-		),
+		slot:  0,
+		root:  oldRoot,
+		block: oldBlock,
 	}
 
 	reorgChainParent := [32]byte{'B'}
@@ -214,3 +214,55 @@ func Test_notifyNewHeadEvent(t *testing.T) {
 		require.DeepSSZEqual(t, wanted, eventHead)
 	})
 }
+
+func TestSaveOrphanedAtts(t *testing.T) {
+	resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+		CorrectlyInsertOrphanedAtts: true,
+	})
+	defer resetCfg()
+
+	genesis, keys := testutil.DeterministicGenesisState(t, 64)
+	b, err := testutil.GenerateFullBlock(genesis, keys, testutil.DefaultBlockGenConfig(), 1)
+	assert.NoError(t, err)
+	r, err := b.Block.HashTreeRoot()
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	beaconDB := testDB.SetupDB(t)
+	service := setupBeaconChain(t, beaconDB)
+	service.genesisTime = time.Now()
+
+	require.NoError(t, service.cfg.BeaconDB.SaveBlock(ctx, wrapper.WrappedPhase0SignedBeaconBlock(b)))
+	require.NoError(t, service.saveOrphanedAtts(ctx, r))
+
+	require.Equal(t, len(b.Block.Body.Attestations), service.cfg.AttPool.AggregatedAttestationCount())
+	savedAtts := service.cfg.AttPool.AggregatedAttestations()
+	atts := b.Block.Body.Attestations
+	require.DeepSSZEqual(t, atts, savedAtts)
+}
+
+func TestSaveOrphanedAtts_CanFilter(t *testing.T) {
+	resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+		CorrectlyInsertOrphanedAtts: true,
+	})
+	defer resetCfg()
+
+	genesis, keys := testutil.DeterministicGenesisState(t, 64)
+	b, err := testutil.GenerateFullBlock(genesis, keys, testutil.DefaultBlockGenConfig(), 1)
+	assert.NoError(t, err)
+	r, err := b.Block.HashTreeRoot()
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	beaconDB := testDB.SetupDB(t)
+	service := setupBeaconChain(t, beaconDB)
+	service.genesisTime = time.Now().Add(time.Duration(-1*int64(params.BeaconConfig().SlotsPerEpoch+1)*int64(params.BeaconConfig().SecondsPerSlot)) * time.Second)
+
+	require.NoError(t, service.cfg.BeaconDB.SaveBlock(ctx, wrapper.WrappedPhase0SignedBeaconBlock(b)))
+	require.NoError(t, service.saveOrphanedAtts(ctx, r))
+
+	require.Equal(t, 0, service.cfg.AttPool.AggregatedAttestationCount())
+	savedAtts := service.cfg.AttPool.AggregatedAttestations()
+	atts := b.Block.Body.Attestations
+	require.DeepNotSSZEqual(t, atts, savedAtts)
+}

beacon-chain/blockchain/metrics.go

@@ -100,6 +100,10 @@ var (
 		Name: "beacon_reorg_total",
 		Help: "Count the number of times beacon chain has a reorg",
 	})
+	saveOrphanedAttCount = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "saved_orphaned_att_total",
+		Help: "Count the number of times an orphaned attestation is saved",
+	})
 	attestationInclusionDelay = promauto.NewHistogram(
 		prometheus.HistogramOpts{
 			Name:    "attestation_inclusion_delay_slots",

shared/featureconfig/config.go

@@ -75,6 +75,7 @@ type Flags struct {
 	EnableSlashingProtectionPruning bool
 
 	// Bug fixes related flags.
+	CorrectlyInsertOrphanedAtts bool
 	CorrectlyPruneCanonicalAtts bool
 }
 
@@ -199,6 +200,10 @@ func ConfigureBeaconChain(ctx *cli.Context) {
 		log.WithField(enableOptimizedBalanceUpdate.Name, enableOptimizedBalanceUpdate.Usage).Warn(enabledFeatureFlag)
 		cfg.EnableOptimizedBalanceUpdate = true
 	}
+	if ctx.Bool(correctlyInsertOrphanedAtts.Name) {
+		log.WithField(correctlyInsertOrphanedAtts.Name, correctlyInsertOrphanedAtts.Usage).Warn(enabledFeatureFlag)
+		cfg.CorrectlyInsertOrphanedAtts = true
+	}
 	if ctx.Bool(correctlyPruneCanonicalAtts.Name) {
 		log.WithField(correctlyPruneCanonicalAtts.Name, correctlyPruneCanonicalAtts.Usage).Warn(enabledFeatureFlag)
 		cfg.CorrectlyPruneCanonicalAtts = true

shared/featureconfig/flags.go

@@ -128,6 +128,11 @@ var (
 			"a foolproof method to find duplicate instances in the network. Your validator will still be" +
 			" vulnerable if it is being run in unsafe configurations.",
 	}
+	correctlyInsertOrphanedAtts = &cli.BoolFlag{
+		Name: "correctly-insert-orphaned-atts",
+		Usage: "This fixes a bug where orphaned attestations don't get reinserted back to mem pool. This improves validator profitability and overall network health," +
+			"see issue #9441 for further detail",
+	}
 	correctlyPruneCanonicalAtts = &cli.BoolFlag{
 		Name: "correctly-prune-canonical-atts",
 		Usage: "This fixes a bug where any block attestations can get incorrectly pruned. This improves validator profitability and overall network health," +
@@ -142,6 +147,7 @@ var devModeFlags = []cli.Flag{
 	forceOptMaxCoverAggregationStategy,
 	updateHeadTimely,
 	enableOptimizedBalanceUpdate,
+	correctlyInsertOrphanedAtts,
 	correctlyPruneCanonicalAtts,
 }
 
@@ -195,6 +201,7 @@ var BeaconChainFlags = append(deprecatedFlags, []cli.Flag{
 	updateHeadTimely,
 	disableProposerAttsSelectionUsingMaxCover,
 	enableOptimizedBalanceUpdate,
+	correctlyInsertOrphanedAtts,
 	correctlyPruneCanonicalAtts,
 }...)
 
@@ -203,5 +210,6 @@ var E2EBeaconChainFlags = []string{
 	"--attestation-aggregation-strategy=opt_max_cover",
 	"--dev",
 	"--use-check-point-cache",
+	"--correctly-insert-orphaned-atts",
 	"--correctly-prune-canonical-atts",
 }