From 4b136abff8fef02f670d63e7de7ac0cab73f0b9e Mon Sep 17 00:00:00 2001 From: Gowtam Lal Date: Mon, 6 Mar 2023 05:22:31 -0500 Subject: [PATCH] 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 | 14 ++++++++++ app/vmctl/main.go | 2 ++ app/vmctl/native/client.go | 56 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+) diff --git a/app/vmctl/flags.go b/app/vmctl/flags.go index 0f2b6da1d..afebd79a8 100644 --- a/app/vmctl/flags.go +++ b/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, diff --git a/app/vmctl/main.go b/app/vmctl/main.go index 384d918ff..aa5bad045 100644 --- a/app/vmctl/main.go +++ b/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), diff --git a/app/vmctl/native/client.go b/app/vmctl/native/client.go index 908251b9a..ad73224c5 100644 --- a/app/vmctl/native/client.go +++ b/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 +}