From e5bb65c62e64d4e7cdd4e1b329ddce5a97253da6 Mon Sep 17 00:00:00 2001
From: Conor Mongey <conor@mongey.net>
Date: Tue, 5 Jan 2021 01:07:35 +0000
Subject: [PATCH] Allow setting of headers in api client

---
 api/api.go      |  9 +++++++++
 api/api_test.go | 16 ++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/api/api.go b/api/api.go
index e54d37bb9..f0d4022b8 100644
--- a/api/api.go
+++ b/api/api.go
@@ -155,6 +155,8 @@ type Config struct {
 	//
 	// TLSConfig is ignored if HttpClient is set.
 	TLSConfig *TLSConfig
+
+	Headers map[string]string
 }
 
 // ClientConfig copies the configuration with a new client address, region, and
@@ -527,6 +529,7 @@ type request struct {
 	body   io.Reader
 	obj    interface{}
 	ctx    context.Context
+	header http.Header
 }
 
 // setQueryOptions is used to annotate the request with
@@ -621,6 +624,7 @@ func (r *request) toHTTP() (*http.Request, error) {
 		req.SetBasicAuth(r.config.HttpAuth.Username, r.config.HttpAuth.Password)
 	}
 
+	req.Header = r.header
 	req.Header.Add("Accept-Encoding", "gzip")
 	if r.token != "" {
 		req.Header.Set("X-Nomad-Token", r.token)
@@ -649,6 +653,7 @@ func (c *Client) newRequest(method, path string) (*request, error) {
 			Path:    u.Path,
 			RawPath: u.RawPath,
 		},
+		header: make(http.Header),
 		params: make(map[string][]string),
 	}
 	if c.config.Region != "" {
@@ -671,6 +676,10 @@ func (c *Client) newRequest(method, path string) (*request, error) {
 		}
 	}
 
+	for key, value := range c.config.Headers {
+		r.header.Set(key, value)
+	}
+
 	return r, nil
 }
 
diff --git a/api/api_test.go b/api/api_test.go
index d46098796..e2dcd6321 100644
--- a/api/api_test.go
+++ b/api/api_test.go
@@ -341,6 +341,22 @@ func TestParseWriteMeta(t *testing.T) {
 	}
 }
 
+func TestClientHeader(t *testing.T) {
+	t.Parallel()
+	c, s := makeClient(t, func(c *Config) {
+		c.Headers = map[string]string{
+			"Hello": "World",
+		}
+	}, nil)
+	defer s.Stop()
+
+	r, _ := c.newRequest("GET", "/v1/jobs")
+
+	if r.header.Get("Hello") != "World" {
+		t.Fatalf("bad: %v", r.header)
+	}
+}
+
 func TestQueryString(t *testing.T) {
 	t.Parallel()
 	c, s := makeClient(t, nil, nil)