Add unit tests for ExtractAddressFromReverse (#328)

This fix adds unit tests for `ExtractAddressFromReverse`. It also
updates the implementation of `ExtractAddressFromReverse` so that
only a valid IPv4 or IPv6 address (in string format) is returned.

If the addresses is not valid and can not be handled by golang's
`net.ParseIP` then an empty string (`""`) will be returned.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

2 files changed, 84 insertions, 5 deletions

diff --git a/middleware/pkg/dnsutil/reverse.go b/middleware/pkg/dnsutil/reverse.go
index a360432f3..daf9cc600 100644
--- a/middleware/pkg/dnsutil/reverse.go
+++ b/middleware/pkg/dnsutil/reverse.go
@@ -1,6 +1,9 @@
 package dnsutil
 
-import "strings"
+import (
+	"net"
+	"strings"
+)
 
 // ExtractAddressFromReverse turns a standard PTR reverse record name
 // into an IP address. This works for ipv4 or ipv6.
@@ -10,26 +13,51 @@ import "strings"
 func ExtractAddressFromReverse(reverseName string) string {
 	search := ""
 
+	f := reverse
+
 	switch {
 	case strings.HasSuffix(reverseName, v4arpaSuffix):
 		search = strings.TrimSuffix(reverseName, v4arpaSuffix)
 	case strings.HasSuffix(reverseName, v6arpaSuffix):
 		search = strings.TrimSuffix(reverseName, v6arpaSuffix)
+		f = reverse6
 	default:
 		return ""
 	}
 
 	// Reverse the segments and then combine them.
-	segments := reverse(strings.Split(search, "."))
-	return strings.Join(segments, ".")
+	return f(strings.Split(search, "."))
 }
 
-func reverse(slice []string) []string {
+func reverse(slice []string) string {
 	for i := 0; i < len(slice)/2; i++ {
 		j := len(slice) - i - 1
 		slice[i], slice[j] = slice[j], slice[i]
 	}
-	return slice
+	ip := net.ParseIP(strings.Join(slice, ".")).To4()
+	if ip == nil {
+		return ""
+	}
+	return ip.String()
+}
+
+// reverse6 reverse the segments and combine them according to RFC3596:
+// b.a.9.8.7.6.5.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.8.b.d.0.1.0.0.2
+// is reversed to 2001:db8::567:89ab
+func reverse6(slice []string) string {
+	for i := 0; i < len(slice)/2; i++ {
+		j := len(slice) - i - 1
+		slice[i], slice[j] = slice[j], slice[i]
+	}
+	slice6 := []string{}
+	for i := 0; i < len(slice)/4; i++ {
+		slice6 = append(slice6, strings.Join(slice[i*4:i*4+4], ""))
+	}
+	ip := net.ParseIP(strings.Join(slice6, ":")).To16()
+	if ip == nil {
+		return ""
+	}
+	return ip.String()
 }
 
 const (
diff --git a/middleware/pkg/dnsutil/reverse_test.go b/middleware/pkg/dnsutil/reverse_test.go
new file mode 100644
index 000000000..25bd897ac
--- /dev/null
+++ b/middleware/pkg/dnsutil/reverse_test.go
@@ -0,0 +1,51 @@
+package dnsutil
+
+import (
+	"testing"
+)
+
+func TestExtractAddressFromReverse(t *testing.T) {
+	tests := []struct {
+		reverseName     string
+		expectedAddress string
+	}{
+		{
+			"54.119.58.176.in-addr.arpa.",
+			"176.58.119.54",
+		},
+		{
+			".58.176.in-addr.arpa.",
+			"",
+		},
+		{
+			"b.a.9.8.7.6.5.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.8.b.d.0.1.0.0.2.in-addr.arpa.",
+			"",
+		},
+		{
+			"b.a.9.8.7.6.5.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.8.b.d.0.1.0.0.2.ip6.arpa.",
+			"2001:db8::567:89ab",
+		},
+		{
+			"d.0.1.0.0.2.ip6.arpa.",
+			"",
+		},
+		{
+			"54.119.58.176.ip6.arpa.",
+			"",
+		},
+		{
+			"NONAME",
+			"",
+		},
+		{
+			"",
+			"",
+		},
+	}
+	for i, test := range tests {
+		got := ExtractAddressFromReverse(test.reverseName)
+		if got != test.expectedAddress {
+			t.Errorf("Test %d, expected '%s', got '%s'", i, test.expectedAddress, got)
+		}
+	}
+}