Add CNAME blocking support

resolver/ipinfo.go

@@ -16,13 +16,92 @@ var (
 	})
 )
 
+// ResolvedDomain holds a Domain name and a list of
+// CNAMES that have been resolved.
+type ResolvedDomain struct {
+	// Domain is the domain as requested by the application.
+	Domain string
+
+	// CNAMEs is a list of CNAMEs that have been resolved for
+	// Domain.
+	CNAMEs []string
+}
+
+// String returns a string representation of ResolvedDomain including
+// the CNAME chain. It implements fmt.Stringer
+func (resolved *ResolvedDomain) String() string {
+	ret := resolved.Domain
+	cnames := ""
+
+	if len(resolved.CNAMEs) > 0 {
+		cnames = " (-> " + strings.Join(resolved.CNAMEs, "->") + ")"
+	}
+
+	return ret + cnames
+}
+
+// ResolvedDomains is a helper type for operating on a slice
+// of ResolvedDomain
+type ResolvedDomains []ResolvedDomain
+
+// String returns a string representation of all domains joined
+// to a single string.
+func (rds ResolvedDomains) String() string {
+	var domains []string
+	for _, n := range rds {
+		domains = append(domains, n.String())
+	}
+	return strings.Join(domains, " or ")
+}
+
+// MostRecentDomain returns the most recent domain.
+func (rds ResolvedDomains) MostRecentDomain() *ResolvedDomain {
+	if len(rds) == 0 {
+		return nil
+	}
+	// TODO(ppacher): we could also do that by using ResolvedAt()
+	mostRecent := rds[len(rds)-1]
+	return &mostRecent
+}
+
 // IPInfo represents various information about an IP.
 type IPInfo struct {
 	record.Base
 	sync.Mutex
 
-	IP      string
-	Domains []string
+	// IP holds the acutal IP address.
+	IP string
+
+	// Domains holds a list of domains that have been
+	// resolved to IP. This field is deprecated and should
+	// be removed.
+	// DEPRECATED: remove with alpha.
+	Domains []string `json:"Domains,omitempty"`
+
+	// ResolvedDomain is a slice of domains that
+	// have been requested by various applications
+	// and have been resolved to IP.
+	ResolvedDomains ResolvedDomains
+}
+
+// AddDomain adds a new resolved domain to ipi.
+func (ipi *IPInfo) AddDomain(resolved ResolvedDomain) bool {
+	for idx, d := range ipi.ResolvedDomains {
+		if d.Domain == resolved.Domain {
+			if utils.StringSliceEqual(d.CNAMEs, resolved.CNAMEs) {
+				return false
+			}
+
+			// we have a different CNAME chain now, remove the previous
+			// entry and add it at the end.
+			ipi.ResolvedDomains = append(ipi.ResolvedDomains[:idx], ipi.ResolvedDomains[idx+1:]...)
+			ipi.ResolvedDomains = append(ipi.ResolvedDomains, resolved)
+			return true
+		}
+	}
+
+	ipi.ResolvedDomains = append(ipi.ResolvedDomains, resolved)
+	return true
 }
 
 func makeIPInfoKey(ip string) string {
@@ -46,6 +125,19 @@ func GetIPInfo(ip string) (*IPInfo, error) {
 		if err != nil {
 			return nil, err
 		}
+
+		// Legacy support,
+		// DEPRECATED: remove with alpha
+		if len(new.Domains) > 0 && len(new.ResolvedDomains) == 0 {
+			for _, d := range new.Domains {
+				new.ResolvedDomains = append(new.ResolvedDomains, ResolvedDomain{
+					Domain: d,
+					// rest is empty...
+				})
+			}
+			new.Domains = nil // clean up so we remove it from the database
+		}
+
 		return new, nil
 	}
 
@@ -57,17 +149,6 @@ func GetIPInfo(ip string) (*IPInfo, error) {
 	return new, nil
 }
 
-// AddDomain adds a domain to the list and reports back if it was added, or was already present.
-func (ipi *IPInfo) AddDomain(domain string) (added bool) {
-	ipi.Lock()
-	defer ipi.Unlock()
-	if !utils.StringInSlice(ipi.Domains, domain) {
-		ipi.Domains = append([]string{domain}, ipi.Domains...)
-		return true
-	}
-	return false
-}
-
 // Save saves the IPInfo record to the database.
 func (ipi *IPInfo) Save() error {
 	ipi.Lock()
@@ -75,17 +156,21 @@ func (ipi *IPInfo) Save() error {
 		ipi.SetKey(makeIPInfoKey(ipi.IP))
 	}
 	ipi.Unlock()
-	return ipInfoDatabase.Put(ipi)
-}
 
-// FmtDomains returns a string consisting of the domains that have seen to use this IP, joined by " or "
-func (ipi *IPInfo) FmtDomains() string {
-	return strings.Join(ipi.Domains, " or ")
+	// Legacy support
+	// Ensure we don't write new Domain fields into the
+	// database.
+	// DEPRECATED: remove with alpha
+	if len(ipi.Domains) > 0 {
+		ipi.Domains = nil
+	}
+
+	return ipInfoDatabase.Put(ipi)
 }
 
 // FmtDomains returns a string consisting of the domains that have seen to use this IP, joined by " or "
 func (ipi *IPInfo) String() string {
 	ipi.Lock()
 	defer ipi.Unlock()
-	return fmt.Sprintf("<IPInfo[%s] %s: %s", ipi.Key(), ipi.IP, ipi.FmtDomains())
+	return fmt.Sprintf("<IPInfo[%s] %s: %s", ipi.Key(), ipi.IP, ipi.ResolvedDomains.String())
 }

resolver/ipinfo_test.go

@@ -1,25 +1,48 @@
 package resolver
 
-import "testing"
+import (
+	"testing"
 
-func testDomains(t *testing.T, ipi *IPInfo, expectedDomains string) {
-	if ipi.FmtDomains() != expectedDomains {
-		t.Errorf("unexpected domains '%s', expected '%s'", ipi.FmtDomains(), expectedDomains)
-	}
-}
+	"github.com/stretchr/testify/assert"
+)
 
 func TestIPInfo(t *testing.T) {
-	ipi := &IPInfo{
-		IP:      "1.2.3.4",
-		Domains: []string{"example.com.", "sub.example.com."},
+	example := ResolvedDomain{
+		Domain: "example.com.",
+	}
+	subExample := ResolvedDomain{
+		Domain: "sub1.example.com",
+		CNAMEs: []string{"example.com"},
 	}
 
-	testDomains(t, ipi, "example.com. or sub.example.com.")
-	ipi.AddDomain("added.example.com.")
-	testDomains(t, ipi, "added.example.com. or example.com. or sub.example.com.")
-	ipi.AddDomain("sub.example.com.")
-	testDomains(t, ipi, "added.example.com. or example.com. or sub.example.com.")
-	ipi.AddDomain("added.example.com.")
-	testDomains(t, ipi, "added.example.com. or example.com. or sub.example.com.")
+	ipi := &IPInfo{
+		IP: "1.2.3.4",
+		ResolvedDomains: ResolvedDomains{
+			example,
+			subExample,
+		},
+	}
 
+	sub2Example := ResolvedDomain{
+		Domain: "sub2.example.com",
+		CNAMEs: []string{"sub1.example.com", "example.com"},
+	}
+	added := ipi.AddDomain(sub2Example)
+
+	assert.True(t, added)
+	assert.Equal(t, ResolvedDomains{example, subExample, sub2Example}, ipi.ResolvedDomains)
+
+	// try again, should do nothing now
+	added = ipi.AddDomain(sub2Example)
+	assert.False(t, added)
+	assert.Equal(t, ResolvedDomains{example, subExample, sub2Example}, ipi.ResolvedDomains)
+
+	subOverWrite := ResolvedDomain{
+		Domain: "sub1.example.com",
+		CNAMEs: []string{}, // now without CNAMEs
+	}
+
+	added = ipi.AddDomain(subOverWrite)
+	assert.True(t, added)
+	assert.Equal(t, ResolvedDomains{example, sub2Example, subOverWrite}, ipi.ResolvedDomains)
 }