Add CLI flag to add additional HTTP headers

Author: martialblog

.golangci.yml

@@ -24,6 +24,7 @@ linters:
     - tagliatelle
     - varnamelen
     - wrapcheck
+    - funlen
   exclusions:
     generated: lax
     presets:

README.md

@@ -26,6 +26,8 @@ Flags:
       --cert-file string   Specify the Certificate File for TLS authentication (CHECK_PROMETHEUS_CERT_FILE)
       --key-file string    Specify the Key File for TLS authentication (CHECK_PROMETHEUS_KEY_FILE)
   -t, --timeout int        Timeout in seconds for the CheckPlugin (default 30)
+      --header strings     Additional HTTP header to include in the request. Can be used multiple times.
+                           Keys and values are separated by a colon (--header "X-Custom: example").
   -h, --help               help for check_prometheus
   -v, --version            version for check_prometheus
 ```

cmd/config.go

@@ -23,6 +23,7 @@ type Config struct {
 	KeyFile   string `env:"CHECK_PROMETHEUS_KEY_FILE"`
 	Hostname  string `env:"CHECK_PROMETHEUS_HOSTNAME"`
 	URL       string `env:"CHECK_PROMETHEUS_URL"`
+	Headers   []string
 	Port      int
 	Info      bool
 	Insecure  bool
@@ -106,6 +107,20 @@ func (c *Config) NewClient() *client.Client {
 		rt = config.NewBasicAuthRoundTripper(u, p, rt)
 	}
 
+	// If extra headers are set, parse them and add them to the request
+	if len(c.Headers) > 0 {
+		headers := make(map[string]string)
+
+		for _, h := range c.Headers {
+			head := strings.Split(h, ":")
+			if len(head) == 2 {
+				headers[strings.TrimSpace(head[0])] = strings.TrimSpace(head[1])
+			}
+		}
+
+		rt = client.NewHeadersRoundTripper(headers, rt)
+	}
+
 	return client.NewClient(u.String(), rt)
 }
 

cmd/health_test.go

@@ -144,6 +144,21 @@ func TestHealthCmd(t *testing.T) {
 			args:     []string{"run", "../main.go", "--user", "passwordmissing", "health"},
 			expected: "[UNKNOWN] - specify the user name and password for server authentication <user:password> (*errors.errorString)\nexit status 3\n",
 		},
+		{
+			name: "health-extra-header",
+			server: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				foobar := r.Header.Get("X-Foobar")
+				if foobar == "Barfoo" {
+					w.WriteHeader(http.StatusOK)
+					w.Write([]byte(`Prometheus Server is Healthy.`))
+					return
+				}
+				w.WriteHeader(http.StatusUnauthorized)
+				w.Write([]byte(`Wrong Header!`))
+			})),
+			args:     []string{"run", "../main.go", "--header", "X-Foobar: Barfoo", "health"},
+			expected: "[OK] - states: ok=1\n\\_ [OK] Prometheus Server is Healthy.\n\n",
+		},
 	}
 
 	for _, test := range tests {

cmd/root.go

@@ -61,6 +61,9 @@ func init() {
 		"Specify the Key File for TLS authentication (CHECK_PROMETHEUS_KEY_FILE)")
 	pfs.IntVarP(&Timeout, "timeout", "t", Timeout,
 		"Timeout in seconds for the CheckPlugin")
+	pfs.StringSliceVarP(&cliConfig.Headers, "header", "", nil,
+		`Additional HTTP header to include in the request. Can be used multiple times.
+Keys and values are separated by a colon (--header "X-Custom: example").`)
 
 	rootCmd.Flags().SortFlags = false
 	pfs.SortFlags = false

internal/client/client.go

@@ -85,3 +85,39 @@ func (c *Client) GetStatus(ctx context.Context, endpoint string) (returncode int
 
 	return check.Unknown, resp.StatusCode, respBody, err
 }
+
+type headersRoundTripper struct {
+	headers map[string]string
+	rt      http.RoundTripper
+}
+
+// NewHeadersRoundTripper adds the given headers to a request
+func NewHeadersRoundTripper(headers map[string]string, rt http.RoundTripper) http.RoundTripper {
+	return &headersRoundTripper{headers, rt}
+}
+
+func (rt *headersRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	// RoundTrip should not modify the request, except for
+	// consuming and closing the Request's Body.
+	req = cloneRequest(req)
+
+	for key, value := range rt.headers {
+		req.Header.Add(key, value)
+	}
+
+	return rt.rt.RoundTrip(req)
+}
+
+// cloneRequest returns a clone of the provided *http.Request
+func cloneRequest(r *http.Request) *http.Request {
+	// Shallow copy of the struct.
+	r2 := new(http.Request)
+	*r2 = *r
+	// Deep copy of the Header.
+	r2.Header = make(http.Header)
+	for k, s := range r.Header {
+		r2.Header[k] = s
+	}
+
+	return r2
+}