diff --git a/pkg/cmd/roachprod/BUILD.bazel b/pkg/cmd/roachprod/BUILD.bazel index 556b0e350f25..edd51e48081c 100644 --- a/pkg/cmd/roachprod/BUILD.bazel +++ b/pkg/cmd/roachprod/BUILD.bazel @@ -19,6 +19,7 @@ go_library( "//pkg/roachprod/fluentbit", "//pkg/roachprod/install", "//pkg/roachprod/opentelemetry", + "//pkg/roachprod/promhelperclient", "//pkg/roachprod/ssh", "//pkg/roachprod/ui", "//pkg/roachprod/vm", diff --git a/pkg/cmd/roachprod/grafana/BUILD.bazel b/pkg/cmd/roachprod/grafana/BUILD.bazel index 403485b37838..9132b0464d07 100644 --- a/pkg/cmd/roachprod/grafana/BUILD.bazel +++ b/pkg/cmd/roachprod/grafana/BUILD.bazel @@ -22,6 +22,7 @@ go_library( importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachprod/grafana", visibility = ["//visibility:public"], deps = [ + "//pkg/roachprod/promhelperclient", "//pkg/util/httputil", "@com_github_cockroachdb_errors//:errors", "@com_github_go_openapi_strfmt//:strfmt", diff --git a/pkg/cmd/roachprod/grafana/annotations.go b/pkg/cmd/roachprod/grafana/annotations.go index 55788da13c89..5387a45e7acf 100644 --- a/pkg/cmd/roachprod/grafana/annotations.go +++ b/pkg/cmd/roachprod/grafana/annotations.go @@ -13,9 +13,9 @@ package grafana import ( "context" "fmt" - "os" "strings" + "github.com/cockroachdb/cockroach/pkg/roachprod/promhelperclient" "github.com/cockroachdb/cockroach/pkg/util/httputil" "github.com/cockroachdb/errors" "github.com/go-openapi/strfmt" @@ -24,9 +24,6 @@ import ( "google.golang.org/api/idtoken" ) -const ServiceAccountJson = "GRAFANA_SERVICE_ACCOUNT_JSON" -const ServiceAccountAudience = "GRAFANA_SERVICE_ACCOUNT_AUDIENCE" - // newGrafanaClient is a helper function that creates an HTTP client to // create grafana api calls with. If secure is true, it tries to get a // GCS identity token by using the service account specified by the env @@ -42,25 +39,15 @@ func newGrafanaClient( scheme = "https" // Read in the service account key and audience, so we can retrieve the identity token. - grafanaKey := os.Getenv(ServiceAccountJson) - if grafanaKey == "" { - return nil, errors.Newf("%s env variable was not found", ServiceAccountJson) - } - grafanaAudience := os.Getenv(ServiceAccountAudience) - if grafanaAudience == "" { - return nil, errors.Newf("%s env variable was not found", ServiceAccountAudience) + if _, err := promhelperclient.SetPromHelperCredsEnv(ctx, false); err != nil { + return nil, err } - ts, err := idtoken.NewTokenSource(ctx, grafanaAudience, idtoken.WithCredentialsJSON([]byte(grafanaKey))) + token, err := promhelperclient.GetToken(ctx, idtoken.NewTokenSource) if err != nil { - return nil, errors.Wrap(err, "Error creating GCS oauth token source from specified credential") + return nil, err } - token, err := ts.Token() - if err != nil { - return nil, errors.Wrap(err, "Error getting identity token") - } - - headers["Authorization"] = fmt.Sprintf("Bearer %s", token.AccessToken) + headers["Authorization"] = fmt.Sprintf("Bearer %s", token) } headers[httputil.ContentTypeHeader] = httputil.JSONContentType diff --git a/pkg/cmd/roachprod/main.go b/pkg/cmd/roachprod/main.go index b1c8804eb8d0..02bf4516dbe4 100644 --- a/pkg/cmd/roachprod/main.go +++ b/pkg/cmd/roachprod/main.go @@ -31,6 +31,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/roachprod/config" rperrors "github.com/cockroachdb/cockroach/pkg/roachprod/errors" "github.com/cockroachdb/cockroach/pkg/roachprod/install" + "github.com/cockroachdb/cockroach/pkg/roachprod/promhelperclient" "github.com/cockroachdb/cockroach/pkg/roachprod/ui" "github.com/cockroachdb/cockroach/pkg/roachprod/vm" "github.com/cockroachdb/cockroach/pkg/roachprod/vm/gce" @@ -1358,7 +1359,7 @@ creates an annotation over time range. Example: # Create an annotation over time range 1-100 on the centralized grafana instance, which needs authentication. roachprod grafana-annotation grafana.testeng.crdb.io example-annotation-event --tags my-cluster --tags test-run-1 --dashboard-uid overview --time-range 1,100 -`, grafana.ServiceAccountJson, grafana.ServiceAccountAudience), +`, promhelperclient.ServiceAccountJson, promhelperclient.ServiceAccountAudience), Args: cobra.ExactArgs(2), Run: wrap(func(cmd *cobra.Command, args []string) error { req := grafana.AddAnnotationRequest{ diff --git a/pkg/roachprod/promhelperclient/client.go b/pkg/roachprod/promhelperclient/client.go index b864be16f792..7076d13e0db6 100644 --- a/pkg/roachprod/promhelperclient/client.go +++ b/pkg/roachprod/promhelperclient/client.go @@ -20,7 +20,6 @@ import ( "fmt" "io" "net/http" - "os" "strconv" "strings" @@ -37,8 +36,8 @@ import ( const ( resourceName = "instance-configs" resourceVersion = "v1" - serviceAccountJson = "PROM_HELPER_SERVICE_ACCOUNT_JSON" - serviceAccountAudience = "PROM_HELPER_SERVICE_ACCOUNT_AUDIENCE" + ServiceAccountJson = "PROM_HELPER_SERVICE_ACCOUNT_JSON" + ServiceAccountAudience = "PROM_HELPER_SERVICE_ACCOUNT_AUDIENCE" ) // SupportedPromProjects are the projects supported for prometheus target @@ -228,18 +227,14 @@ func (c *PromClient) getToken( return "", nil } // Read in the service account key and audience, so we can retrieve the identity token. - if _, err := setPromHelperCredsEnv(ctx, forceFetchCreds, l); err != nil { + if credSrc, err := SetPromHelperCredsEnv(ctx, forceFetchCreds); err != nil { return "", err + } else { + l.Printf("Prometheus credentials obtained from %s", credSrc) } - grafanaKey := os.Getenv(serviceAccountJson) - grafanaAudience := os.Getenv(serviceAccountAudience) - ts, err := c.newTokenSource(ctx, grafanaAudience, idtoken.WithCredentialsJSON([]byte(grafanaKey))) + token, err := GetToken(ctx, c.newTokenSource) if err != nil { - return "", errors.Wrap(err, "error creating GCS oauth token source from specified credential") - } - token, err := ts.Token() - if err != nil { - return "", errors.Wrap(err, "error getting identity token") + return "", err } - return fmt.Sprintf("Bearer %s", token.AccessToken), nil + return fmt.Sprintf("Bearer %s", token), nil } diff --git a/pkg/roachprod/promhelperclient/client_test.go b/pkg/roachprod/promhelperclient/client_test.go index e5c927067862..13f36b41537b 100644 --- a/pkg/roachprod/promhelperclient/client_test.go +++ b/pkg/roachprod/promhelperclient/client_test.go @@ -150,9 +150,9 @@ func Test_getToken(t *testing.T) { require.Empty(t, token) }) t.Run("invalid credentials", func(t *testing.T) { - err := os.Setenv(serviceAccountJson, "{}") + err := os.Setenv(ServiceAccountJson, "{}") require.Nil(t, err) - err = os.Setenv(serviceAccountAudience, "dummy_audience") + err = os.Setenv(ServiceAccountAudience, "dummy_audience") require.Nil(t, err) c.newTokenSource = func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error) { return nil, fmt.Errorf("invalid") @@ -164,9 +164,9 @@ func Test_getToken(t *testing.T) { require.Equal(t, "error creating GCS oauth token source from specified credential: invalid", err.Error()) }) t.Run("invalid token", func(t *testing.T) { - err := os.Setenv(serviceAccountJson, "{}") + err := os.Setenv(ServiceAccountJson, "{}") require.Nil(t, err) - err = os.Setenv(serviceAccountAudience, "dummy_audience") + err = os.Setenv(ServiceAccountAudience, "dummy_audience") require.Nil(t, err) c.newTokenSource = func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error) { return &mockToken{token: "", err: fmt.Errorf("failed")}, nil @@ -178,9 +178,9 @@ func Test_getToken(t *testing.T) { require.Equal(t, "error getting identity token: failed", err.Error()) }) t.Run("success", func(t *testing.T) { - err := os.Setenv(serviceAccountJson, "{}") + err := os.Setenv(ServiceAccountJson, "{}") require.Nil(t, err) - err = os.Setenv(serviceAccountAudience, "dummy_audience") + err = os.Setenv(ServiceAccountAudience, "dummy_audience") require.Nil(t, err) c.newTokenSource = func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error) { return &mockToken{token: "token"}, nil diff --git a/pkg/roachprod/promhelperclient/promhelper_utils.go b/pkg/roachprod/promhelperclient/promhelper_utils.go index 15010074ff4c..c3f5ee6e682d 100644 --- a/pkg/roachprod/promhelperclient/promhelper_utils.go +++ b/pkg/roachprod/promhelperclient/promhelper_utils.go @@ -12,14 +12,15 @@ package promhelperclient import ( "context" - "fmt" "io" "os" "path/filepath" "strings" "cloud.google.com/go/storage" - "github.com/cockroachdb/cockroach/pkg/roachprod/logger" + "github.com/cockroachdb/errors" + "golang.org/x/oauth2" + "google.golang.org/api/idtoken" "google.golang.org/api/option" ) @@ -47,9 +48,9 @@ const ( objectLocation = "promhelpers-secrets" ) -// setPromHelperCredsEnv sets the environment variables serviceAccountAudience and -// serviceAccountJson based on the following conditions: -// > check if forFetch is false +// SetPromHelperCredsEnv sets the environment variables ServiceAccountAudience and +// ServiceAccountJson based on the following conditions: +// > check if forceFetch is false // // > forceFetch is false // > if env is set return @@ -58,64 +59,76 @@ const ( // > read the creds from secrets manager // // > set the env variable and save the creds to the promCredFile -func setPromHelperCredsEnv( - ctx context.Context, forceFetch bool, l *logger.Logger, -) (FetchedFrom, error) { +func SetPromHelperCredsEnv(ctx context.Context, forceFetch bool) (FetchedFrom, error) { creds := "" fetchedFrom := Env if !forceFetch { // bypass environment and creds file if forceFetch is false // check if environment is set - audience := os.Getenv(serviceAccountAudience) - saJson := os.Getenv(serviceAccountJson) + audience := os.Getenv(ServiceAccountAudience) + saJson := os.Getenv(ServiceAccountJson) if audience != "" && saJson != "" { - l.Printf("Secrets obtained from environment.") return fetchedFrom, nil } // check if the secrets file is available b, err := os.ReadFile(promCredFile) if err == nil { - l.Printf("Secrets obtained from temp file: %s", promCredFile) creds = string(b) fetchedFrom = File } } if creds == "" { // creds == "" means (env is not set and the file does not have the creds) or forceFetch is true - l.Printf("creds need to be fetched from store.") options := []option.ClientOption{option.WithScopes(storage.ScopeReadOnly)} cj := os.Getenv("GOOGLE_EPHEMERAL_CREDENTIALS") - if len(cj) != 0 { - options = append(options, option.WithCredentialsJSON([]byte(cj))) - } else { - l.Printf("GOOGLE_EPHEMERAL_CREDENTIALS env is not set.") + if cj == "" { + return "", errors.New("SetPromHelperCredsEnv: GOOGLE_EPHEMERAL_CREDENTIALS env is not set") } + options = append(options, option.WithCredentialsJSON([]byte(cj))) + client, err := storage.NewClient(ctx, options...) if err != nil { - return fetchedFrom, err + return "", err } defer func() { _ = client.Close() }() fetchedFrom = Store obj := client.Bucket(bucket).Object(objectLocation) r, err := obj.NewReader(ctx) if err != nil { - return fetchedFrom, err + return "", err } defer func() { _ = r.Close() }() body, err := io.ReadAll(r) creds = string(body) if err != nil { - return fetchedFrom, err - } - err = os.WriteFile(promCredFile, body, 0700) - if err != nil { - l.Errorf("error writing to the credential file: %v", err) + return "", err } + _ = os.WriteFile(promCredFile, body, 0700) } secretValues := strings.Split(creds, secretsDelimiter) if len(secretValues) == 2 { - _ = os.Setenv(serviceAccountAudience, secretValues[0]) - _ = os.Setenv(serviceAccountJson, secretValues[1]) + _ = os.Setenv(ServiceAccountAudience, secretValues[0]) + _ = os.Setenv(ServiceAccountJson, secretValues[1]) return fetchedFrom, nil } - return fetchedFrom, fmt.Errorf("invalid secret values - %s", creds) + return "", errors.Newf("invalid secret values - %s", creds) +} + +// GetToken returns a GCS oauth token based on the service account key and audience +// set through the ServiceAccountJson and ServiceAccountAudience env vars. +// Assumes that the env vars have been set already, i.e. through SetPromHelperCredsEnv. +func GetToken( + ctx context.Context, + tokenSource func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error), +) (string, error) { + key := os.Getenv(ServiceAccountJson) + audience := os.Getenv(ServiceAccountAudience) + ts, err := tokenSource(ctx, audience, idtoken.WithCredentialsJSON([]byte(key))) + if err != nil { + return "", errors.Wrap(err, "error creating GCS oauth token source from specified credential") + } + token, err := ts.Token() + if err != nil { + return "", errors.Wrap(err, "error getting identity token") + } + return token.AccessToken, nil }