Support URLs for nested templates

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

4 files changed, 358 insertions, 74 deletions

diff --git a/internal/config/configfile.go b/internal/config/configfile.go
index 687c232b..8bfdd0b7 100644
--- a/internal/config/configfile.go
+++ b/internal/config/configfile.go
@@ -36,6 +36,15 @@ type Config struct {
 	Stdout io.Writer `yaml:"-"`
 	Stderr io.Writer `yaml:"-"`
 
+	DataSources map[string]DataSource   `yaml:"datasources,omitempty"`
+	Context     map[string]DataSource   `yaml:"context,omitempty"`
+	Plugins     map[string]PluginConfig `yaml:"plugins,omitempty"`
+	Templates   Templates               `yaml:"templates,omitempty"`
+
+	// Extra HTTP headers not attached to pre-defined datsources. Potentially
+	// used by datasources defined in the template.
+	ExtraHeaders map[string]http.Header `yaml:"-"`
+
 	// internal use only, can't be injected in YAML
 	PostExecInput io.Reader `yaml:"-"`
 
@@ -54,16 +63,6 @@ type Config struct {
 
 	PostExec []string `yaml:"postExec,omitempty,flow"`
 
-	DataSources map[string]DataSource   `yaml:"datasources,omitempty"`
-	Context     map[string]DataSource   `yaml:"context,omitempty"`
-	Plugins     map[string]PluginConfig `yaml:"plugins,omitempty"`
-
-	// Extra HTTP headers not attached to pre-defined datsources. Potentially
-	// used by datasources defined in the template.
-	ExtraHeaders map[string]http.Header `yaml:"-"`
-
-	Templates []string `yaml:"templates,omitempty"`
-
 	PluginTimeout time.Duration `yaml:"pluginTimeout,omitempty"`
 
 	ExecPipe      bool `yaml:"execPipe,omitempty"`
@@ -251,8 +250,10 @@ func (c *Config) MergeFrom(o *Config) *Config {
 	if !isZero(o.RDelim) {
 		c.RDelim = o.RDelim
 	}
-	if !isZero(o.Templates) {
+	if c.Templates == nil {
 		c.Templates = o.Templates
+	} else {
+		mergeDataSources(c.Templates, o.Templates)
 	}
 	if c.DataSources == nil {
 		c.DataSources = o.DataSources
@@ -273,9 +274,44 @@ func (c *Config) MergeFrom(o *Config) *Config {
 	return c
 }
 
-// ParseDataSourceFlags - sets the DataSources and Context fields from the
-// key=value format flags as provided at the command-line
-func (c *Config) ParseDataSourceFlags(datasources, contexts, headers []string) error {
+// ParseDataSourceFlags - sets DataSources, Context, and Templates fields from
+// the key=value format flags as provided at the command-line
+// Unreferenced headers will be set in c.ExtraHeaders
+func (c *Config) ParseDataSourceFlags(datasources, contexts, templates, headers []string) error {
+	err := c.parseResources(datasources, contexts, templates)
+	if err != nil {
+		return err
+	}
+
+	hdrs, err := parseHeaderArgs(headers)
+	if err != nil {
+		return err
+	}
+
+	for k, v := range hdrs {
+		if d, ok := c.Context[k]; ok {
+			d.Header = v
+			c.Context[k] = d
+			delete(hdrs, k)
+		}
+		if d, ok := c.DataSources[k]; ok {
+			d.Header = v
+			c.DataSources[k] = d
+			delete(hdrs, k)
+		}
+		if t, ok := c.Templates[k]; ok {
+			t.Header = v
+			c.Templates[k] = t
+			delete(hdrs, k)
+		}
+	}
+	if len(hdrs) > 0 {
+		c.ExtraHeaders = hdrs
+	}
+	return nil
+}
+
+func (c *Config) parseResources(datasources, contexts, templates []string) error {
 	for _, d := range datasources {
 		k, ds, err := parseDatasourceArg(d)
 		if err != nil {
@@ -296,27 +332,17 @@ func (c *Config) ParseDataSourceFlags(datasources, contexts, headers []string) e
 		}
 		c.Context[k] = ds
 	}
-
-	hdrs, err := parseHeaderArgs(headers)
-	if err != nil {
-		return err
-	}
-
-	for k, v := range hdrs {
-		if d, ok := c.Context[k]; ok {
-			d.Header = v
-			c.Context[k] = d
-			delete(hdrs, k)
+	for _, t := range templates {
+		k, ds, err := parseTemplateArg(t)
+		if err != nil {
+			return err
 		}
-		if d, ok := c.DataSources[k]; ok {
-			d.Header = v
-			c.DataSources[k] = d
-			delete(hdrs, k)
+		if c.Templates == nil {
+			c.Templates = map[string]DataSource{}
 		}
+		c.Templates[k] = ds
 	}
-	if len(hdrs) > 0 {
-		c.ExtraHeaders = hdrs
-	}
+
 	return nil
 }
 
@@ -336,21 +362,20 @@ func (c *Config) ParsePluginFlags(plugins []string) error {
 	return nil
 }
 
-func parseDatasourceArg(value string) (key string, ds DataSource, err error) {
-	parts := strings.SplitN(value, "=", 2)
-	if len(parts) == 1 {
-		f := parts[0]
-		key = strings.SplitN(value, ".", 2)[0]
-		if path.Base(f) != f {
-			err = fmt.Errorf("invalid datasource (%s): must provide an alias with files not in working directory", value)
-			return key, ds, err
+func parseDatasourceArg(value string) (alias string, ds DataSource, err error) {
+	alias, u, _ := strings.Cut(value, "=")
+	if u == "" {
+		u = alias
+		alias, _, _ = strings.Cut(value, ".")
+		if path.Base(u) != u {
+			err = fmt.Errorf("invalid argument (%s): must provide an alias with files not in working directory", value)
+			return alias, ds, err
 		}
-		ds.URL, err = ParseSourceURL(f)
-	} else if len(parts) == 2 {
-		key = parts[0]
-		ds.URL, err = ParseSourceURL(parts[1])
 	}
-	return key, ds, err
+
+	ds.URL, err = ParseSourceURL(u)
+
+	return alias, ds, err
 }
 
 func parseHeaderArgs(headerArgs []string) (map[string]http.Header, error) {
diff --git a/internal/config/configfile_test.go b/internal/config/configfile_test.go
index 567ded60..458306a1 100644
--- a/internal/config/configfile_test.go
+++ b/internal/config/configfile_test.go
@@ -383,21 +383,65 @@ func TestMergeFrom(t *testing.T) {
 	}
 
 	assert.EqualValues(t, expected, cfg.MergeFrom(other))
+
+	// test template merging & a few other things
+	cfg = &Config{
+		InputDir:    "indir/",
+		ExcludeGlob: []string{"*.txt"},
+		Templates: Templates{
+			"foo": {
+				URL: mustURL("file:///foo.yaml"),
+			},
+			"bar": {
+				URL:    mustURL("stdin:///"),
+				Header: http.Header{"Accept": {"application/json"}},
+			},
+		},
+	}
+	other = &Config{
+		ExcludeGlob: []string{"*.yaml"},
+		OutputMap:   "${ .in }.out",
+		OutMode:     "600",
+		LDelim:      "${",
+		RDelim:      "}",
+		Templates: Templates{
+			"foo": {URL: mustURL("https://example.com/foo.yaml")},
+			"baz": {URL: mustURL("vault:///baz")},
+		},
+	}
+	expected = &Config{
+		InputDir:    "indir/",
+		ExcludeGlob: []string{"*.yaml"},
+		OutputMap:   "${ .in }.out",
+		OutMode:     "600",
+		LDelim:      "${",
+		RDelim:      "}",
+		Templates: Templates{
+			"foo": {URL: mustURL("https://example.com/foo.yaml")},
+			"bar": {
+				URL:    mustURL("stdin:///"),
+				Header: http.Header{"Accept": {"application/json"}},
+			},
+			"baz": {URL: mustURL("vault:///baz")},
+		},
+	}
+
+	assert.EqualValues(t, expected, cfg.MergeFrom(other))
 }
 
 func TestParseDataSourceFlags(t *testing.T) {
 	t.Parallel()
 	cfg := &Config{}
-	err := cfg.ParseDataSourceFlags(nil, nil, nil)
+	err := cfg.ParseDataSourceFlags(nil, nil, nil, nil)
 	assert.NoError(t, err)
 	assert.EqualValues(t, &Config{}, cfg)
 
 	cfg = &Config{}
-	err = cfg.ParseDataSourceFlags([]string{"foo/bar/baz.json"}, nil, nil)
+	err = cfg.ParseDataSourceFlags([]string{"foo/bar/baz.json"}, nil, nil, nil)
 	assert.Error(t, err)
 
 	cfg = &Config{}
-	err = cfg.ParseDataSourceFlags([]string{"baz=foo/bar/baz.json"}, nil, nil)
+	err = cfg.ParseDataSourceFlags([]string{"baz=foo/bar/baz.json"}, nil, nil, nil)
 	assert.NoError(t, err)
 	expected := &Config{
 		DataSources: map[string]DataSource{
@@ -410,6 +454,7 @@ func TestParseDataSourceFlags(t *testing.T) {
 	err = cfg.ParseDataSourceFlags(
 		[]string{"baz=foo/bar/baz.json"},
 		nil,
+		nil,
 		[]string{"baz=Accept: application/json"})
 	assert.NoError(t, err)
 	assert.EqualValues(t, &Config{
@@ -427,6 +472,7 @@ func TestParseDataSourceFlags(t *testing.T) {
 	err = cfg.ParseDataSourceFlags(
 		[]string{"baz=foo/bar/baz.json"},
 		[]string{"foo=http://example.com"},
+		nil,
 		[]string{"foo=Accept: application/json",
 			"bar=Authorization: Basic xxxxx"})
 	assert.NoError(t, err)
@@ -446,6 +492,28 @@ func TestParseDataSourceFlags(t *testing.T) {
 			"bar": {"Authorization": {"Basic xxxxx"}},
 		},
 	}, cfg)
+
+	cfg = &Config{}
+	err = cfg.ParseDataSourceFlags(
+		nil,
+		nil,
+		[]string{"foo=http://example.com", "file.tmpl", "tmpldir/"},
+		[]string{"foo=Accept: application/json",
+			"bar=Authorization: Basic xxxxx"})
+	assert.NoError(t, err)
+	assert.EqualValues(t, &Config{
+		Templates: Templates{
+			"foo": {
+				URL:    mustURL("http://example.com"),
+				Header: http.Header{"Accept": {"application/json"}},
+			},
+			"file.tmpl": {URL: mustURL("file.tmpl")},
+			"tmpldir/":  {URL: mustURL("tmpldir/")},
+		},
+		ExtraHeaders: map[string]http.Header{
+			"bar": {"Authorization": {"Basic xxxxx"}},
+		},
+	}, cfg)
 }
 
 func TestParsePluginFlags(t *testing.T) {
@@ -478,7 +546,10 @@ pluginTimeout: 5s
 		RDelim:      "R",
 		Input:       "foo",
 		OutputFiles: []string{"-"},
-		Templates:   []string{"foo=foo.t", "bar=bar.t"},
+		Templates: Templates{
+			"foo": {URL: mustURL("https://www.example.com/foo.tmpl")},
+			"bar": {URL: mustURL("/tmp/bar.t")},
+		},
 	}
 	expected = `---
 in: foo
@@ -486,17 +557,22 @@ outputFiles: ['-']
 leftDelim: L
 rightDelim: R
 templates:
-  - foo=foo.t
-  - bar=bar.t
+  foo:
+    url: https://www.example.com/foo.tmpl
+  bar:
+    url: file:///tmp/bar.t
 `
-	assert.Equal(t, expected, c.String())
+	assert.YAMLEq(t, expected, c.String())
 
 	c = &Config{
 		LDelim:      "L",
 		RDelim:      "R",
 		Input:       "long input that should be truncated",
 		OutputFiles: []string{"-"},
-		Templates:   []string{"foo=foo.t", "bar=bar.t"},
+		Templates: Templates{
+			"foo": {URL: mustURL("https://www.example.com/foo.tmpl")},
+			"bar": {URL: mustURL("/tmp/bar.t")},
+		},
 	}
 	expected = `---
 in: long inp...
@@ -504,10 +580,12 @@ outputFiles: ['-']
 leftDelim: L
 rightDelim: R
 templates:
-  - foo=foo.t
-  - bar=bar.t
+  foo:
+    url: https://www.example.com/foo.tmpl
+  bar:
+    url: file:///tmp/bar.t
 `
-	assert.Equal(t, expected, c.String())
+	assert.YAMLEq(t, expected, c.String())
 
 	c = &Config{
 		InputDir:  "in/",
@@ -748,9 +826,9 @@ func TestAbsFileURL(t *testing.T) {
 }
 
 func TestParseDatasourceArgNoAlias(t *testing.T) {
-	key, ds, err := parseDatasourceArg("foo.json")
+	alias, ds, err := parseDatasourceArg("foo.json")
 	assert.NoError(t, err)
-	assert.Equal(t, "foo", key)
+	assert.Equal(t, "foo", alias)
 	assert.Equal(t, "file", ds.URL.Scheme)
 
 	_, _, err = parseDatasourceArg("../foo.json")
@@ -761,60 +839,60 @@ func TestParseDatasourceArgNoAlias(t *testing.T) {
 }
 
 func TestParseDatasourceArgWithAlias(t *testing.T) {
-	key, ds, err := parseDatasourceArg("data=foo.json")
+	alias, ds, err := parseDatasourceArg("data=foo.json")
 	assert.NoError(t, err)
-	assert.Equal(t, "data", key)
+	assert.Equal(t, "data", alias)
 	assert.Equal(t, "file", ds.URL.Scheme)
 	assert.True(t, ds.URL.IsAbs())
 
-	key, ds, err = parseDatasourceArg("data=/otherdir/foo.json")
+	alias, ds, err = parseDatasourceArg("data=/otherdir/foo.json")
 	assert.NoError(t, err)
-	assert.Equal(t, "data", key)
+	assert.Equal(t, "data", alias)
 	assert.Equal(t, "file", ds.URL.Scheme)
 	assert.True(t, ds.URL.IsAbs())
 	assert.Equal(t, "/otherdir/foo.json", ds.URL.Path)
 
 	if runtime.GOOS == "windows" {
-		key, ds, err = parseDatasourceArg("data=foo.json")
+		alias, ds, err = parseDatasourceArg("data=foo.json")
 		assert.NoError(t, err)
-		assert.Equal(t, "data", key)
+		assert.Equal(t, "data", alias)
 		assert.Equal(t, "file", ds.URL.Scheme)
 		assert.True(t, ds.URL.IsAbs())
 		assert.Equalf(t, byte(':'), ds.URL.Path[1], "Path was %s", ds.URL.Path)
 
-		key, ds, err = parseDatasourceArg(`data=\otherdir\foo.json`)
+		alias, ds, err = parseDatasourceArg(`data=\otherdir\foo.json`)
 		assert.NoError(t, err)
-		assert.Equal(t, "data", key)
+		assert.Equal(t, "data", alias)
 		assert.Equal(t, "file", ds.URL.Scheme)
 		assert.True(t, ds.URL.IsAbs())
 		assert.Equal(t, `/otherdir/foo.json`, ds.URL.Path)
 
-		key, ds, err = parseDatasourceArg("data=C:\\windowsdir\\foo.json")
+		alias, ds, err = parseDatasourceArg("data=C:\\windowsdir\\foo.json")
 		assert.NoError(t, err)
-		assert.Equal(t, "data", key)
+		assert.Equal(t, "data", alias)
 		assert.Equal(t, "file", ds.URL.Scheme)
 		assert.True(t, ds.URL.IsAbs())
 		assert.Equal(t, "C:/windowsdir/foo.json", ds.URL.Path)
 
-		key, ds, err = parseDatasourceArg("data=\\\\somehost\\share\\foo.json")
+		alias, ds, err = parseDatasourceArg("data=\\\\somehost\\share\\foo.json")
 		assert.NoError(t, err)
-		assert.Equal(t, "data", key)
+		assert.Equal(t, "data", alias)
 		assert.Equal(t, "file", ds.URL.Scheme)
 		assert.Equal(t, "somehost", ds.URL.Host)
 		assert.True(t, ds.URL.IsAbs())
 		assert.Equal(t, "/share/foo.json", ds.URL.Path)
 	}
 
-	key, ds, err = parseDatasourceArg("data=sftp://example.com/blahblah/foo.json")
+	alias, ds, err = parseDatasourceArg("data=sftp://example.com/blahblah/foo.json")
 	assert.NoError(t, err)
-	assert.Equal(t, "data", key)
+	assert.Equal(t, "data", alias)
 	assert.Equal(t, "sftp", ds.URL.Scheme)
 	assert.True(t, ds.URL.IsAbs())
 	assert.Equal(t, "/blahblah/foo.json", ds.URL.Path)
 
-	key, ds, err = parseDatasourceArg("merged=merge:./foo.yaml|http://example.com/bar.json%3Ffoo=bar")
+	alias, ds, err = parseDatasourceArg("merged=merge:./foo.yaml|http://example.com/bar.json%3Ffoo=bar")
 	assert.NoError(t, err)
-	assert.Equal(t, "merged", key)
+	assert.Equal(t, "merged", alias)
 	assert.Equal(t, "merge", ds.URL.Scheme)
 	assert.Equal(t, "./foo.yaml|http://example.com/bar.json%3Ffoo=bar", ds.URL.Opaque)
 }
diff --git a/internal/config/types.go b/internal/config/types.go
new file mode 100644
index 00000000..042ba533
--- /dev/null
+++ b/internal/config/types.go
@@ -0,0 +1,93 @@
+package config
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+
+	"gopkg.in/yaml.v3"
+)
+
+// Templates - a map of templates. We can't just use map[string]DataSource,
+// because we need to be able to marshal both the old (array of '[k=]v' strings)
+// and the new (proper map) formats.
+//
+// Note that templates use the DataSource type, since they have the exact same
+// shape.
+// TODO: get rid of this and just use map[string]DataSource once the legacy
+// [k=]v array format is no longer supported
+type Templates map[string]DataSource
+
+// UnmarshalYAML - satisfy the yaml.Umarshaler interface
+func (t *Templates) UnmarshalYAML(value *yaml.Node) error {
+	// first attempt to unmarshal as a map[string]DataSource
+	err := value.Decode(map[string]DataSource(*t))
+	if err == nil {
+		return nil
+	}
+
+	// if that fails, try to unmarshal as an array of '[k=]v' strings
+	err = t.unmarshalYAMLArray(value)
+	if err != nil {
+		return fmt.Errorf("could not unmarshal templates as map or array: %w", err)
+	}
+
+	return nil
+}
+
+func (t *Templates) unmarshalYAMLArray(value *yaml.Node) error {
+	a := []string{}
+	err := value.Decode(&a)
+	if err != nil {
+		return fmt.Errorf("could not unmarshal templates as array: %w", err)
+	}
+
+	ts := Templates{}
+	for _, s := range a {
+		alias, pth, _ := strings.Cut(s, "=")
+		if pth == "" {
+			// when alias is omitted, the path and alias are identical
+			pth = alias
+		}
+
+		u, err := ParseSourceURL(pth)
+		if err != nil {
+			return fmt.Errorf("could not parse template URL %q: %w", pth, err)
+		}
+
+		ts[alias] = DataSource{
+			URL: u,
+		}
+	}
+
+	*t = ts
+
+	return nil
+}
+
+func (t Templates) MarshalYAML() (interface{}, error) {
+	type rawTemplate struct {
+		Header http.Header `yaml:"header,omitempty,flow"`
+		URL    string      `yaml:"url"`
+	}
+
+	m := map[string]rawTemplate{}
+	for k, v := range t {
+		m[k] = rawTemplate{
+			Header: v.Header,
+			URL:    v.URL.String(),
+		}
+	}
+	return m, nil
+}
+
+func parseTemplateArg(value string) (alias string, ds DataSource, err error) {
+	alias, u, _ := strings.Cut(value, "=")
+	if u == "" {
+		u = alias
+	}
+
+	ds.URL, err = ParseSourceURL(u)
+
+	return alias, ds, err
+}
diff --git a/internal/config/types_test.go b/internal/config/types_test.go
new file mode 100644
index 00000000..93afd068
--- /dev/null
+++ b/internal/config/types_test.go
@@ -0,0 +1,88 @@
+package config
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gopkg.in/yaml.v3"
+)
+
+func TestTemplates_UnmarshalYAML(t *testing.T) {
+	in := `t:
+  url: foo/bar/helloworld.tmpl
+templatedir:
+  url: templatedir/
+dir:
+  url: foo/bar/
+mytemplate.t:
+  url: mytemplate.t
+remote:
+  url: https://example.com/foo/bar/helloworld.tmpl
+  header:
+    Accept: [text/plain, text/template]`
+	out := Templates{}
+	err := yaml.Unmarshal([]byte(in), &out)
+	assert.NoError(t, err)
+	assert.EqualValues(t, Templates{
+		"t":            {URL: mustURL("foo/bar/helloworld.tmpl")},
+		"templatedir":  {URL: mustURL("templatedir/")},
+		"dir":          {URL: mustURL("foo/bar/")},
+		"mytemplate.t": {URL: mustURL("mytemplate.t")},
+		"remote": {
+			URL:    mustURL("https://example.com/foo/bar/helloworld.tmpl"),
+			Header: http.Header{"Accept": {"text/plain", "text/template"}},
+		},
+	}, out)
+
+	// legacy array format
+	in = `- t=foo/bar/helloworld.tmpl
+- templatedir/
+- dir=foo/bar/
+- mytemplate.t
+- remote=https://example.com/foo/bar/helloworld.tmpl`
+	out = Templates{}
+	err = yaml.Unmarshal([]byte(in), &out)
+	assert.NoError(t, err)
+	assert.EqualValues(t, Templates{
+		"t":            {URL: mustURL("foo/bar/helloworld.tmpl")},
+		"templatedir/": {URL: mustURL("templatedir/")},
+		"dir":          {URL: mustURL("foo/bar/")},
+		"mytemplate.t": {URL: mustURL("mytemplate.t")},
+		"remote":       {URL: mustURL("https://example.com/foo/bar/helloworld.tmpl")},
+	}, out)
+
+	// invalid format
+	in = `"neither an array nor a map"`
+	out = Templates{}
+	err = yaml.Unmarshal([]byte(in), &out)
+	assert.Error(t, err)
+
+	// invalid URL
+	in = `- t="not a:valid url"`
+	out = Templates{}
+	err = yaml.Unmarshal([]byte(in), &out)
+	assert.Error(t, err)
+}
+
+func TestParseTemplateArg(t *testing.T) {
+	data := []struct {
+		ds    DataSource
+		in    string
+		alias string
+	}{
+		{in: "t=foo/bar/helloworld.tmpl", alias: "t", ds: DataSource{URL: mustURL("foo/bar/helloworld.tmpl")}},
+		{in: "templatedir/", alias: "templatedir/", ds: DataSource{URL: mustURL("templatedir/")}},
+		{in: "dir=foo/bar/", alias: "dir", ds: DataSource{URL: mustURL("foo/bar/")}},
+		{in: "mytemplate.t", alias: "mytemplate.t", ds: DataSource{URL: mustURL("mytemplate.t")}},
+		{in: "remote=https://example.com/foo/bar/helloworld.tmpl",
+			alias: "remote", ds: DataSource{URL: mustURL("https://example.com/foo/bar/helloworld.tmpl")}},
+	}
+
+	for _, d := range data {
+		alias, ds, err := parseTemplateArg(d.in)
+		assert.NoError(t, err)
+		assert.Equal(t, d.alias, alias)
+		assert.EqualValues(t, d.ds, ds)
+	}
+}