From cfbfccb203eba6c41d21276dc88ef348dcb993ab Mon Sep 17 00:00:00 2001
From: james-prysm <90280386+james-prysm@users.noreply.github.com>
Date: Mon, 7 Oct 2024 09:35:10 -0500
Subject: [PATCH] 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 <pvanloon@offchainlabs.com>

* fixing linting

---------

Co-authored-by: Preston Van Loon <pvanloon@offchainlabs.com>
---
 validator/client/validator.go      |  5 ++-
 validator/client/validator_test.go | 60 ++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/validator/client/validator.go b/validator/client/validator.go
index d6eca9ae6fee..b0b62f004d47 100644
--- 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
index 56abfd651b6b..a707e0bcab48 100644
--- 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)
+	}
+}