From ac493072c738bc9c0dcb9a51f47b0b35ac7813e9 Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Wed, 26 Aug 2015 12:19:22 +0100
Subject: [PATCH] Correctly handle CNAME/DNAMEs

---
 sanitize.go      | 17 ++++++++++-------
 sanitize_test.go | 10 ++++------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/sanitize.go b/sanitize.go
index b4f555d9..cbe8c291 100644
--- a/sanitize.go
+++ b/sanitize.go
@@ -44,7 +44,8 @@ func Dedup(rrs []RR) []RR {
 	}
 	// If the length of the result map equals the amount of RRs we got,
 	// it means they were all different. We can then just return the original rrset.
-	if len(m) == len(rrs) {
+	// We can only do this when we haven't found a CNAME or DNAME.
+	if len(m) == len(rrs) && len(cname) == 0 && len(dname) == 0 {
 		return rrs
 	}
 
@@ -117,22 +118,24 @@ func normalizedString(r RR) (string, int) {
 // needsDeletion checks if the RR is masked by either a CNAME or a DNAME.
 // If so it return true.
 func needsDeletion(r RR, s string, cname, dname []string) bool {
-	// TODO(miek): fix.
-	// This is too broad, but we have to be care full not to delete ourselves.
-	if r.Header().Rrtype == TypeCNAME || r.Header().Rrtype == TypeDNAME {
-		return false
-	}
-
 	// For CNAME we can do strings.HasPrefix with s.
 	// For DNAME we can do strings.Contains with s.
 	// Either signals a removal of this RR.
 	for _, c := range cname {
 		if strings.HasPrefix(s, c) {
+			if r.Header().Rrtype == TypeCNAME {
+				// don't delete yourself
+				continue
+			}
 			return true
 		}
 	}
 	for _, d := range dname {
 		if strings.Contains(s, d) {
+			if r.Header().Rrtype == TypeDNAME && strings.HasPrefix(s, d) {
+				// don't delete yourself
+				continue
+			}
 			return true
 		}
 	}
diff --git a/sanitize_test.go b/sanitize_test.go
index 4b7c10ce..1eae1a30 100644
--- a/sanitize_test.go
+++ b/sanitize_test.go
@@ -39,15 +39,13 @@ func TestDedup(t *testing.T) {
 		},
 	}
 
-	T := 0
 	for rr, expected := range testcases {
 		out := Dedup([]RR{rr[0], rr[1], rr[2]})
 		for i, o := range out {
 			if o.String() != expected[i] {
-				t.Fatalf("test %d, expected %v, got %v", T, expected[i], o.String())
+				t.Fatalf("expected %v, got %v", expected[i], o.String())
 			}
 		}
-		T++
 	}
 }
 
@@ -88,18 +86,18 @@ func TestDedupWithCNAMEDNAME(t *testing.T) {
 		}: []string{"miek.nl.\t3600\tIN\tDNAME\ta."},
 	}
 
-	T := 0
 	for rr, expected := range testcases {
 		out := Dedup([]RR{rr[0], rr[1], rr[2], rr[3]})
 		for i, o := range out {
 			if o.String() != expected[i] {
-				t.Fatalf("test %d, expected %v, got %v", T, expected[i], o.String())
+				t.Fatalf("expected %v, got %v", expected[i], o.String())
 			}
 		}
-		T++
 	}
 }
 
+// BenchMark test as well TODO(miek)
+
 func TestNormalizedString(t *testing.T) {
 	tests := map[RR]string{
 		newRR(t, "mIEk.Nl. 3600 IN A 127.0.0.1"):     "miek.nl.\tIN\tA\t127.0.0.1",