app/vmctl: Add ability to set headers for vm-native HTTP requests. (#3906)

app/vmctl: Add ability to set headers for vm-native HTTP requests

app/vmctl/flags.go

@@ -328,10 +328,12 @@ const (
 	vmNativeSrcAddr     = "vm-native-src-addr"
 	vmNativeSrcUser     = "vm-native-src-user"
 	vmNativeSrcPassword = "vm-native-src-password"
+	vmNativeSrcHeaders  = "vm-native-src-headers"
 
 	vmNativeDstAddr     = "vm-native-dst-addr"
 	vmNativeDstUser     = "vm-native-dst-user"
 	vmNativeDstPassword = "vm-native-dst-password"
+	vmNativeDstHeaders  = "vm-native-dst-headers"
 )
 
 var (
@@ -373,6 +375,12 @@ var (
 			Usage:   "VictoriaMetrics password for basic auth",
 			EnvVars: []string{"VM_NATIVE_SRC_PASSWORD"},
 		},
+		&cli.StringFlag{
+			Name: vmNativeSrcHeaders,
+			Usage: "Optional HTTP headers to send with each request to the corresponding source address. \n" +
+				"For example, --vm-native-src-headers='My-Auth:foobar' would send 'My-Auth: foobar' HTTP header with every request to the corresponding source address. \n" +
+				"Multiple headers must be delimited by '^^': --vm-native-src-headers='header1:value1^^header2:value2'",
+		},
 		&cli.StringFlag{
 			Name: vmNativeDstAddr,
 			Usage: "VictoriaMetrics address to perform import to. \n" +
@@ -390,6 +398,12 @@ var (
 			Usage:   "VictoriaMetrics password for basic auth",
 			EnvVars: []string{"VM_NATIVE_DST_PASSWORD"},
 		},
+		&cli.StringFlag{
+			Name: vmNativeDstHeaders,
+			Usage: "Optional HTTP headers to send with each request to the corresponding destination address. \n" +
+				"For example, --vm-native-dst-headers='My-Auth:foobar' would send 'My-Auth: foobar' HTTP header with every request to the corresponding destination address. \n" +
+				"Multiple headers must be delimited by '^^': --vm-native-dst-headers='header1:value1^^header2:value2'",
+		},
 		&cli.StringSliceFlag{
 			Name:  vmExtraLabel,
 			Value: nil,

app/vmctl/main.go

@@ -212,12 +212,14 @@ func main() {
 							Addr:     strings.Trim(c.String(vmNativeSrcAddr), "/"),
 							User:     c.String(vmNativeSrcUser),
 							Password: c.String(vmNativeSrcPassword),
+							Headers:  c.String(vmNativeSrcHeaders),
 						},
 						dst: &native.Client{
 							Addr:        strings.Trim(c.String(vmNativeDstAddr), "/"),
 							User:        c.String(vmNativeDstUser),
 							Password:    c.String(vmNativeDstPassword),
 							ExtraLabels: c.StringSlice(vmExtraLabel),
+							Headers:     c.String(vmNativeDstHeaders),
 						},
 						backoff: backoff.New(),
 						cc:      c.Int(vmConcurrency),

app/vmctl/native/client.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 )
 
 const (
@@ -21,6 +22,7 @@ type Client struct {
 	User        string
 	Password    string
 	ExtraLabels []string
+	Headers     string
 }
 
 // LabelValues represents series from api/v1/series response
@@ -89,6 +91,16 @@ func (c *Client) ImportPipe(ctx context.Context, dstURL string, pr *io.PipeReade
 	if err != nil {
 		return fmt.Errorf("cannot create import request to %q: %s", c.Addr, err)
 	}
+
+	parsedHeaders, err := parseHeaders(c.Headers)
+	if err != nil {
+		return err
+	}
+
+	for _, header := range parsedHeaders {
+		req.Header.Set(header.key, header.value)
+	}
+
 	importResp, err := c.do(req, http.StatusNoContent)
 	if err != nil {
 		return fmt.Errorf("import request failed: %s", err)
@@ -118,6 +130,16 @@ func (c *Client) ExportPipe(ctx context.Context, url string, f Filter) (io.ReadC
 
 	// disable compression since it is meaningless for native format
 	req.Header.Set("Accept-Encoding", "identity")
+
+	parsedHeaders, err := parseHeaders(c.Headers)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, header := range parsedHeaders {
+		req.Header.Set(header.key, header.value)
+	}
+
 	resp, err := c.do(req, http.StatusOK)
 	if err != nil {
 		return nil, fmt.Errorf("export request failed: %w", err)
@@ -142,6 +164,15 @@ func (c *Client) GetSourceTenants(ctx context.Context, f Filter) ([]string, erro
 	}
 	req.URL.RawQuery = params.Encode()
 
+	parsedHeaders, err := parseHeaders(c.Headers)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, header := range parsedHeaders {
+		req.Header.Set(header.key, header.value)
+	}
+
 	resp, err := c.do(req, http.StatusOK)
 	if err != nil {
 		return nil, fmt.Errorf("tenants request failed: %s", err)
@@ -179,3 +210,28 @@ func (c *Client) do(req *http.Request, expSC int) (*http.Response, error) {
 	}
 	return resp, err
 }
+
+type keyValue struct {
+	key   string
+	value string
+}
+
+func parseHeaders(headers string) ([]keyValue, error) {
+	if len(headers) == 0 {
+		return nil, nil
+	}
+
+	var headersSplitByDelimiter = strings.Split(headers, "^^")
+
+	kvs := make([]keyValue, len(headersSplitByDelimiter))
+	for i, h := range headersSplitByDelimiter {
+		n := strings.IndexByte(h, ':')
+		if n < 0 {
+			return nil, fmt.Errorf(`missing ':' in header %q; expecting "key: value" format`, h)
+		}
+		kv := &kvs[i]
+		kv.key = strings.TrimSpace(h[:n])
+		kv.value = strings.TrimSpace(h[n+1:])
+	}
+	return kvs, nil
+}