Support arbitrary headers with HTTP datasources

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

README.md

@@ -585,6 +585,13 @@ $ echo 'Hello there, {{(datasource "foo").headers.Host}}...' | gomplate -d foo=h
 Hello there, httpbin.org...
 ```
 
+Additional headers can be provided with the `--datasource-header`/`-H` option:
+
+```console
+$ gomplate -d foo=https://httpbin.org/get -H 'foo=Foo: bar' -i '{{(datasource "foo").headers.Foo}}'
+bar
+```
+
 ###### Usage with Vault data
 
 The special `vault://` URL scheme can be used to retrieve data from [Hashicorp

data.go

@@ -17,6 +17,9 @@ import (
 	"github.com/hairyhenderson/gomplate/vault"
 )
 
+// logFatal is defined so log.Fatal calls can be overridden for testing
+var logFatalf = log.Fatalf
+
 func init() {
 	// Add some types we want to be able to handle which can be missing by default
 	err := mime.AddExtensionType(".json", "application/json")
@@ -55,14 +58,16 @@ type Data struct {
 }
 
 // NewData - constructor for Data
-func NewData(datasourceArgs []string) *Data {
+func NewData(datasourceArgs []string, headerArgs []string) *Data {
 	sources := make(map[string]*Source)
+	headers := parseHeaderArgs(headerArgs)
 	for _, v := range datasourceArgs {
 		s, err := ParseSource(v)
 		if err != nil {
 			log.Fatalf("error parsing datasource %v", err)
 			return nil
 		}
+		s.Header = headers[s.Alias]
 		sources[s.Alias] = s
 	}
 	return &Data{
@@ -72,13 +77,14 @@ func NewData(datasourceArgs []string) *Data {
 
 // Source - a data source
 type Source struct {
-	Alias string
-	URL   *url.URL
-	Ext   string
-	Type  string
-	FS    vfs.Filesystem // used for file: URLs, nil otherwise
-	HC    *http.Client   // used for http[s]: URLs, nil otherwise
-	VC    *vault.Client  //used for vault: URLs, nil otherwise
+	Alias  string
+	URL    *url.URL
+	Ext    string
+	Type   string
+	FS     vfs.Filesystem // used for file: URLs, nil otherwise
+	HC     *http.Client   // used for http[s]: URLs, nil otherwise
+	VC     *vault.Client  //used for vault: URLs, nil otherwise
+	Header http.Header    // used for http[s]: URLs, nil otherwise
 }
 
 // NewSource - builds a &Source
@@ -236,7 +242,12 @@ func readHTTP(source *Source, args ...string) ([]byte, error) {
 	if source.HC == nil {
 		source.HC = &http.Client{Timeout: time.Second * 5}
 	}
-	res, err := source.HC.Get(source.URL.String())
+	req, err := http.NewRequest("GET", source.URL.String(), nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header = source.Header
+	res, err := source.HC.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -286,3 +297,37 @@ func readVault(source *Source, args ...string) ([]byte, error) {
 
 	return data, nil
 }
+
+func parseHeaderArgs(headerArgs []string) map[string]http.Header {
+	headers := make(map[string]http.Header)
+	for _, v := range headerArgs {
+		ds, name, value := splitHeaderArg(v)
+		if _, ok := headers[ds]; !ok {
+			headers[ds] = make(http.Header)
+		}
+		headers[ds][name] = append(headers[ds][name], strings.TrimSpace(value))
+	}
+	return headers
+}
+
+func splitHeaderArg(arg string) (datasourceAlias, name, value string) {
+	parts := strings.SplitN(arg, "=", 2)
+	if len(parts) != 2 {
+		logFatalf("Invalid datasource-header option '%s'", arg)
+		return "", "", ""
+	}
+	datasourceAlias = parts[0]
+	name, value = splitHeader(parts[1])
+	return datasourceAlias, name, value
+}
+
+func splitHeader(header string) (name, value string) {
+	parts := strings.SplitN(header, ":", 2)
+	if len(parts) != 2 {
+		logFatalf("Invalid HTTP Header format '%s'", header)
+		return "", ""
+	}
+	name = http.CanonicalHeaderKey(parts[0])
+	value = parts[1]
+	return name, value
+}

data_test.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -12,6 +14,22 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+var spyLogFatalfMsg string
+
+func restoreLogFatalf() {
+	logFatalf = log.Fatalf
+}
+
+func mockLogFatalf(msg string, args ...interface{}) {
+	spyLogFatalfMsg = msg
+	panic(spyLogFatalfMsg)
+}
+
+func setupMockLogFatalf() {
+	logFatalf = mockLogFatalf
+	spyLogFatalfMsg = ""
+}
+
 func TestNewSource(t *testing.T) {
 	s := NewSource("foo", &url.URL{
 		Scheme: "file",
@@ -37,6 +55,26 @@ func TestNewSource(t *testing.T) {
 	assert.Equal(t, ".json", s.Ext)
 }
 
+func TestNewData(t *testing.T) {
+	d := NewData(nil, nil)
+	assert.Len(t, d.Sources, 0)
+
+	d = NewData([]string{"foo=http:///foo.json"}, nil)
+	assert.Equal(t, "/foo.json", d.Sources["foo"].URL.Path)
+
+	d = NewData([]string{"foo=http:///foo.json"}, []string{})
+	assert.Equal(t, "/foo.json", d.Sources["foo"].URL.Path)
+	assert.Empty(t, d.Sources["foo"].Header)
+
+	d = NewData([]string{"foo=http:///foo.json"}, []string{"bar=Accept: blah"})
+	assert.Equal(t, "/foo.json", d.Sources["foo"].URL.Path)
+	assert.Empty(t, d.Sources["foo"].Header)
+
+	d = NewData([]string{"foo=http:///foo.json"}, []string{"foo=Accept: blah"})
+	assert.Equal(t, "/foo.json", d.Sources["foo"].URL.Path)
+	assert.Equal(t, "blah", d.Sources["foo"].Header["Accept"][0])
+}
+
 func TestParseSourceNoAlias(t *testing.T) {
 	s, err := ParseSource("foo.json")
 	assert.NoError(t, err)
@@ -112,9 +150,15 @@ func TestDatasourceExists(t *testing.T) {
 
 func setupHTTP(code int, mimetype string, body string) (*httptest.Server, *http.Client) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
 		w.Header().Set("Content-Type", mimetype)
 		w.WriteHeader(code)
-		fmt.Fprintln(w, body)
+		if body == "" {
+			// mirror back the headers
+			fmt.Fprintln(w, marshalObj(r.Header, json.Marshal))
+		} else {
+			fmt.Fprintln(w, body)
+		}
 	}))
 
 	client := &http.Client{
@@ -150,3 +194,80 @@ func TestHTTPFile(t *testing.T) {
 	actual := data.Datasource("foo")
 	assert.Equal(t, expected["hello"], actual["hello"])
 }
+
+func TestHTTPFileWithHeaders(t *testing.T) {
+	server, client := setupHTTP(200, "application/json", "")
+	defer server.Close()
+
+	sources := make(map[string]*Source)
+	sources["foo"] = &Source{
+		Alias: "foo",
+		URL: &url.URL{
+			Scheme: "http",
+			Host:   "example.com",
+			Path:   "/foo",
+		},
+		HC: client,
+		Header: http.Header{
+			"Foo":             {"bar"},
+			"foo":             {"baz"},
+			"User-Agent":      {},
+			"Accept-Encoding": {"test"},
+		},
+	}
+	data := &Data{
+		Sources: sources,
+	}
+	expected := http.Header{
+		"Accept-Encoding": {"test"},
+		"Foo":             {"bar", "baz"},
+	}
+	actual := data.Datasource("foo")
+	assert.Equal(t, marshalObj(expected, json.Marshal), marshalObj(actual, json.Marshal))
+}
+
+func TestParseHeaderArgs(t *testing.T) {
+	args := []string{
+		"foo=Accept: application/json",
+		"bar=Authorization: Bearer supersecret",
+	}
+	expected := map[string]http.Header{
+		"foo": {
+			"Accept": {"application/json"},
+		},
+		"bar": {
+			"Authorization": {"Bearer supersecret"},
+		},
+	}
+	assert.Equal(t, expected, parseHeaderArgs(args))
+
+	defer restoreLogFatalf()
+	setupMockLogFatalf()
+	assert.Panics(t, func() {
+		parseHeaderArgs([]string{"foo"})
+	})
+
+	defer restoreLogFatalf()
+	setupMockLogFatalf()
+	assert.Panics(t, func() {
+		parseHeaderArgs([]string{"foo=bar"})
+	})
+
+	args = []string{
+		"foo=Accept: application/json",
+		"foo=Foo: bar",
+		"foo=foo: baz",
+		"foo=fOO: qux",
+		"bar=Authorization: Bearer  supersecret",
+	}
+	expected = map[string]http.Header{
+		"foo": {
+			"Accept": {"application/json"},
+			"Foo":    {"bar", "baz", "qux"},
+		},
+		"bar": {
+			"Authorization": {"Bearer  supersecret"},
+		},
+	}
+	assert.Equal(t, expected, parseHeaderArgs(args))
+}

main.go

@@ -117,7 +117,7 @@ func openOutFile(filename string) (out *os.File, err error) {
 
 func runTemplate(c *cli.Context) error {
 	defer runCleanupHooks()
-	data := NewData(c.StringSlice("datasource"))
+	data := NewData(c.StringSlice("datasource"), c.StringSlice("datasource-header"))
 	lDelim := c.String("left-delim")
 	rDelim := c.String("right-delim")
 
@@ -166,6 +166,10 @@ func main() {
 			Name:  "datasource, d",
 			Usage: "Data source in alias=URL form. Specify multiple times to add multiple sources.",
 		},
+		cli.StringSliceFlag{
+			Name:  "datasource-header, H",
+			Usage: "HTTP Header field in 'alias=Name: value' form to be provided on HTTP-based data sources. Multiples can be set.",
+		},
 		cli.StringFlag{
 			Name:   "left-delim",
 			Usage:  "Override the default left-delimiter `{{`",

test/integration/datasources_http.bats

@@ -0,0 +1,12 @@
+#!/usr/bin/env bats
+
+load helper
+
+@test "HTTP datasource with headers" {
+  gomplate \
+    -d foo=http://httpbin.org/get \
+    -H foo=Foo:bar \
+    -i '{{ (datasource "foo").headers.Foo }}'
+  [ "$status" -eq 0 ]
+  [[ "${output}" == "bar" ]]
+}