Implement custom headers: #10200

Allows specifying extra headers to send on the command line

Author: ire4ever1190

cmd/logcli/main.go

@@ -593,6 +593,7 @@ func newQueryClient(app *kingpin.Application) client.Client {
 	app.Flag("proxy-url", "The http or https proxy to use when making requests. Can also be set using LOKI_HTTP_PROXY_URL env var.").Default("").Envar("LOKI_HTTP_PROXY_URL").StringVar(&client.ProxyURL)
 	app.Flag("compress", "Request that Loki compress returned data in transit. Can also be set using LOKI_HTTP_COMPRESSION env var.").Default("false").Envar("LOKI_HTTP_COMPRESSION").BoolVar(&client.Compression)
 	app.Flag("envproxy", "Use ProxyFromEnvironment to use net/http ProxyFromEnvironment configuration, eg HTTP_PROXY").Default("false").Envar("LOKI_ENV_PROXY").BoolVar(&client.EnvironmentProxy)
+	app.Flag("header", "Add custom HTTP headers to requests. Can be specified multiple times. Format: 'Header-Name: value'").StringsVar(&client.CustomHeaders)
 
 	return client
 }

pkg/logcli/client/client.go

@@ -94,6 +94,7 @@ type DefaultClient struct {
 	BackoffConfig    BackoffConfig
 	Compression      bool
 	EnvironmentProxy bool
+	CustomHeaders    []string
 }
 
 // Query uses the /api/v1/query endpoint to execute an instant query
@@ -648,6 +649,22 @@ func (c *DefaultClient) getHTTPRequestHeader() (http.Header, error) {
 		h.Set(HTTPQueryTags, c.QueryTags)
 	}
 
+	// Add custom headers
+	if c.CustomHeaders != nil {
+		for _, header := range c.CustomHeaders {
+			parts := strings.SplitN(header, ":", 2)
+			if len(parts) != 2 {
+				return nil, fmt.Errorf("invalid header format: %q. Expected format: 'Header-Name: value'", header)
+			}
+			key := strings.TrimSpace(parts[0])
+			value := strings.TrimSpace(parts[1])
+			if key == "" {
+				return nil, fmt.Errorf("header name cannot be empty in header: %q", header)
+			}
+			h.Set(key, value)
+		}
+	}
+
 	if (c.Username != "" || c.Password != "") && (len(c.BearerToken) > 0 || len(c.BearerTokenFile) > 0) {
 		return nil, fmt.Errorf("at most one of HTTP basic auth (username/password), bearer-token & bearer-token-file is allowed to be configured")
 	}

pkg/logcli/client/client_test.go

@@ -59,6 +59,24 @@ func Test_getHTTPRequestHeader(t *testing.T) {
 		}, http.Header{
 			"Authorization": []string{"Bearer " + "secureToken"},
 		}, false},
+		{"custom-headers", DefaultClient{
+			CustomHeaders: []string{"X-Custom-Header: custom-value", "X-Another-Header: another-value"},
+		}, http.Header{
+			"X-Custom-Header":  []string{"custom-value"},
+			"X-Another-Header": []string{"another-value"},
+		}, false},
+		{"custom-headers-with-spaces", DefaultClient{
+			CustomHeaders: []string{"X-Custom-Header:  custom-value  ", "X-Another-Header:  another-value  "},
+		}, http.Header{
+			"X-Custom-Header":  []string{"custom-value"},
+			"X-Another-Header": []string{"another-value"},
+		}, false},
+		{"custom-headers-invalid-format", DefaultClient{
+			CustomHeaders: []string{"InvalidHeader"},
+		}, nil, true},
+		{"custom-headers-empty-name", DefaultClient{
+			CustomHeaders: []string{" : value"},
+		}, nil, true},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -68,6 +86,12 @@ func Test_getHTTPRequestHeader(t *testing.T) {
 				return
 			}
 
+			// If we expect an error, we shouldn't have any headers
+			if tt.wantErr {
+				assert.Nil(t, got)
+				return
+			}
+
 			// User-Agent should be set all the time.
 			assert.Equal(t, got["User-Agent"], []string{userAgent})