Fix: validator status cache does not clear upon key removal (#14504)

* adding fix and unit test * Update validator.go Co-authored-by: Preston Van Loon <[email protected]> * fixing linting --------- Co-authored-by: Preston Van Loon <[email protected]>
prysmaticlabs · Oct 7, 2024 · cfbfccb · cfbfccb
1 parent 884b663
commit cfbfccb
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 1 deletion.
diff --git a/validator/client/validator.go b/validator/client/validator.go
@@ -1262,13 +1262,16 @@ func (v *validator) updateValidatorStatusCache(ctx context.Context, pubkeys [][f
 	if len(resp.Statuses) != len(resp.Indices) {
 		return fmt.Errorf("expected %d indices in status, received %d", len(resp.Statuses), len(resp.Indices))
 	}
+	pubkeyToStatus := make(map[[fieldparams.BLSPubkeyLength]byte]*validatorStatus, len(resp.Statuses))
 	for i, s := range resp.Statuses {
-		v.pubkeyToStatus[bytesutil.ToBytes48(resp.PublicKeys[i])] = &validatorStatus{
+		pubkeyToStatus[bytesutil.ToBytes48(resp.PublicKeys[i])] = &validatorStatus{
 			publicKey: resp.PublicKeys[i],
 			status:    s,
 			index:     resp.Indices[i],
 		}
 	}
+	v.pubkeyToStatus = pubkeyToStatus
+
 	return nil
 }
 

diff --git a/validator/client/validator_test.go b/validator/client/validator_test.go
@@ -2908,3 +2908,63 @@ func TestValidator_ChangeHost(t *testing.T) {
 	v.ChangeHost()
 	assert.Equal(t, uint64(0), v.currentHostIndex)
 }
+
+func TestUpdateValidatorStatusCache(t *testing.T) {
+	ctx := context.Background()
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	pubkeys := [][fieldparams.BLSPubkeyLength]byte{
+		{0x01},
+		{0x02},
+	}
+	statusRequestKeys := [][]byte{
+		pubkeys[0][:],
+		pubkeys[1][:],
+	}
+
+	client := validatormock.NewMockValidatorClient(ctrl)
+	mockResponse := &ethpb.MultipleValidatorStatusResponse{
+		PublicKeys: statusRequestKeys,
+		Statuses: []*ethpb.ValidatorStatusResponse{
+			{
+				Status: ethpb.ValidatorStatus_ACTIVE,
+			}, {
+				Status: ethpb.ValidatorStatus_EXITING,
+			}},
+		Indices: []primitives.ValidatorIndex{1, 2},
+	}
+	client.EXPECT().MultipleValidatorStatus(
+		gomock.Any(),
+		gomock.Any()).Return(mockResponse, nil)
+
+	v := &validator{
+		validatorClient:  client,
+		beaconNodeHosts:  []string{"http://localhost:8080", "http://localhost:8081"},
+		currentHostIndex: 0,
+		pubkeyToStatus: map[[fieldparams.BLSPubkeyLength]byte]*validatorStatus{
+			[fieldparams.BLSPubkeyLength]byte{0x03}: &validatorStatus{ // add non existant key and status to cache, should be fully removed on update
+				publicKey: []byte{0x03},
+				status: &ethpb.ValidatorStatusResponse{
+					Status: ethpb.ValidatorStatus_ACTIVE,
+				},
+				index: 3,
+			},
+		},
+	}
+
+	err := v.updateValidatorStatusCache(ctx, pubkeys)
+	assert.NoError(t, err)
+
+	// make sure the nonexistant key is fully removed
+	_, ok := v.pubkeyToStatus[[fieldparams.BLSPubkeyLength]byte{0x03}]
+	require.Equal(t, false, ok)
+	// make sure we only have the added values
+	assert.Equal(t, 2, len(v.pubkeyToStatus))
+	for i, pk := range pubkeys {
+		status, exists := v.pubkeyToStatus[pk]
+		require.Equal(t, true, exists)
+		require.DeepEqual(t, pk[:], status.publicKey)
+		require.Equal(t, mockResponse.Statuses[i], status.status)
+		require.Equal(t, mockResponse.Indices[i], status.index)
+	}
+}