Fix domain list generation within public suffix

service/intel/entity.go

@@ -397,28 +397,32 @@ func (e *Entity) getDomainLists(ctx context.Context) {
 }
 
 func splitDomain(domain string) []string {
-	domain = strings.Trim(domain, ".")
+	// Get suffix.
-	suffix, _ := publicsuffix.PublicSuffix(domain)
+	d := strings.TrimSuffix(domain, ".")
-	if suffix == domain {
+	suffix, icann := publicsuffix.PublicSuffix(d)
+	if suffix == d {
 		return []string{domain}
 	}
 
-	domainWithoutSuffix := domain[:len(domain)-len(suffix)]
+	// Split all subdomain into labels.
-	domainWithoutSuffix = strings.Trim(domainWithoutSuffix, ".")
+	labels := strings.FieldsFunc(d[:len(d)-len(suffix)], func(r rune) bool {
-
-	splitted := strings.FieldsFunc(domainWithoutSuffix, func(r rune) bool {
 		return r == '.'
 	})
 
-	domains := make([]string, 0, len(splitted))
+	// Build list of all domains up to the public suffix.
-	for idx := range splitted {
+	domains := make([]string, 0, len(labels)+1)
-
+	for idx := range labels {
-		d := strings.Join(splitted[idx:], ".") + "." + suffix
+		domains = append(
-		if d[len(d)-1] != '.' {
+			domains,
-			d += "."
+			strings.Join(labels[idx:], ".")+"."+suffix+".",
-		}
+		)
-		domains = append(domains, d)
 	}
+
+	// If the suffix is not a real TLD, but a public suffix, add it to the list.
+	if !icann {
+		domains = append(domains, suffix+".")
+	}
+
 	return domains
 }
 

service/intel/entity_test.go

@@ -0,0 +1,31 @@
+package intel
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var splitDomainTestCases = [][]string{
+	// Regular registered domains and subdomains.
+	{"example.com."},
+	{"www.example.com.", "example.com."},
+	{"sub.domain.example.com.", "domain.example.com.", "example.com."},
+	{"example.co.uk."},
+	{"www.example.co.uk.", "example.co.uk."},
+
+	// TLD or public suffix: Return as is.
+	{"com."},
+	{"googleapis.com."},
+
+	// Public suffix domains: Return including public suffix.
+	{"test.googleapis.com.", "googleapis.com."},
+	{"sub.domain.googleapis.com.", "domain.googleapis.com.", "googleapis.com."},
+}
+
+func TestSplitDomain(t *testing.T) {
+	for _, testCase := range splitDomainTestCases {
+		splitted := splitDomain(testCase[0])
+		assert.Equal(t, testCase, splitted, "result must match")
+	}
+}