From 6eb158c16a0d96953b894e2f048e00e89f845789 Mon Sep 17 00:00:00 2001
From: Nishant Das <nishdas93@gmail.com>
Date: Mon, 27 Sep 2021 16:27:11 +0800
Subject: [PATCH] Check Head State For New Head Methods (#9676)

* check head state

* add tests
---
 beacon-chain/blockchain/chain_info.go      |  8 +++++--
 beacon-chain/blockchain/chain_info_test.go | 28 ++++++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/beacon-chain/blockchain/chain_info.go b/beacon-chain/blockchain/chain_info.go
index 0972eafba..09359d4c5 100644
--- a/beacon-chain/blockchain/chain_info.go
+++ b/beacon-chain/blockchain/chain_info.go
@@ -293,7 +293,9 @@ func (s *Service) ChainHeads() ([][32]byte, []types.Slot) {
 func (s *Service) HeadPublicKeyToValidatorIndex(ctx context.Context, pubKey [48]byte) (types.ValidatorIndex, bool) {
 	s.headLock.RLock()
 	defer s.headLock.RUnlock()
-
+	if !s.hasHeadState() {
+		return 0, false
+	}
 	return s.headState(ctx).ValidatorIndexByPubkey(pubKey)
 }
 
@@ -301,7 +303,9 @@ func (s *Service) HeadPublicKeyToValidatorIndex(ctx context.Context, pubKey [48]
 func (s *Service) HeadValidatorIndexToPublicKey(_ context.Context, index types.ValidatorIndex) ([48]byte, error) {
 	s.headLock.RLock()
 	defer s.headLock.RUnlock()
-
+	if !s.hasHeadState() {
+		return [48]byte{}, nil
+	}
 	v, err := s.headValidatorAtIndex(index)
 	if err != nil {
 		return [48]byte{}, err
diff --git a/beacon-chain/blockchain/chain_info_test.go b/beacon-chain/blockchain/chain_info_test.go
index 6d09abe3f..c4d0b5603 100644
--- a/beacon-chain/blockchain/chain_info_test.go
+++ b/beacon-chain/blockchain/chain_info_test.go
@@ -313,6 +313,20 @@ func TestService_HeadPublicKeyToValidatorIndex(t *testing.T) {
 	require.Equal(t, types.ValidatorIndex(0), i)
 }
 
+func TestService_HeadPublicKeyToValidatorIndexNil(t *testing.T) {
+	c := &Service{}
+	c.head = nil
+
+	idx, e := c.HeadPublicKeyToValidatorIndex(context.Background(), [48]byte{})
+	require.Equal(t, false, e)
+	require.Equal(t, types.ValidatorIndex(0), idx)
+
+	c.head = &head{state: nil}
+	i, e := c.HeadPublicKeyToValidatorIndex(context.Background(), [48]byte{})
+	require.Equal(t, false, e)
+	require.Equal(t, types.ValidatorIndex(0), i)
+}
+
 func TestService_HeadValidatorIndexToPublicKey(t *testing.T) {
 	s, _ := util.DeterministicGenesisState(t, 10)
 	c := &Service{}
@@ -326,3 +340,17 @@ func TestService_HeadValidatorIndexToPublicKey(t *testing.T) {
 
 	require.Equal(t, bytesutil.ToBytes48(v.PublicKey), p)
 }
+
+func TestService_HeadValidatorIndexToPublicKeyNil(t *testing.T) {
+	c := &Service{}
+	c.head = nil
+
+	p, err := c.HeadValidatorIndexToPublicKey(context.Background(), 0)
+	require.NoError(t, err)
+	require.Equal(t, [48]byte{}, p)
+
+	c.head = &head{state: nil}
+	p, err = c.HeadValidatorIndexToPublicKey(context.Background(), 0)
+	require.NoError(t, err)
+	require.Equal(t, [48]byte{}, p)
+}