From 701b34dc555f1b34ab8756833e62baee5e5ef24a Mon Sep 17 00:00:00 2001
From: Dylan Meissner <dylanmei@gmail.com>
Date: Mon, 28 Aug 2017 16:08:50 -0700
Subject: [PATCH] HTTP headers can be added to InfluxDB output (#3182)

---
 plugins/outputs/influxdb/README.md           |  4 ++++
 plugins/outputs/influxdb/client/http.go      | 10 ++++++++++
 plugins/outputs/influxdb/client/http_test.go | 10 ++++++++++
 plugins/outputs/influxdb/influxdb.go         | 14 +++++++++++---
 4 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/plugins/outputs/influxdb/README.md b/plugins/outputs/influxdb/README.md
index 36d64c219..31326c918 100644
--- a/plugins/outputs/influxdb/README.md
+++ b/plugins/outputs/influxdb/README.md
@@ -44,6 +44,9 @@ This plugin writes to [InfluxDB](https://www.influxdb.com) via HTTP or UDP.
   ## HTTP Proxy Config
   # http_proxy = "http://corporate.proxy:3128"
 
+  ## Optional HTTP headers
+  # http_headers = {"X-Special-Header" = "Special-Value"}
+
   ## Compress each HTTP request payload using GZIP.
   # content_encoding = "gzip"
 ```
@@ -70,4 +73,5 @@ to write to. Each URL should start with either `http://` or `udp://`
 * `ssl_key`: SSL key
 * `insecure_skip_verify`: Use SSL but skip chain & host verification (default: false)
 * `http_proxy`: HTTP Proxy URI
+* `http_headers`: HTTP headers to add to each HTTP request
 * `content_encoding`: Compress each HTTP request payload using gzip if set to: "gzip"
diff --git a/plugins/outputs/influxdb/client/http.go b/plugins/outputs/influxdb/client/http.go
index 957c490ed..2587b2da5 100644
--- a/plugins/outputs/influxdb/client/http.go
+++ b/plugins/outputs/influxdb/client/http.go
@@ -68,6 +68,8 @@ func NewHTTP(config HTTPConfig, defaultWP WriteParams) (Client, error) {
 	}, nil
 }
 
+type HTTPHeaders map[string]string
+
 type HTTPConfig struct {
 	// URL should be of the form "http://host:port" (REQUIRED)
 	URL string
@@ -95,6 +97,9 @@ type HTTPConfig struct {
 	// Proxy URL should be of the form "http://host:port"
 	HTTPProxy string
 
+	// HTTP headers to append to HTTP requests.
+	HTTPHeaders HTTPHeaders
+
 	// The content encoding mechanism to use for each request.
 	ContentEncoding string
 }
@@ -253,6 +258,11 @@ func (c *httpClient) makeRequest(uri string, body io.Reader) (*http.Request, err
 	if err != nil {
 		return nil, err
 	}
+
+	for header, value := range c.config.HTTPHeaders {
+		req.Header.Set(header, value)
+	}
+
 	req.Header.Set("Content-Type", "text/plain")
 	req.Header.Set("User-Agent", c.config.UserAgent)
 	if c.config.Username != "" && c.config.Password != "" {
diff --git a/plugins/outputs/influxdb/client/http_test.go b/plugins/outputs/influxdb/client/http_test.go
index 7f3bd7e34..d094078ea 100644
--- a/plugins/outputs/influxdb/client/http_test.go
+++ b/plugins/outputs/influxdb/client/http_test.go
@@ -55,6 +55,13 @@ func TestHTTPClient_Write(t *testing.T) {
 				fmt.Fprintln(w, `{"results":[{}],"error":"basic auth incorrect"}`)
 			}
 
+			// test that user-specified http header is set properly
+			if r.Header.Get("X-Test-Header") != "Test-Value" {
+				w.WriteHeader(http.StatusTeapot)
+				w.Header().Set("Content-Type", "application/json")
+				fmt.Fprintln(w, `{"results":[{}],"error":"wrong http header value"}`)
+			}
+
 			// Validate Content-Length Header
 			if r.ContentLength != 13 {
 				w.WriteHeader(http.StatusTeapot)
@@ -90,6 +97,9 @@ func TestHTTPClient_Write(t *testing.T) {
 		UserAgent: "test-agent",
 		Username:  "test-user",
 		Password:  "test-password",
+		HTTPHeaders: HTTPHeaders{
+			"X-Test-Header": "Test-Value",
+		},
 	}
 	wp := WriteParams{
 		Database:        "test",
diff --git a/plugins/outputs/influxdb/influxdb.go b/plugins/outputs/influxdb/influxdb.go
index b4567af99..f2f946513 100644
--- a/plugins/outputs/influxdb/influxdb.go
+++ b/plugins/outputs/influxdb/influxdb.go
@@ -32,9 +32,10 @@ type InfluxDB struct {
 	RetentionPolicy  string
 	WriteConsistency string
 	Timeout          internal.Duration
-	UDPPayload       int    `toml:"udp_payload"`
-	HTTPProxy        string `toml:"http_proxy"`
-	ContentEncoding  string `toml:"content_encoding"`
+	UDPPayload       int               `toml:"udp_payload"`
+	HTTPProxy        string            `toml:"http_proxy"`
+	HTTPHeaders      map[string]string `toml:"http_headers"`
+	ContentEncoding  string            `toml:"content_encoding"`
 
 	// Path to CA file
 	SSLCA string `toml:"ssl_ca"`
@@ -89,6 +90,9 @@ var sampleConfig = `
   ## HTTP Proxy Config
   # http_proxy = "http://corporate.proxy:3128"
 
+  ## Optional HTTP headers
+  # http_headers = {"X-Special-Header" = "Special-Value"}
+
   ## Compress each HTTP request payload using GZIP.
   # content_encoding = "gzip"
 `
@@ -132,8 +136,12 @@ func (i *InfluxDB) Connect() error {
 				Username:        i.Username,
 				Password:        i.Password,
 				HTTPProxy:       i.HTTPProxy,
+				HTTPHeaders:     client.HTTPHeaders{},
 				ContentEncoding: i.ContentEncoding,
 			}
+			for header, value := range i.HTTPHeaders {
+				config.HTTPHeaders[header] = value
+			}
 			wp := client.WriteParams{
 				Database:        i.Database,
 				RetentionPolicy: i.RetentionPolicy,