From 9e712e4598dce20b7b0786d46f64371bbd8cd3ca Mon Sep 17 00:00:00 2001
From: terence tsao <terence@prysmaticlabs.com>
Date: Tue, 13 Oct 2020 11:26:43 -0700
Subject: [PATCH] Add justified block insertion to forkchoice when missed
 (#7520)

* Add justified block insertion to fork choice when missed

* Add regression test
---
 beacon-chain/blockchain/head.go      | 16 ++++++++++++++++
 beacon-chain/blockchain/head_test.go | 17 +++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/beacon-chain/blockchain/head.go b/beacon-chain/blockchain/head.go
index 42b649c583..67ee31ecd6 100644
--- a/beacon-chain/blockchain/head.go
+++ b/beacon-chain/blockchain/head.go
@@ -51,6 +51,22 @@ func (s *Service) updateHead(ctx context.Context, balances []uint64) error {
 	if headStartRoot == params.BeaconConfig().ZeroHash {
 		headStartRoot = s.genesisRoot
 	}
+
+	// In order to process head, fork choice store requires justified info.
+	// If the fork choice store is missing justified block info, a node should
+	// re-initiate fork choice store using the latest justified info.
+	// This recovers a fatal condition and should not happen in run time.
+	if !s.forkChoiceStore.HasNode(headStartRoot) {
+		jb, err := s.beaconDB.Block(ctx, headStartRoot)
+		if err != nil {
+			return err
+		}
+		s.forkChoiceStore = protoarray.New(j.Epoch, f.Epoch, bytesutil.ToBytes32(f.Root))
+		if err := s.insertBlockToForkChoiceStore(ctx, jb.Block, headStartRoot, f, j); err != nil {
+			return err
+		}
+	}
+
 	headRoot, err := s.forkChoiceStore.Head(ctx, j.Epoch, headStartRoot, balances, f.Epoch)
 	if err != nil {
 		return err
diff --git a/beacon-chain/blockchain/head_test.go b/beacon-chain/blockchain/head_test.go
index e968a4ad76..b593d9dd1a 100644
--- a/beacon-chain/blockchain/head_test.go
+++ b/beacon-chain/blockchain/head_test.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"testing"
 
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	testDB "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
@@ -137,3 +138,19 @@ func TestCacheJustifiedStateBalances_CanCache(t *testing.T) {
 	require.NoError(t, service.cacheJustifiedStateBalances(context.Background(), r))
 	require.DeepEqual(t, service.getJustifiedBalances(), state.Balances(), "Incorrect justified balances")
 }
+
+func TestUpdateHead_MissingJustifiedRoot(t *testing.T) {
+	db, sc := testDB.SetupDB(t)
+	service := setupBeaconChain(t, db, sc)
+
+	b := testutil.NewBeaconBlock()
+	require.NoError(t, service.beaconDB.SaveBlock(context.Background(), b))
+	r, err := b.Block.HashTreeRoot()
+	require.NoError(t, err)
+
+	service.justifiedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	service.finalizedCheckpt = &ethpb.Checkpoint{}
+	service.bestJustifiedCheckpt = &ethpb.Checkpoint{}
+
+	require.NoError(t, service.updateHead(context.Background(), []uint64{}))
+}