vmalert: allow configuring custom headers per group (#2901)

See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/2860

Signed-off-by: hagen1778 <roman@victoriametrics.com>

app/vmalert/README.md

@@ -125,6 +125,16 @@ name: <string>
 params:
   [ <string>: [<string>, ...]]
 
+# Optional list of HTTP headers in form `header-name: value`
+# applied for all rules requests within a group
+# For example:
+#  headers:
+#    - "CustomHeader: foo"
+#    - "CustomHeader2: bar"
+# Headers set via this param have priority over headers set via `-datasource.headers` flag. 
+headers:
+  [ <string>, ...]
+
 # Optional list of labels added to every rule within a group.
 # It has priority over the external labels.
 # Labels are commonly used for adding environment

app/vmalert/alerting.go

@@ -75,6 +75,7 @@ func newAlertingRule(qb datasource.QuerierBuilder, group *Group, cfg config.Rule
 			DataSourceType:     &group.Type,
 			EvaluationInterval: group.Interval,
 			QueryParams:        group.Params,
+			Headers:            group.Headers,
 		}),
 		alerts:  make(map[uint64]*notifier.Alert),
 		metrics: &alertingRuleMetrics{},

app/vmalert/config/config.go

@@ -38,6 +38,8 @@ type Group struct {
 	Checksum string
 	// Optional HTTP URL parameters added to each rule request
 	Params url.Values `yaml:"params"`
+	// Headers contains optional HTTP headers added to each rule request
+	Headers []datasource.Header `yaml:"headers,omitempty"`
 
 	// Catches all undefined fields and must be empty after parsing.
 	XXX map[string]interface{} `yaml:",inline"`

app/vmalert/config/config_test.go

@@ -60,6 +60,10 @@ func TestParseBad(t *testing.T) {
 			[]string{"testdata/rules/rules1-bad.rules"},
 			"bad graphite expr",
 		},
+		{
+			[]string{"testdata/dir/rules6-bad.rules"},
+			"missing ':' in header",
+		},
 	}
 	for _, tc := range testCases {
 		_, err := Parse(tc.path, true, true)
@@ -505,6 +509,24 @@ rules:
 `, `
 name: TestGroup
 limit: 10
+rules:
+  - alert: foo
+    expr: sum by(job) (up == 1)
+`)
+	})
+
+	t.Run("`headers` change", func(t *testing.T) {
+		f(t, `
+name: TestGroup
+headers:
+  - "TenantID: foo"
+rules:
+  - alert: foo
+    expr: sum by(job) (up == 1)
+`, `
+name: TestGroup
+headers:
+  - "TenantID: bar"
 rules:
   - alert: foo
     expr: sum by(job) (up == 1)

app/vmalert/config/testdata/dir/rules6-bad.rules

@@ -0,0 +1,7 @@
+groups:
+  - name: group
+    headers:
+    - 'foobar'
+    rules:
+      - alert: rows
+        expr: vm_rows > 0

app/vmalert/config/testdata/rules/rules2-good.rules

@@ -3,9 +3,10 @@ groups:
     interval: 2s
     concurrency: 2
     limit: 1000
+    headers:
+      - "MyHeader: foo"
     params:
       denyPartialResponse: ["true"]
-      extra_label: ["env=dev"]
     rules:
       - alert: Conns
         expr: sum(vm_tcplistener_conns) by(instance) > 1

app/vmalert/datasource/datasource.go

@@ -22,6 +22,7 @@ type QuerierParams struct {
 	DataSourceType     *Type
 	EvaluationInterval time.Duration
 	QueryParams        url.Values
+	Headers            []Header
 }
 
 // Metric is the basic entity which should be return by datasource

app/vmalert/datasource/init.go

@@ -15,7 +15,7 @@ var (
 		"E.g. http://127.0.0.1:8428 . See also -remoteRead.disablePathAppend")
 	appendTypePrefix = flag.Bool("datasource.appendTypePrefix", false, "Whether to add type prefix to -datasource.url based on the query type. Set to true if sending different query types to the vmselect URL.")
 
-	headers = flag.String("datasource.headers", "", "Optional HTTP headers to send with each request to the corresponding -datasource.url. "+
+	headers = flag.String("datasource.headers", "", "Optional HTTP extraHeaders to send with each request to the corresponding -datasource.url. "+
 		"For example, -datasource.headers='My-Auth:foobar' would send 'My-Auth: foobar' HTTP header with every request to the corresponding -datasource.url. "+
 		"Multiple headers must be delimited by '^^': -datasource.headers='header1:value1^^header2:value2'")
 

app/vmalert/datasource/type.go

@@ -2,6 +2,7 @@ package datasource
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vmselect/graphiteql"
 	"github.com/VictoriaMetrics/metricsql"
@@ -89,3 +90,27 @@ func (t *Type) UnmarshalYAML(unmarshal func(interface{}) error) error {
 func (t Type) MarshalYAML() (interface{}, error) {
 	return t.name, nil
 }
+
+// Header is a Key - Value struct for holding an HTTP header.
+type Header struct {
+	Key   string
+	Value string
+}
+
+// UnmarshalYAML implements the yaml.Unmarshaler interface.
+func (h *Header) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var s string
+	if err := unmarshal(&s); err != nil {
+		return err
+	}
+	if s == "" {
+		return nil
+	}
+	n := strings.IndexByte(s, ':')
+	if n < 0 {
+		return fmt.Errorf(`missing ':' in header %q; expecting "key: value" format`, s)
+	}
+	h.Key = strings.TrimSpace(s[:n])
+	h.Value = strings.TrimSpace(s[n+1:])
+	return nil
+}

app/vmalert/datasource/vm.go

@@ -24,6 +24,7 @@ type VMStorage struct {
 	dataSourceType     Type
 	evaluationInterval time.Duration
 	extraParams        url.Values
+	extraHeaders       []Header
 }
 
 // Clone makes clone of VMStorage, shares http client.
@@ -46,6 +47,7 @@ func (s *VMStorage) ApplyParams(params QuerierParams) *VMStorage {
 	}
 	s.evaluationInterval = params.EvaluationInterval
 	s.extraParams = params.QueryParams
+	s.extraHeaders = params.Headers
 	return s
 }
 
@@ -148,5 +150,8 @@ func (s *VMStorage) newRequestPOST() (*http.Request, error) {
 	if s.authCfg != nil {
 		s.authCfg.SetHeaders(req, true)
 	}
+	for _, h := range s.extraHeaders {
+		req.Header.Set(h.Key, h.Value)
+	}
 	return req, nil
 }

app/vmalert/datasource/vm_test.go

@@ -487,7 +487,7 @@ func TestRequestParams(t *testing.T) {
 	}
 }
 
-func TestAuthConfig(t *testing.T) {
+func TestHeaders(t *testing.T) {
 	var testCases = []struct {
 		name    string
 		vmFn    func() *VMStorage
@@ -527,6 +527,40 @@ func TestAuthConfig(t *testing.T) {
 				checkEqualString(t, "foo", token)
 			},
 		},
+		{
+			name: "custom extraHeaders",
+			vmFn: func() *VMStorage {
+				return &VMStorage{extraHeaders: []Header{
+					{Key: "Foo", Value: "bar"},
+					{Key: "Baz", Value: "qux"},
+				}}
+			},
+			checkFn: func(t *testing.T, r *http.Request) {
+				h1 := r.Header.Get("Foo")
+				checkEqualString(t, "bar", h1)
+				h2 := r.Header.Get("Baz")
+				checkEqualString(t, "qux", h2)
+			},
+		},
+		{
+			name: "custom header overrides basic auth",
+			vmFn: func() *VMStorage {
+				cfg, err := utils.AuthConfig(utils.WithBasicAuth("foo", "bar", ""))
+				if err != nil {
+					t.Errorf("Error get auth config: %s", err)
+				}
+				return &VMStorage{
+					authCfg: cfg,
+					extraHeaders: []Header{
+						{Key: "Authorization", Value: "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="},
+					}}
+			},
+			checkFn: func(t *testing.T, r *http.Request) {
+				u, p, _ := r.BasicAuth()
+				checkEqualString(t, "Aladdin", u)
+				checkEqualString(t, "open sesame", p)
+			},
+		},
 	}
 	for _, tt := range testCases {
 		t.Run(tt.name, func(t *testing.T) {

app/vmalert/group.go

@@ -35,8 +35,9 @@ type Group struct {
 	Checksum       string
 	LastEvaluation time.Time
 
-	Labels map[string]string
+	Labels  map[string]string
-	Params url.Values
+	Params  url.Values
+	Headers []datasource.Header
 
 	doneCh     chan struct{}
 	finishedCh chan struct{}
@@ -96,6 +97,7 @@ func newGroup(cfg config.Group, qb datasource.QuerierBuilder, defaultInterval ti
 		Concurrency: cfg.Concurrency,
 		Checksum:    cfg.Checksum,
 		Params:      cfg.Params,
+		Headers:     cfg.Headers,
 		Labels:      cfg.Labels,
 
 		doneCh:     make(chan struct{}),
@@ -217,6 +219,7 @@ func (g *Group) updateWith(newGroup *Group) error {
 	g.Type = newGroup.Type
 	g.Concurrency = newGroup.Concurrency
 	g.Params = newGroup.Params
+	g.Headers = newGroup.Headers
 	g.Labels = newGroup.Labels
 	g.Limit = newGroup.Limit
 	g.Checksum = newGroup.Checksum

app/vmalert/manager.go

@@ -170,6 +170,7 @@ func (g *Group) toAPI() APIGroup {
 		LastEvaluation: g.LastEvaluation,
 		Concurrency:    g.Concurrency,
 		Params:         urlValuesToStrings(g.Params),
+		Headers:        headersToStrings(g.Headers),
 		Labels:         g.Labels,
 	}
 	for _, r := range g.Rules {
@@ -198,3 +199,14 @@ func urlValuesToStrings(values url.Values) []string {
 	}
 	return res
 }
+
+func headersToStrings(headers []datasource.Header) []string {
+	if len(headers) < 1 {
+		return nil
+	}
+	var res []string
+	for _, h := range headers {
+		res = append(res, fmt.Sprintf("%s: %s", h.Key, h.Value))
+	}
+	return res
+}

app/vmalert/recording.go

@@ -73,6 +73,7 @@ func newRecordingRule(qb datasource.QuerierBuilder, group *Group, cfg config.Rul
 			DataSourceType:     &group.Type,
 			EvaluationInterval: group.Interval,
 			QueryParams:        group.Params,
+			Headers:            group.Headers,
 		}),
 	}
 

app/vmalert/web.qtpl

@@ -57,6 +57,13 @@
                     {% endfor %}
                     </div>
                 {% endif %}
+                {% if len(g.Headers) > 0 %}
+                    <div class="fs-6 fw-lighter">Extra headers
+                    {% for _, header := range g.Headers %}
+                            <span class="float-left badge bg-primary">{%s header %}</span>
+                    {% endfor %}
+                    </div>
+                {% endif %}
             </div>
             <div class="collapse" id="rules-{%s g.ID %}">
                 <table class="table table-striped table-hover table-sm">

app/vmalert/web_types.go

@@ -70,6 +70,8 @@ type APIGroup struct {
 	Concurrency int `json:"concurrency"`
 	// Params contains HTTP URL parameters added to each Rule's request
 	Params []string `json:"params,omitempty"`
+	// Headers contains HTTP headers added to each Rule's request
+	Headers []string `json:"headers,omitempty"`
 	// Labels is a set of label value pairs, that will be added to every rule.
 	Labels map[string]string `json:"labels,omitempty"`
 }

docs/CHANGELOG.md

@@ -15,7 +15,7 @@ The following tip changes can be tested by building VictoriaMetrics components f
 
 ## tip
 
-* FEATURE: [vmalert](https://docs.victoriametrics.com/vmalert.html): allow configuring additional headers for `datasource.url`, `remoteWrite.url` and `remoteRead.url` URLs. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/2860) for details.
+* FEATURE: [vmalert](https://docs.victoriametrics.com/vmalert.html): allow configuring additional headers for `datasource.url`, `remoteWrite.url` and `remoteRead.url` URLs. Headers also can be set on group level via `headers` param. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/2860) for details.
 * FEATURE: [MetricsQL](https://docs.victoriametrics.com/MetricsQL.html): execute left and right sides of certain operations in parallel. For example, `q1 or q2`, `aggr_func(q1) <op> q2`, `q1 <op> aggr_func(q1)`. This may improve query performance if VictoriaMetrics has enough free resources for parallel processing of both sides of the operation. See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/2886).
 * FEATURE: [vmauth](https://docs.victoriametrics.com/vmagent.html): allow duplicate username records with different passwords at configuration file. It should allow password rotation without username change.
 

docs/vmalert.md

@@ -129,6 +129,16 @@ name: <string>
 params:
   [ <string>: [<string>, ...]]
 
+# Optional list of HTTP headers in form `header-name: value`
+# applied for all rules requests within a group
+# For example:
+#  headers:
+#    - "CustomHeader: foo"
+#    - "CustomHeader2: bar"
+# Headers set via this param have priority over headers set via `-datasource.headers` flag. 
+headers:
+  [ <string>, ...]
+
 # Optional list of labels added to every rule within a group.
 # It has priority over the external labels.
 # Labels are commonly used for adding environment

lib/promauth/config.go

@@ -308,7 +308,7 @@ func (ac *Config) HeadersNoAuthString() string {
 	return strings.Join(a, "")
 }
 
-// SetHeaders sets the configuted ac headers to req.
+// SetHeaders sets the configured ac headers to req.
 func (ac *Config) SetHeaders(req *http.Request, setAuthHeader bool) {
 	reqHeaders := req.Header
 	for _, h := range ac.headers {