Merge pull request #33 from NETWAYS/chore/cleanup

martialblog · web-flow · commit 5e90d0d61de2 · 2025-04-16T11:31:27.000+02:00
Some various cleanup
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -19,7 +19,8 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.24
+          go-version: stable
+          check-latest: true
 
       - name: Test
         run: go test -v ./...
diff --git a/.gitignore b/.gitignore
@@ -1,4 +1,4 @@
-.idea
+.idea/
 check_influxdb
 vendor/
 coverage.html
diff --git a/.goreleaser.yml b/.goreleaser.yml
@@ -8,6 +8,7 @@ builds:
       - CGO_ENABLED=0
     goarch:
       - amd64
+      - arm64
     goos:
       - linux
       - windows
@@ -30,6 +31,7 @@ archives:
       {{- if eq .Os "windows" }}Windows{{ end }}
       {{- if eq .Os "darwin" }}Darwin{{ end }}
       {{- if eq .Arch "amd64" }}_x86_64{{ end }}
+      {{- if eq .Arch "arm64" }}_arm64{{ end }}
 checksum:
   name_template: 'checksums.txt'
 snapshot:
diff --git a/.idea/.gitignore b/.idea/.gitignore
diff --git a/.idea/check_influxdb.iml b/.idea/check_influxdb.iml
diff --git a/.idea/modules.xml b/.idea/modules.xml
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/README.md b/README.md
@@ -31,6 +31,8 @@ Global Flags:
       --key-file string    Specify the Key File for TLS authentication (CHECK_INFLUXDB_KEY_FILE)
   -i, --insecure           Skip the verification of the server's TLS certificate
   -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_influxdb
   -v, --version            version for check_influxdb
 ```
diff --git a/cmd/config.go b/cmd/config.go
@@ -24,6 +24,7 @@ type Config struct {
 	KeyFile      string `env:"CHECK_INFLUXDB_KEY_FILE"`
 	Token        string `env:"CHECK_INFLUXDB_TOKEN"`
 	Organization string `env:"CHECK_INFLUXDB_ORGANISATION"`
+	Headers      []string
 	Port         int
 	Insecure     bool
 	Secure       bool
@@ -82,6 +83,20 @@ func (c *Config) NewClient() *client.Client {
 		rt = checkhttpconfig.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(), c.Token, c.Organization, rt)
 }
 
diff --git a/cmd/health_test.go b/cmd/health_test.go
@@ -63,6 +63,21 @@ func TestHealthCmd(t *testing.T) {
 			args:     []string{"run", "../main.go", "health"},
 			expected: "[UNKNOWN] - InfluxDB Version 3 not supported (*errors.errorString)\nexit status 3\n",
 		},
+		{
+			name: "health-extra-header",
+			httpreturn: func(w http.ResponseWriter, r *http.Request) {
+				foobar := r.Header.Get("X-Foobar")
+				if foobar == "Barfoo" {
+					w.WriteHeader(http.StatusOK)
+					w.Write([]byte(`{"checks":[],"message":"ready for queries and writes","name":"influxdb","status":"pass","version":"1.8.10"}`))
+					return
+				}
+				w.WriteHeader(http.StatusUnauthorized)
+				w.Write([]byte(`Wrong Header!`))
+			},
+			args:     []string{"run", "../main.go", "--header", "X-Foobar: Barfoo", "health"},
+			expected: "[OK] - InfluxDB Status: pass\n",
+		},
 	}
 
 	for _, test := range tests {
diff --git a/cmd/root.go b/cmd/root.go
@@ -59,6 +59,9 @@ func init() {
 		"Skip the verification of the server's TLS certificate")
 	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
diff --git a/internal/client/client.go b/internal/client/client.go
@@ -69,3 +69,39 @@ func (c *Client) Version() (influxdb.APIVersion, error) {
 
 	return v, nil
 }
+
+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
+}

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-.idea`
	`1`	`+.idea/`
`2`	`2`	`check_influxdb`
`3`	`3`	`vendor/`
`4`	`4`	`coverage.html`