Unify mimirtool authentication options and add extra-headers support for commands that depend on MimirClient

CHANGELOG.md

@@ -28,6 +28,7 @@
 ### Mimirtool
 
 * [BUGFIX] Fix issue where `MIMIR_HTTP_PREFIX` environment variable was ignored and the value from `MIMIR_MIMIR_HTTP_PREFIX` was used instead. #10207
+* [ENHANCEMENT] Unify mimirtool authentication options and add extra-headers support for commands that depend on MimirClient. #10178
 
 ### Mimir Continuous Test
 

docs/sources/mimir/manage/tools/mimirtool.md

@@ -70,12 +70,13 @@ chmod +x mimirtool
 
 For Mimirtools to interact with Grafana Mimir, Grafana Enterprise Metrics, Prometheus, or Grafana, set the following environment variables or CLI flags.
 
-| Environment variable | Flag        | Description                                                                                                                                                                                      |
-| -------------------- | ----------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `MIMIR_ADDRESS`      | `--address` | Sets the address of the API of the Grafana Mimir cluster.                                                                                                                                        |
-| `MIMIR_API_USER`     | `--user`    | Sets the basic auth username. If this variable is empty and `MIMIR_API_KEY` is set, the system uses `MIMIR_TENANT_ID` instead. If you're using Grafana Cloud, this variable is your instance ID. |
-| `MIMIR_API_KEY`      | `--key`     | Sets the basic auth password. If you're using Grafana Cloud, this variable is your API key.                                                                                                      |
-| `MIMIR_TENANT_ID`    | `--id`      | Sets the tenant ID of the Grafana Mimir instance that Mimirtools interacts with.                                                                                                                 |
+| Environment variable  | Flag              | Description                                                                                                                                                                                                |
+| --------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `MIMIR_ADDRESS`       | `--address`       | Sets the address of the API of the Grafana Mimir cluster.                                                                                                                                                  |
+| `MIMIR_API_USER`      | `--user`          | Sets the basic authentication username. If this variable is empty and `MIMIR_API_KEY` is set, the system uses `MIMIR_TENANT_ID` instead. If you're using Grafana Cloud, this variable is your instance ID. |
+| `MIMIR_API_KEY`       | `--key`           | Sets the basic authentication password. If you're using Grafana Cloud, this variable is your API key.                                                                                                      |
+| `MIMIR_TENANT_ID`     | `--id`            | Sets the tenant ID of the Grafana Mimir instance that Mimirtool interacts with.                                                                                                                            |
+| `MIMIR_EXTRA_HEADERS` | `--extra-headers` | Extra headers to add to the requests in header=value format. You can specify this flag multiple times. You must newline separate environment values.                                                       |
 
 It is also possible to set TLS-related options with the following environment variables or CLI flags:
 
@@ -309,10 +310,9 @@ For more information, refer to the [documentation of Mimirtool Github Action](ht
 
 Configuration options relevant to rules commands:
 
-| Flag              | Description                                                                               |
-| ----------------- | ----------------------------------------------------------------------------------------- |
-| `--auth-token`    | Authentication token for bearer token or JWT auth.                                        |
-| `--extra-headers` | Extra headers to add to the requests in header=value format. (Can specify multiple times) |
+| Flag           | Description                                        |
+| -------------- | -------------------------------------------------- |
+| `--auth-token` | Authentication token for bearer token or JWT auth. |
 
 #### List rules
 

pkg/mimirtool/client/client_test.go

@@ -8,7 +8,9 @@ package client
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"testing"
 
@@ -100,3 +102,114 @@ func TestBuildURL(t *testing.T) {
 	}
 
 }
+
+func TestDoRequest(t *testing.T) {
+	requestCh := make(chan *http.Request, 1)
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestCh <- r
+		fmt.Fprintln(w, "hello")
+	}))
+	defer ts.Close()
+
+	for _, tc := range []struct {
+		name         string
+		user         string
+		key          string
+		id           string
+		authToken    string
+		extraHeaders map[string]string
+		expectedErr  string
+		validate     func(t *testing.T, req *http.Request)
+	}{
+		{
+			name:        "errors because user, key and authToken are provided",
+			user:        "my-ba-user",
+			key:         "my-ba-password",
+			authToken:   "RandomJwt",
+			expectedErr: "at most one of basic auth or auth token should be configured",
+		},
+		{
+			name: "user provided so uses key as password",
+			user: "my-ba-user",
+			key:  "my-ba-password",
+			id:   "my-tenant-id",
+			validate: func(t *testing.T, req *http.Request) {
+				user, pass, ok := req.BasicAuth()
+				require.True(t, ok)
+				require.Equal(t, "my-ba-user", user)
+				require.Equal(t, "my-ba-password", pass)
+				require.Equal(t, "my-tenant-id", req.Header.Get("X-Scope-OrgID"))
+			},
+		},
+		{
+			name: "user not provided so uses id as username and key as password",
+			key:  "my-ba-password",
+			id:   "my-tenant-id",
+			validate: func(t *testing.T, req *http.Request) {
+				user, pass, ok := req.BasicAuth()
+				require.True(t, ok)
+				require.Equal(t, "my-tenant-id", user)
+				require.Equal(t, "my-ba-password", pass)
+				require.Equal(t, "my-tenant-id", req.Header.Get("X-Scope-OrgID"))
+			},
+		},
+		{
+			name:      "authToken is provided",
+			id:        "my-tenant-id",
+			authToken: "RandomJwt",
+			validate: func(t *testing.T, req *http.Request) {
+				require.Equal(t, "Bearer RandomJwt", req.Header.Get("Authorization"))
+				require.Equal(t, "my-tenant-id", req.Header.Get("X-Scope-OrgID"))
+			},
+		},
+		{
+			name: "no auth options and tenant are provided",
+			validate: func(t *testing.T, req *http.Request) {
+				require.Empty(t, req.Header.Get("Authorization"))
+				require.Empty(t, req.Header.Get("X-Scope-OrgID"))
+			},
+		},
+		{
+			name: "extraHeaders are added",
+			id:   "my-tenant-id",
+			extraHeaders: map[string]string{
+				"key1":          "value1",
+				"key2":          "value2",
+				"X-Scope-OrgID": "first-tenant-id",
+			},
+			validate: func(t *testing.T, req *http.Request) {
+				require.Equal(t, "value1", req.Header.Get("key1"))
+				require.Equal(t, "value2", req.Header.Get("key2"))
+				require.Equal(t, []string{"first-tenant-id", "my-tenant-id"}, req.Header.Values("X-Scope-OrgID"))
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			client, err := New(Config{
+				Address:      ts.URL,
+				User:         tc.user,
+				Key:          tc.key,
+				AuthToken:    tc.authToken,
+				ID:           tc.id,
+				ExtraHeaders: tc.extraHeaders,
+			})
+			require.NoError(t, err)
+
+			res, err := client.doRequest(ctx, "/test", http.MethodGet, nil, -1)
+
+			// Validate errors
+			if tc.expectedErr != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tc.expectedErr)
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, http.StatusOK, res.StatusCode)
+			req := <-requestCh
+			tc.validate(t, req)
+		})
+	}
+}

pkg/mimirtool/commands/alerts.go

@@ -95,6 +95,8 @@ func (a *AlertmanagerCommand) Register(app *kingpin.Application, envVars EnvVarN
 	for _, cmd := range []*kingpin.CmdClause{getAlertsCmd, deleteCmd, loadalertCmd} {
 		cmd.Flag("address", "Address of the Grafana Mimir cluster; alternatively, set "+envVars.Address+".").Envar(envVars.Address).Required().StringVar(&a.ClientConfig.Address)
 		cmd.Flag("id", "Grafana Mimir tenant ID; alternatively, set "+envVars.TenantID+". Used for X-Scope-OrgID HTTP header. Also used for basic auth if --user is not provided.").Envar(envVars.TenantID).Required().StringVar(&a.ClientConfig.ID)
+		a.ClientConfig.ExtraHeaders = map[string]string{}
+		cmd.Flag("extra-headers", "Extra headers to add to the requests in header=value format, alternatively set newline separated "+envVars.ExtraHeaders+".").Envar(envVars.ExtraHeaders).StringMapVar(&a.ClientConfig.ExtraHeaders)
 	}
 
 	migrateCmd := alertCmd.Command("migrate-utf8", "Migrate the Alertmanager tenant configuration for UTF-8.").Action(a.migrateConfig)
@@ -270,6 +272,8 @@ func (a *AlertCommand) Register(app *kingpin.Application, envVars EnvVarNames, r
 	alertCmd.Flag("user", fmt.Sprintf("Basic auth username to use when contacting Grafana Mimir, alternatively set %s. If empty, %s will be used instead. ", envVars.APIUser, envVars.TenantID)).Default("").Envar(envVars.APIUser).StringVar(&a.ClientConfig.User)
 	alertCmd.Flag("key", "Basic auth password to use when contacting Grafana Mimir; alternatively, set "+envVars.APIKey+".").Default("").Envar(envVars.APIKey).StringVar(&a.ClientConfig.Key)
 	alertCmd.Flag("auth-token", "Authentication token for bearer token or JWT auth, alternatively set "+envVars.AuthToken+".").Default("").Envar(envVars.AuthToken).StringVar(&a.ClientConfig.AuthToken)
+	a.ClientConfig.ExtraHeaders = map[string]string{}
+	alertCmd.Flag("extra-headers", "Extra headers to add to the requests in header=value format, alternatively set newline separated "+envVars.ExtraHeaders+".").Envar(envVars.ExtraHeaders).StringMapVar(&a.ClientConfig.ExtraHeaders)
 
 	verifyAlertsCmd := alertCmd.Command("verify", "Verifies whether or not alerts in an Alertmanager cluster are deduplicated; useful for verifying correct configuration when transferring from Prometheus to Grafana Mimir alert evaluation.").Action(a.verifyConfig)
 	verifyAlertsCmd.Flag("ignore-alerts", "A comma separated list of Alert names to ignore in deduplication checks.").StringVar(&a.IgnoreString)

pkg/mimirtool/commands/analyse.go

@@ -6,6 +6,7 @@
 package commands
 
 import (
+	"fmt"
 	"runtime"
 	"strconv"
 
@@ -30,9 +31,13 @@ func (cmd *AnalyzeCommand) Register(app *kingpin.Application, envVars EnvVarName
 		Default("").
 		Envar(envVars.AuthToken).
 		StringVar(&paCmd.authToken)
-	prometheusAnalyzeCmd.Flag("id", "Basic auth username to use when contacting Prometheus or Grafana Mimir, also set as tenant ID; alternatively, set "+envVars.TenantID+".").
+	prometheusAnalyzeCmd.Flag("id", "Grafana Mimir tenant ID; alternatively, set "+envVars.TenantID+". Used for X-Scope-OrgID HTTP header. Also used for basic auth if --user is not provided.").
 		Envar(envVars.TenantID).
 		Default("").
+		StringVar(&paCmd.tenantID)
+	prometheusAnalyzeCmd.Flag("user", fmt.Sprintf("Basic auth API user to use when contacting Prometheus or Grafana Mimir; alternatively, set %s. If empty, %s is used instead.", envVars.APIUser, envVars.TenantID)).
+		Envar(envVars.APIUser).
+		Default("").
 		StringVar(&paCmd.username)
 	prometheusAnalyzeCmd.Flag("key", "Basic auth password to use when contacting Prometheus or Grafana Mimir; alternatively, set "+envVars.APIKey+"").
 		Envar(envVars.APIKey).
@@ -81,10 +86,14 @@ func (cmd *AnalyzeCommand) Register(app *kingpin.Application, envVars EnvVarName
 		Envar(envVars.Address).
 		Required().
 		StringVar(&raCmd.ClientConfig.Address)
-	rulerAnalyzeCmd.Flag("id", "Basic auth username and X-Scope-OrgID value to use when contacting Prometheus or Grafana Mimir; alternatively, set "+envVars.TenantID+".").
+	rulerAnalyzeCmd.Flag("id", "Mimir tenant id, alternatively set "+envVars.TenantID+". Used for X-Scope-OrgID HTTP header. Also used for basic auth if --user is not provided.").
 		Envar(envVars.TenantID).
 		Default("").
 		StringVar(&raCmd.ClientConfig.ID)
+	rulerAnalyzeCmd.Flag("user", fmt.Sprintf("Basic auth username to use when contacting Prometheus or Grafana Mimir, alternatively set %s. If empty, %s will be used instead. ", envVars.APIUser, envVars.TenantID)).
+		Envar(envVars.APIUser).
+		Default("").
+		StringVar(&raCmd.ClientConfig.User)
 	rulerAnalyzeCmd.Flag("key", "Basic auth password to use when contacting Prometheus or Grafana Mimir; alternatively, set "+envVars.APIKey+".").
 		Envar(envVars.APIKey).
 		Default("").
@@ -96,6 +105,10 @@ func (cmd *AnalyzeCommand) Register(app *kingpin.Application, envVars EnvVarName
 		Default("").
 		Envar(envVars.AuthToken).
 		StringVar(&raCmd.ClientConfig.AuthToken)
+	raCmd.ClientConfig.ExtraHeaders = map[string]string{}
+	rulerAnalyzeCmd.Flag("extra-headers", "Extra headers to add to the requests in header=value format, alternatively set newline separated "+envVars.ExtraHeaders+".").
+		Envar(envVars.ExtraHeaders).
+		StringMapVar(&raCmd.ClientConfig.ExtraHeaders)
 
 	daCmd := &DashboardAnalyzeCommand{}
 	dashboardAnalyzeCmd := analyzeCmd.Command("dashboard", "Analyze and output the metrics used in Grafana dashboard files").Action(daCmd.run)

pkg/mimirtool/commands/analyse_prometheus.go

@@ -34,6 +34,7 @@ type PrometheusAnalyzeCommand struct {
 	address              string
 	prometheusHTTPPrefix string
 	username             string
+	tenantID             string
 	password             string
 	authToken            string
 	readTimeout          time.Duration
@@ -90,14 +91,21 @@ func (cmd *PrometheusAnalyzeCommand) parseUsedMetrics() (model.LabelValues, erro
 func (cmd *PrometheusAnalyzeCommand) newAPI() (v1.API, error) {
 	rt := api.DefaultRoundTripper
 	rt = config.NewUserAgentRoundTripper(client.UserAgent(), rt)
-	if cmd.authToken != "" {
-		rt = config.NewAuthorizationCredentialsRoundTripper("Bearer", config.NewInlineSecret(cmd.authToken), rt)
-	} else if cmd.username != "" {
-		rt = &setTenantIDTransport{
-			RoundTripper: rt,
-			tenantID:     cmd.username,
-		}
+
+	switch {
+	case cmd.username != "":
 		rt = config.NewBasicAuthRoundTripper(config.NewInlineSecret(cmd.username), config.NewInlineSecret(cmd.password), rt)
+
+	case cmd.password != "":
+		rt = config.NewBasicAuthRoundTripper(config.NewInlineSecret(cmd.tenantID), config.NewInlineSecret(cmd.password), rt)
+
+	case cmd.authToken != "":
+		rt = config.NewAuthorizationCredentialsRoundTripper("Bearer", config.NewInlineSecret(cmd.authToken), rt)
+	}
+
+	rt = &setTenantIDTransport{
+		RoundTripper: rt,
+		tenantID:     cmd.tenantID,
 	}
 
 	address, err := url.JoinPath(cmd.address, cmd.prometheusHTTPPrefix)

pkg/mimirtool/commands/backfill.go

@@ -69,6 +69,11 @@ func (c *BackfillCommand) Register(app *kingpin.Application, envVars EnvVarNames
 		Envar(envVars.APIKey).
 		StringVar(&c.clientConfig.Key)
 
+	c.clientConfig.ExtraHeaders = map[string]string{}
+	cmd.Flag("extra-headers", "Extra headers to add to the requests in header=value format, alternatively set newline separated "+envVars.ExtraHeaders+".").
+		Envar(envVars.ExtraHeaders).
+		StringMapVar(&c.clientConfig.ExtraHeaders)
+
 	cmd.Flag("tls-ca-path", "TLS CA certificate to verify Grafana Mimir API as part of mTLS; alternatively, set "+envVars.TLSCAPath+".").
 		Default("").
 		Envar(envVars.TLSCAPath).