From e118db6c20172783a3bb573d2f78670f0207ec1e Mon Sep 17 00:00:00 2001
From: nisdas <nishdas93@gmail.com>
Date: Mon, 14 Mar 2022 19:26:35 +0800
Subject: [PATCH] fix broken conditional checks

---
 beacon-chain/blockchain/optimistic_sync.go | 16 +++-------
 beacon-chain/blockchain/process_block.go   | 36 ++++++++++++++++------
 2 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/beacon-chain/blockchain/optimistic_sync.go b/beacon-chain/blockchain/optimistic_sync.go
index cd2c93e074fa..1b016aa7419b 100644
--- a/beacon-chain/blockchain/optimistic_sync.go
+++ b/beacon-chain/blockchain/optimistic_sync.go
@@ -12,7 +12,6 @@ import (
 	"github.com/prysmaticlabs/prysm/config/params"
 	"github.com/prysmaticlabs/prysm/encoding/bytesutil"
 	enginev1 "github.com/prysmaticlabs/prysm/proto/engine/v1"
-	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1/block"
 	"github.com/prysmaticlabs/prysm/runtime/version"
 	"github.com/sirupsen/logrus"
@@ -80,7 +79,7 @@ func (s *Service) notifyForkchoiceUpdate(ctx context.Context, headBlk block.Beac
 }
 
 // notifyForkchoiceUpdate signals execution engine on a new payload
-func (s *Service) notifyNewPayload(ctx context.Context, preStateVersion int, header *ethpb.ExecutionPayloadHeader, postState state.BeaconState, blk block.SignedBeaconBlock) error {
+func (s *Service) notifyNewPayload(ctx context.Context, preStateIsMerged, postStateIsMerged bool, postState state.BeaconState, blk block.SignedBeaconBlock) error {
 	if postState == nil {
 		return errors.New("pre and post states must not be nil")
 	}
@@ -92,11 +91,7 @@ func (s *Service) notifyNewPayload(ctx context.Context, preStateVersion int, hea
 		return err
 	}
 	body := blk.Block().Body()
-	enabled, err := blocks.ExecutionEnabled(postState, blk.Block().Body())
-	if err != nil {
-		return errors.Wrap(err, "could not determine if execution is enabled")
-	}
-	if !enabled {
+	if !postStateIsMerged {
 		return nil
 	}
 	payload, err := body.ExecutionPayload()
@@ -118,13 +113,10 @@ func (s *Service) notifyNewPayload(ctx context.Context, preStateVersion int, hea
 	}
 
 	// During the transition event, the transition block should be verified for sanity.
-	if isPreBellatrix(preStateVersion) {
+	if isPreBellatrix(postState.Version()) {
 		return nil
 	}
-	atTransition, err := blocks.IsMergeTransitionBlockUsingPayloadHeader(header, body)
-	if err != nil {
-		return errors.Wrap(err, "could not check if merge block is terminal")
-	}
+	atTransition := !preStateIsMerged && postStateIsMerged
 	if !atTransition {
 		return nil
 	}
diff --git a/beacon-chain/blockchain/process_block.go b/beacon-chain/blockchain/process_block.go
index a024c9d14ae5..bcd5cd87e90f 100644
--- a/beacon-chain/blockchain/process_block.go
+++ b/beacon-chain/blockchain/process_block.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/blocks"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/feed"
 	statefeed "github.com/prysmaticlabs/prysm/beacon-chain/core/feed/state"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
@@ -98,16 +99,24 @@ func (s *Service) onBlock(ctx context.Context, signed block.SignedBeaconBlock, b
 	if err != nil {
 		return err
 	}
-
-	preStateVersion, preStateHeader, err := getStateVersionAndPayload(preState)
-	if err != nil {
-		return err
+	preStateIsMerged, postStateIsMerged := false, false
+	if preState.Version() >= version.Bellatrix {
+		preStateIsMerged, err = blocks.MergeTransitionComplete(preState)
+		if err != nil {
+			return err
+		}
 	}
 	postState, err := transition.ExecuteStateTransition(ctx, preState, signed)
 	if err != nil {
 		return err
 	}
-	if err := s.notifyNewPayload(ctx, preStateVersion, preStateHeader, postState, signed); err != nil {
+	if postState.Version() >= version.Bellatrix {
+		postStateIsMerged, err = blocks.MergeTransitionComplete(postState)
+		if err != nil {
+			return err
+		}
+	}
+	if err := s.notifyNewPayload(ctx, preStateIsMerged, postStateIsMerged, postState, signed); err != nil {
 		return errors.Wrap(err, "could not verify new payload")
 	}
 	// TODO(10261) Check optimistic status
@@ -310,15 +319,24 @@ func (s *Service) onBlockBatch(ctx context.Context, blks []block.SignedBeaconBlo
 	var set *bls.SignatureBatch
 	boundaries := make(map[[32]byte]state.BeaconState)
 	for i, b := range blks {
-		preStateVersion, preStateHeader, err := getStateVersionAndPayload(preState)
-		if err != nil {
-			return nil, nil, nil, err
+		preStateIsMerged, postStateIsMerged := false, false
+		if preState.Version() >= version.Bellatrix {
+			preStateIsMerged, err = blocks.MergeTransitionComplete(preState)
+			if err != nil {
+				return nil, nil, nil, err
+			}
 		}
 		set, preState, err = transition.ExecuteStateTransitionNoVerifyAnySig(ctx, preState, b)
 		if err != nil {
 			return nil, nil, nil, err
 		}
-		if err := s.notifyNewPayload(ctx, preStateVersion, preStateHeader, preState, b); err != nil {
+		if preState.Version() >= version.Bellatrix {
+			postStateIsMerged, err = blocks.MergeTransitionComplete(preState)
+			if err != nil {
+				return nil, nil, nil, err
+			}
+		}
+		if err := s.notifyNewPayload(ctx, preStateIsMerged, postStateIsMerged, preState, b); err != nil {
 			return nil, nil, nil, err
 		}