config: careful with addressability and replacing variables

This commit is contained in:
Mitchell Hashimoto 2014-05-23 23:15:00 -07:00
parent 537fa6cc87
commit 2ecf1b500f
2 changed files with 137 additions and 0 deletions

View File

@ -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
}

View File

@ -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)
}
}
}