process grpc header flag for exit (#8334)

* process grpc header flag for exit

* ran bazel //:gazelle

* Add regression test

* Re-run "bazel run //:gazelle"

* whitespace

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>

* gofmt

Co-authored-by: Danny Joyce <djoyce@bisontrails.co>
Co-authored-by: Danny Joyce <djoyahoy@users.noreply.github.com>
Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>

validator/accounts/BUILD.bazel

@@ -54,6 +54,7 @@ go_library(
         "@com_github_urfave_cli_v2//:go_default_library",
         "@com_github_wealdtech_go_eth2_wallet_encryptor_keystorev4//:go_default_library",
         "@org_golang_google_grpc//:go_default_library",
+        "@org_golang_google_grpc//metadata:go_default_library",
     ],
 )
 
@@ -98,5 +99,6 @@ go_test(
         "@com_github_sirupsen_logrus//hooks/test:go_default_library",
         "@com_github_urfave_cli_v2//:go_default_library",
         "@com_github_wealdtech_go_eth2_wallet_encryptor_keystorev4//:go_default_library",
+        "@org_golang_google_grpc//metadata:go_default_library",
     ],
 )

validator/accounts/accounts_exit.go

@@ -19,6 +19,7 @@ import (
 	"github.com/prysmaticlabs/prysm/validator/keymanager"
 	"github.com/urfave/cli/v2"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 type performExitCfg struct {
@@ -183,13 +184,22 @@ func prepareClients(cliCtx *cli.Context) (*ethpb.BeaconNodeValidatorClient, *eth
 	if dialOpts == nil {
 		return nil, nil, errors.New("failed to construct dial options")
 	}
+	for _, hdr := range strings.Split(cliCtx.String(flags.GrpcHeadersFlag.Name), ",") {
+		if hdr != "" {
+			ss := strings.Split(hdr, "=")
+			if len(ss) < 2 {
+				log.Warnf("Incorrect gRPC header flag format. Skipping %v", ss[0])
+				continue
+			}
+			cliCtx.Context = metadata.AppendToOutgoingContext(cliCtx.Context, ss[0], strings.Join(ss[1:], "="))
+		}
+	}
 	conn, err := grpc.DialContext(cliCtx.Context, cliCtx.String(flags.BeaconRPCProviderFlag.Name), dialOpts...)
 	if err != nil {
 		return nil, nil, errors.Wrapf(err, "could not dial endpoint %s", flags.BeaconRPCProviderFlag.Name)
 	}
 	validatorClient := ethpb.NewBeaconNodeValidatorClient(conn)
 	nodeClient := ethpb.NewNodeClient(conn)
-
 	return &validatorClient, &nodeClient, nil
 }
 

validator/accounts/accounts_exit_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/prysmaticlabs/prysm/validator/accounts/wallet"
 	"github.com/prysmaticlabs/prysm/validator/keymanager"
 	"github.com/prysmaticlabs/prysm/validator/keymanager/imported"
+	"google.golang.org/grpc/metadata"
 )
 
 func TestExitAccountsCli_Ok(t *testing.T) {
@@ -123,3 +124,28 @@ func TestPrepareWallet_EmptyWalletReturnsError(t *testing.T) {
 	_, _, err = prepareWallet(cliCtx)
 	assert.ErrorContains(t, "wallet is empty", err)
 }
+
+func TestPrepareClients_AddsGRPCHeaders(t *testing.T) {
+	imported.ResetCaches()
+	walletDir, _, passwordFilePath := setupWalletAndPasswordsDir(t)
+	cliCtx := setupWalletCtx(t, &testWalletConfig{
+		walletDir:           walletDir,
+		keymanagerKind:      keymanager.Imported,
+		walletPasswordFile:  passwordFilePath,
+		accountPasswordFile: passwordFilePath,
+		grpcHeaders:         "Authorization=Basic some-token,Some-Other-Header=some-value",
+	})
+	_, err := CreateWalletWithKeymanager(cliCtx.Context, &CreateWalletConfig{
+		WalletCfg: &wallet.Config{
+			WalletDir:      walletDir,
+			KeymanagerKind: keymanager.Imported,
+			WalletPassword: password,
+		},
+	})
+	require.NoError(t, err)
+	_, _, err = prepareClients(cliCtx)
+	require.NoError(t, err)
+	md, _ := metadata.FromOutgoingContext(cliCtx.Context)
+	assert.Equal(t, "Basic some-token", md.Get("Authorization")[0])
+	assert.Equal(t, "some-value", md.Get("Some-Other-Header")[0])
+}

validator/accounts/wallet_create_test.go

@@ -50,6 +50,7 @@ type testWalletConfig struct {
 	walletPasswordFile      string
 	accountPasswordFile     string
 	privateKeyFile          string
+	grpcHeaders             string
 	skipDepositConfirm      bool
 	numAccounts             int64
 	keymanagerKind          keymanager.Kind
@@ -76,6 +77,7 @@ func setupWalletCtx(
 	set.Int64(flags.NumAccountsFlag.Name, cfg.numAccounts, "")
 	set.Bool(flags.SkipDepositConfirmationFlag.Name, cfg.skipDepositConfirm, "")
 	set.Bool(flags.SkipMnemonic25thWordCheckFlag.Name, true, "")
+	set.String(flags.GrpcHeadersFlag.Name, cfg.grpcHeaders, "")
 
 	if cfg.privateKeyFile != "" {
 		set.String(flags.ImportPrivateKeyFileFlag.Name, cfg.privateKeyFile, "")
@@ -96,6 +98,7 @@ func setupWalletCtx(
 	assert.NoError(tb, set.Set(flags.AccountPasswordFileFlag.Name, cfg.accountPasswordFile))
 	assert.NoError(tb, set.Set(flags.NumAccountsFlag.Name, strconv.Itoa(int(cfg.numAccounts))))
 	assert.NoError(tb, set.Set(flags.SkipDepositConfirmationFlag.Name, strconv.FormatBool(cfg.skipDepositConfirm)))
+	assert.NoError(tb, set.Set(flags.GrpcHeadersFlag.Name, cfg.grpcHeaders))
 	return cli.NewContext(&app, set, nil)
 }