ValidatorIndexByPubkey length check (#9424)

* `ValidatorIndexByPubkey` length check

* gaz

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/state/v2/BUILD.bazel

@@ -63,5 +63,6 @@ go_test(
         "//shared/params:go_default_library",
         "//shared/testutil/assert:go_default_library",
         "//shared/testutil/require:go_default_library",
+        "@com_github_prysmaticlabs_eth2_types//:go_default_library",
     ],
 )

beacon-chain/state/v2/getters.go

@@ -581,7 +581,12 @@ func (b *BeaconState) ValidatorIndexByPubkey(key [48]byte) (types.ValidatorIndex
 	}
 	b.lock.RLock()
 	defer b.lock.RUnlock()
+	numOfVals := len(b.state.Validators)
+
 	idx, ok := b.valMapHandler.Get(key)
+	if ok && numOfVals <= int(idx) {
+		return types.ValidatorIndex(0), false
+	}
 	return idx, ok
 }
 

beacon-chain/state/v2/getters_test.go

@@ -5,7 +5,9 @@ import (
 	"sync"
 	"testing"
 
+	types "github.com/prysmaticlabs/eth2-types"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
 )
 
@@ -86,3 +88,105 @@ func TestNilState_NoPanic(t *testing.T) {
 	_, err = st.NextSyncCommittee()
 	_ = err
 }
+
+func TestBeaconState_ValidatorByPubkey(t *testing.T) {
+	keyCreator := func(input []byte) [48]byte {
+		nKey := [48]byte{}
+		copy(nKey[:1], input)
+		return nKey
+	}
+
+	tests := []struct {
+		name            string
+		modifyFunc      func(b *BeaconState, k [48]byte)
+		exists          bool
+		expectedIdx     types.ValidatorIndex
+		largestIdxInSet types.ValidatorIndex
+	}{
+		{
+			name: "retrieve validator",
+			modifyFunc: func(b *BeaconState, key [48]byte) {
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
+			},
+			exists:      true,
+			expectedIdx: 0,
+		},
+		{
+			name: "retrieve validator with multiple validators from the start",
+			modifyFunc: func(b *BeaconState, key [48]byte) {
+				key1 := keyCreator([]byte{'C'})
+				key2 := keyCreator([]byte{'D'})
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
+			},
+			exists:      true,
+			expectedIdx: 0,
+		},
+		{
+			name: "retrieve validator with multiple validators",
+			modifyFunc: func(b *BeaconState, key [48]byte) {
+				key1 := keyCreator([]byte{'C'})
+				key2 := keyCreator([]byte{'D'})
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
+			},
+			exists:      true,
+			expectedIdx: 2,
+		},
+		{
+			name: "retrieve validator with multiple validators from the start with shared state",
+			modifyFunc: func(b *BeaconState, key [48]byte) {
+				key1 := keyCreator([]byte{'C'})
+				key2 := keyCreator([]byte{'D'})
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
+				_ = b.Copy()
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
+			},
+			exists:      true,
+			expectedIdx: 0,
+		},
+		{
+			name: "retrieve validator with multiple validators with shared state",
+			modifyFunc: func(b *BeaconState, key [48]byte) {
+				key1 := keyCreator([]byte{'C'})
+				key2 := keyCreator([]byte{'D'})
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
+				n := b.Copy()
+				// Append to another state
+				assert.NoError(t, n.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
+
+			},
+			exists:      false,
+			expectedIdx: 0,
+		},
+		{
+			name: "retrieve validator with multiple validators with shared state at boundary",
+			modifyFunc: func(b *BeaconState, key [48]byte) {
+				key1 := keyCreator([]byte{'C'})
+				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
+				n := b.Copy()
+				// Append to another state
+				assert.NoError(t, n.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
+
+			},
+			exists:      false,
+			expectedIdx: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s, err := InitializeFromProto(&ethpb.BeaconStateAltair{})
+			require.NoError(t, err)
+			nKey := keyCreator([]byte{'A'})
+			tt.modifyFunc(s, nKey)
+			idx, ok := s.ValidatorIndexByPubkey(nKey)
+			assert.Equal(t, tt.exists, ok)
+			assert.Equal(t, tt.expectedIdx, idx)
+		})
+	}
+}