From 9615484fe9a0312d76ac68a4a4086e64f0bdfb71 Mon Sep 17 00:00:00 2001
From: terencechain <terence@prysmaticlabs.com>
Date: Mon, 2 May 2022 20:48:58 -0700
Subject: [PATCH] notifyForkchoiceUpdate can prune invalid blocks (#10510)

* notifyForkchoiceUpdate can prune invalid blocks

* define and use ErrInvalidPayload

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
---
 beacon-chain/blockchain/execution_engine.go   | 28 +++++++++++++++++--
 .../blockchain/execution_engine_test.go       |  9 ++++--
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/beacon-chain/blockchain/execution_engine.go b/beacon-chain/blockchain/execution_engine.go
index 75795e1024..2f3809a58d 100644
--- a/beacon-chain/blockchain/execution_engine.go
+++ b/beacon-chain/blockchain/execution_engine.go
@@ -24,7 +24,10 @@ import (
 	"go.opencensus.io/trace"
 )
 
-var ErrUndefinedExecutionEngineError = errors.New("received an undefined ee error")
+var (
+	ErrInvalidPayload                = errors.New("recevied an INVALID payload from execution engine")
+	ErrUndefinedExecutionEngineError = errors.New("received an undefined ee error")
+)
 
 // notifyForkchoiceUpdate signals execution engine the fork choice updates. Execution engine should:
 // 1. Re-organizes the execution payload chain and corresponding state to make head_block_hash the head.
@@ -75,7 +78,7 @@ func (s *Service) notifyForkchoiceUpdate(ctx context.Context, headState state.Be
 		return nil, errors.Wrap(err, "could not get payload attribute")
 	}
 
-	payloadID, _, err := s.cfg.ExecutionEngineCaller.ForkchoiceUpdated(ctx, fcs, attr)
+	payloadID, lastValidHash, err := s.cfg.ExecutionEngineCaller.ForkchoiceUpdated(ctx, fcs, attr)
 	if err != nil {
 		switch err {
 		case powchain.ErrAcceptedSyncingPayloadStatus:
@@ -86,6 +89,20 @@ func (s *Service) notifyForkchoiceUpdate(ctx context.Context, headState state.Be
 				"finalizedPayloadBlockHash": fmt.Sprintf("%#x", bytesutil.Trunc(finalizedHash)),
 			}).Info("Called fork choice updated with optimistic block")
 			return payloadID, s.optimisticCandidateBlock(ctx, headBlk)
+		case powchain.ErrInvalidPayloadStatus:
+			invalidRoots, err := s.ForkChoicer().SetOptimisticToInvalid(ctx, headRoot, bytesutil.ToBytes32(headBlk.ParentRoot()), bytesutil.ToBytes32(lastValidHash))
+			if err != nil {
+				return nil, err
+			}
+			if err := s.removeInvalidBlockAndState(ctx, invalidRoots); err != nil {
+				return nil, err
+			}
+			log.WithFields(logrus.Fields{
+				"slot":         headBlk.Slot(),
+				"blockRoot":    fmt.Sprintf("%#x", headRoot),
+				"invalidCount": len(invalidRoots),
+			}).Warn("Pruned invalid blocks")
+			return nil, ErrInvalidPayload
 		default:
 			return nil, errors.WithMessage(ErrUndefinedExecutionEngineError, err.Error())
 		}
@@ -154,7 +171,12 @@ func (s *Service) notifyNewPayload(ctx context.Context, postStateVersion int,
 		if err := s.removeInvalidBlockAndState(ctx, invalidRoots); err != nil {
 			return false, err
 		}
-		return false, errors.New("could not validate an INVALID payload from execution engine")
+		log.WithFields(logrus.Fields{
+			"slot":         blk.Block().Slot(),
+			"blockRoot":    fmt.Sprintf("%#x", root),
+			"invalidCount": len(invalidRoots),
+		}).Warn("Pruned invalid blocks")
+		return false, ErrInvalidPayload
 	default:
 		return false, errors.WithMessage(ErrUndefinedExecutionEngineError, err.Error())
 	}
diff --git a/beacon-chain/blockchain/execution_engine_test.go b/beacon-chain/blockchain/execution_engine_test.go
index d9c7ab71bc..2a0ae5d2f0 100644
--- a/beacon-chain/blockchain/execution_engine_test.go
+++ b/beacon-chain/blockchain/execution_engine_test.go
@@ -56,10 +56,12 @@ func Test_NotifyForkchoiceUpdate(t *testing.T) {
 	}
 	require.NoError(t, err)
 	require.NoError(t, fcs.InsertOptimisticBlock(ctx, 0, [32]byte{}, [32]byte{}, params.BeaconConfig().ZeroHash, 0, 0))
+	require.NoError(t, fcs.InsertOptimisticBlock(ctx, 1, [32]byte{'a'}, [32]byte{}, params.BeaconConfig().ZeroHash, 0, 0))
 
 	tests := []struct {
 		name             string
 		blk              interfaces.BeaconBlock
+		headRoot         [32]byte
 		finalizedRoot    [32]byte
 		newForkchoiceErr error
 		errString        string
@@ -158,7 +160,8 @@ func Test_NotifyForkchoiceUpdate(t *testing.T) {
 			}(),
 			newForkchoiceErr: powchain.ErrInvalidPayloadStatus,
 			finalizedRoot:    bellatrixBlkRoot,
-			errString:        ErrUndefinedExecutionEngineError.Error(),
+			headRoot:         [32]byte{'a'},
+			errString:        ErrInvalidPayload.Error(),
 		},
 	}
 
@@ -166,7 +169,7 @@ func Test_NotifyForkchoiceUpdate(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			service.cfg.ExecutionEngineCaller = &mockPOW.EngineClient{ErrForkchoiceUpdated: tt.newForkchoiceErr}
 			st, _ := util.DeterministicGenesisState(t, 1)
-			_, err := service.notifyForkchoiceUpdate(ctx, st, tt.blk, service.headRoot(), tt.finalizedRoot)
+			_, err := service.notifyForkchoiceUpdate(ctx, st, tt.blk, tt.headRoot, tt.finalizedRoot)
 			if tt.errString != "" {
 				require.ErrorContains(t, tt.errString, err)
 			} else {
@@ -262,7 +265,7 @@ func Test_NotifyNewPayload(t *testing.T) {
 			postState:      bellatrixState,
 			blk:            bellatrixBlk,
 			newPayloadErr:  powchain.ErrInvalidPayloadStatus,
-			errString:      "could not validate an INVALID payload from execution engine",
+			errString:      ErrInvalidPayload.Error(),
 			isValidPayload: false,
 		},
 		{