Add LookupIPV4, LookupIPV6 template functions (kelseyhightower#684)

Add LookupIPV4, LookupIPV6 template functions

resource/template/template_funcs.go

@@ -33,6 +33,8 @@ func newFuncMap() map[string]interface{} {
 	m["replace"] = strings.Replace
 	m["trimSuffix"] = strings.TrimSuffix
 	m["lookupIP"] = LookupIP
+	m["lookupIPV4"] = LookupIPV4
+	m["lookupIPV6"] = LookupIPV6
 	m["lookupSRV"] = LookupSRV
 	m["fileExists"] = isFileExist
 	m["base64Encode"] = Base64Encode
@@ -181,6 +183,26 @@ func LookupIP(data string) []string {
 	return ipStrings
 }
 
+func LookupIPV6(data string) []string {
+	var addresses []string
+	for _, ip := range LookupIP(data) {
+		if strings.Contains(ip, ":") {
+			addresses = append(addresses, ip)
+		}
+	}
+	return addresses
+}
+
+func LookupIPV4(data string) []string {
+	var addresses []string
+	for _, ip := range LookupIP(data) {
+		if strings.Contains(ip, ".") {
+			addresses = append(addresses, ip)
+		}
+	}
+	return addresses
+}
+
 type sortSRV []*net.SRV
 
 func (s sortSRV) Len() int {

resource/template/template_test.go

@@ -109,7 +109,7 @@ type templateTest struct {
 	desc        string                  // description of the test (for helpful errors)
 	toml        string                  // toml file contents
 	tmpl        string                  // template file contents
-	expected    string                  // expected generated file contents
+	expected    interface{}             // expected generated file contents
 	updateStore func(*TemplateResource) // function for setting values in store
 }
 
@@ -615,6 +615,62 @@ dir: /test/data
 			tr.store.Set("/test/data/def", "child")
 		},
 	},
+	templateTest{
+		desc: "ipv4 lookup test",
+		toml: `
+[template]
+src = "test.conf.tmpl"
+dest = "./tmp/test.conf"
+keys = [
+    "/test/data",
+    "/test/data/abc",
+]
+`,
+		tmpl: `
+{{range lookupIPV4 "localhost"}}
+ip: {{.}}
+{{end}}
+`,
+		expected: `
+
+ip: 127.0.0.1
+
+`,
+		updateStore: func(tr *TemplateResource) {
+			tr.store.Set("/test/data", "parent")
+			tr.store.Set("/test/data/def", "child")
+		},
+	},
+	templateTest{
+		desc: "ipv6 lookup test",
+		toml: `
+[template]
+src = "test.conf.tmpl"
+dest = "./tmp/test.conf"
+keys = [
+    "/test/data",
+    "/test/data/abc",
+]
+`,
+		tmpl: `
+{{range lookupIPV6 "localhost"}}
+ip: {{.}}
+{{end}}
+`,
+		expected: [...]string{
+			`
+ip: ::1
+
+`,
+			`
+
+`,
+},
+		updateStore: func(tr *TemplateResource) {
+			tr.store.Set("/test/data", "parent")
+			tr.store.Set("/test/data/def", "child")
+		},
+	},
 	templateTest{
 		desc: "ip lookup test",
 		toml: `
@@ -631,11 +687,25 @@ keys = [
 ip: {{.}}
 {{end}}
 `,
-		expected: `
+		expected: [...]string{
+			`
+
+ip: 127.0.0.1
+
+`,
+			`
 
 ip: 127.0.0.1
 
+ip: ::1
+
 `,
+			`
+
+ip: ::1
+
+`,
+		},
 		updateStore: func(tr *TemplateResource) {
 			tr.store.Set("/test/data", "parent")
 			tr.store.Set("/test/data/def", "child")
@@ -749,8 +819,18 @@ func ExecuteTestTemplate(tt templateTest, t *testing.T) {
 	if err != nil {
 		t.Errorf(tt.desc + ": failed to read StageFile: " + err.Error())
 	}
-	if string(actual) != tt.expected {
-		t.Errorf(fmt.Sprintf("%v: invalid StageFile. Expected %v, actual %v", tt.desc, tt.expected, string(actual)))
+	switch tt.expected.(type) {
+	case string:
+		if string(actual) != tt.expected.(string) {
+			t.Errorf(fmt.Sprintf("%v: invalid StageFile. Expected %v, actual %v", tt.desc, tt.expected, string(actual)))
+		}
+	case []string:
+		for _, expected := range tt.expected.([]string) {
+			if string(actual) == expected {
+				break
+			}
+		}
+		t.Errorf(fmt.Sprintf("%v: invalid StageFile. Possible expected values %v, actual %v", tt.desc, tt.expected, string(actual)))
 	}
 }