From 2ecf1b500fcc62008414f234fe498de4c7ac3701 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Fri, 23 May 2014 23:15:00 -0700
Subject: [PATCH] config: careful with addressability and replacing variables

---
 config/variable.go      | 83 +++++++++++++++++++++++++++++++++++++++++
 config/variable_test.go | 54 +++++++++++++++++++++++++++
 2 files changed, 137 insertions(+)

diff --git a/config/variable.go b/config/variable.go
index e26fa85bb..571648b32 100644
--- a/config/variable.go
+++ b/config/variable.go
@@ -5,6 +5,8 @@ import (
 	"reflect"
 	"regexp"
 	"strings"
+
+	"github.com/mitchellh/reflectwalk"
 )
 
 // varRegexp is a regexp that matches variables such as ${foo.bar}
@@ -23,6 +25,9 @@ type variableDetectWalker struct {
 
 func (w *variableDetectWalker) Primitive(v reflect.Value) error {
 	// We only care about strings
+	if v.Kind() == reflect.Interface {
+		v = v.Elem()
+	}
 	if v.Kind() != reflect.String {
 		return nil
 	}
@@ -76,3 +81,81 @@ func (w *variableDetectWalker) Primitive(v reflect.Value) error {
 
 	return nil
 }
+
+// variableReplaceWalker implements interfaces for reflectwalk that
+// is used to replace variables with their values.
+//
+// If Values does not have every available value, then the program
+// will _panic_. The variableDetectWalker will tell you all variables
+// you need.
+type variableReplaceWalker struct {
+	Values map[string]string
+
+	loc reflectwalk.Location
+	m, mk reflect.Value
+}
+
+func (w *variableReplaceWalker) Enter(loc reflectwalk.Location) error {
+	w.loc = loc
+	return nil
+}
+
+func (w *variableReplaceWalker) Exit(loc reflectwalk.Location) error {
+	w.loc = reflectwalk.None
+	return nil
+}
+
+func (w *variableReplaceWalker) MapElem(m, k, v reflect.Value) error {
+	w.m = m
+	w.mk = k
+	return nil
+}
+
+func (w *variableReplaceWalker) Primitive(v reflect.Value) error {
+	// We only care about strings
+	setV := v
+	if v.Kind() == reflect.Interface {
+		setV = v
+		v = v.Elem()
+	}
+	if v.Kind() != reflect.String {
+		return nil
+	}
+
+	matches := varRegexp.FindAllStringSubmatch(v.String(), -1)
+	if len(matches) == 0 {
+		return nil
+	}
+
+	result := v.String()
+	for _, match := range matches {
+		dollars := len(match[1])
+
+		// If there are even amounts of dollar signs, then it is escaped
+		if dollars%2 == 0 {
+			continue
+		}
+
+		// Get the key
+		key := match[2]
+		value, ok := w.Values[key]
+		if !ok {
+			panic("no value for variable key: " + key)
+		}
+
+		// Replace
+		result = strings.Replace(result, match[0], value, -1)
+	}
+
+	resultVal := reflect.ValueOf(result)
+	if w.loc == reflectwalk.MapValue {
+		// If we're in a map, then the only way to set a map value is
+		// to set it directly.
+		w.m.SetMapIndex(w.mk, resultVal)
+	} else {
+		// Otherwise, we should be addressable
+		setV.Set(resultVal)
+	}
+
+	return nil
+}
diff --git a/config/variable_test.go b/config/variable_test.go
index a7f11e9f8..564b5dd8f 100644
--- a/config/variable_test.go
+++ b/config/variable_test.go
@@ -3,6 +3,8 @@ package config
 import (
 	"reflect"
 	"testing"
+
+	"github.com/mitchellh/reflectwalk"
 )
 
 func BenchmarkVariableDetectWalker(b *testing.B) {
@@ -83,3 +85,55 @@ func TestVariableDetectWalker_empty(t *testing.T) {
 	}
 }
 
+func TestVariableReplaceWalker(t *testing.T) {
+	w := &variableReplaceWalker{
+		Values: map[string]string{
+			"var.bar": "bar",
+		},
+	}
+
+	cases := []struct {
+		Input  interface{}
+		Output interface{}
+	}{
+		{
+			`foo ${var.bar}`,
+			"foo bar",
+		},
+		{
+			[]string{"foo", "${var.bar}"},
+			[]string{"foo", "bar"},
+		},
+		{
+			map[string]interface{}{
+				"ami": "${var.bar}",
+				"security_groups": []interface{}{
+					"foo",
+					"${var.bar}",
+				},
+			},
+			map[string]interface{}{
+				"ami": "bar",
+				"security_groups": []interface{}{
+					"foo",
+					"bar",
+				},
+			},
+		},
+	}
+
+	for i, tc := range cases {
+		var input interface{} = tc.Input
+		if reflect.ValueOf(tc.Input).Kind() == reflect.String {
+			input = &tc.Input
+		}
+
+		if err := reflectwalk.Walk(input, w); err != nil {
+			t.Fatalf("err: %s", err)
+		}
+
+		if !reflect.DeepEqual(tc.Input, tc.Output) {
+			t.Fatalf("bad %d: %#v", i, tc.Input)
+		}
+	}
+}