Add support for extra-headers to cortextool rules commands (#288)

Author: alexander-akhmetov

CHANGELOG.md

@@ -4,6 +4,8 @@ Order should be `CHANGE`, `FEATURE`, `ENHANCEMENT`, and `BUGFIX`
 
 ## unreleased/master
 
+* [FEATURE] Add `--extra-headers` support for `cortextool rules` commands. #288
+
 ## v0.11.0
 
 * [FEATURE] Support Arm64 on Darwin for all binaries (benchtool etc). #215

pkg/client/client.go

@@ -33,19 +33,21 @@ type Config struct {
 	Address         string `yaml:"address"`
 	ID              string `yaml:"id"`
 	TLS             tls.ClientConfig
-	UseLegacyRoutes bool   `yaml:"use_legacy_routes"`
-	AuthToken       string `yaml:"auth_token"`
+	UseLegacyRoutes bool              `yaml:"use_legacy_routes"`
+	AuthToken       string            `yaml:"auth_token"`
+	ExtraHeaders    map[string]string `yaml:"extra_headers"`
 }
 
 // CortexClient is used to get and load rules into a cortex ruler
 type CortexClient struct {
-	user      string
-	key       string
-	id        string
-	endpoint  *url.URL
-	Client    http.Client
-	apiPath   string
-	authToken string
+	user         string
+	key          string
+	id           string
+	endpoint     *url.URL
+	Client       http.Client
+	apiPath      string
+	authToken    string
+	extraHeaders map[string]string
 }
 
 // New returns a new Client
@@ -87,13 +89,14 @@ func New(cfg Config) (*CortexClient, error) {
 	}
 
 	return &CortexClient{
-		user:      cfg.User,
-		key:       cfg.Key,
-		id:        cfg.ID,
-		endpoint:  endpoint,
-		Client:    client,
-		apiPath:   path,
-		authToken: cfg.AuthToken,
+		user:         cfg.User,
+		key:          cfg.Key,
+		id:           cfg.ID,
+		endpoint:     endpoint,
+		Client:       client,
+		apiPath:      path,
+		authToken:    cfg.AuthToken,
+		extraHeaders: cfg.ExtraHeaders,
 	}, nil
 }
 
@@ -137,6 +140,10 @@ func (r *CortexClient) doRequest(ctx context.Context, path, method string, paylo
 		req.Header.Add("Authorization", "Bearer "+r.authToken)
 	}
 
+	for k, v := range r.extraHeaders {
+		req.Header.Add(k, v)
+	}
+
 	req.Header.Add("X-Scope-OrgID", r.id)
 
 	log.WithFields(log.Fields{

pkg/client/client_test.go

@@ -2,7 +2,9 @@ package client
 
 import (
 	"context"
+	"fmt"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"testing"
 
@@ -94,3 +96,98 @@ 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: "basic headers only, no extra headers",
+			id:   "my-tenant-id",
+			validate: func(t *testing.T, req *http.Request) {
+				require.Equal(t, "my-tenant-id", req.Header.Get("X-Scope-OrgID"))
+				// Verify no extra headers were added
+				require.Empty(t, req.Header.Get("key1"))
+			},
+		},
+		{
+			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"))
+			},
+		},
+		{
+			name:      "auth token with extra headers",
+			id:        "my-tenant-id",
+			authToken: "my-auth-token",
+			extraHeaders: map[string]string{
+				"Custom-Header": "custom-value",
+			},
+			validate: func(t *testing.T, req *http.Request) {
+				require.Equal(t, "Bearer my-auth-token", req.Header.Get("Authorization"))
+				require.Equal(t, "my-tenant-id", req.Header.Get("X-Scope-OrgID"))
+				require.Equal(t, "custom-value", req.Header.Get("Custom-Header"))
+			},
+		},
+		{
+			name:      "authorization header in extra headers is ignored when auth token is set",
+			id:        "my-tenant-id",
+			authToken: "my-auth-token",
+			extraHeaders: map[string]string{
+				"Authorization": "Bearer should-be-ignored",
+			},
+			validate: func(t *testing.T, req *http.Request) {
+				// The Authorization header from extraHeaders should be overwritten
+				require.Equal(t, "Bearer my-auth-token", req.Header.Get("Authorization"))
+			},
+		},
+	} {
+		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)
+
+			if tc.expectedErr != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tc.expectedErr)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, http.StatusOK, res.StatusCode)
+				req := <-requestCh
+				tc.validate(t, req)
+			}
+		})
+	}
+}

pkg/commands/rules.go

@@ -93,6 +93,8 @@ func (r *RuleCommand) Register(app *kingpin.Application) {
 	rulesCmd.Flag("user", "API user to use when contacting cortex, alternatively set CORTEX_API_USER. If empty, CORTEX_TENANT_ID will be used instead.").Default("").Envar("CORTEX_API_USER").StringVar(&r.ClientConfig.User)
 	rulesCmd.Flag("key", "API key to use when contacting cortex, alternatively set CORTEX_API_KEY.").Default("").Envar("CORTEX_API_KEY").StringVar(&r.ClientConfig.Key)
 	rulesCmd.Flag("backend", "Backend type to interact with: <cortex|loki>").Default("cortex").EnumVar(&r.Backend, backends...)
+	r.ClientConfig.ExtraHeaders = map[string]string{}
+	rulesCmd.Flag("extra-headers", "Extra headers to add to the requests in header=value format, alternatively set newline separated CORTEX_EXTRA_HEADERS.").Envar("CORTEX_EXTRA_HEADERS").StringMapVar(&r.ClientConfig.ExtraHeaders)
 
 	// Register rule commands
 	listCmd := rulesCmd.