Initial import

internal/netutil/domain.go

@@ -0,0 +1,99 @@
+package netutil
+
+import (
+	"strings"
+
+	"github.com/miekg/dns"
+)
+
+type DomainTree struct {
+	root *domainTreeNode
+}
+
+type domainTreeNode struct {
+	leaf  map[string]*domainTreeNode
+	isEnd bool
+}
+
+func NewDomainList(domains ...string) *DomainTree {
+	tree := &DomainTree{
+		root: &domainTreeNode{leaf: make(map[string]*domainTreeNode)},
+	}
+	for _, domain := range domains {
+		tree.Add(domain)
+	}
+	return tree
+}
+
+func (tree *DomainTree) Add(domain string) {
+	domain = normalizeDomain(domain)
+	if domain == "" {
+		return
+	}
+
+	labels := dns.SplitDomainName(domain)
+	if len(labels) == 0 {
+		return
+	}
+
+	node := tree.root
+	for i := len(labels) - 1; i >= 0; i-- {
+		label := labels[i]
+		if label == "" {
+			continue
+		}
+		if node.leaf == nil {
+			node.leaf = make(map[string]*domainTreeNode)
+		}
+		if node.leaf[label] == nil {
+			node.leaf[label] = &domainTreeNode{}
+		}
+		node = node.leaf[label]
+	}
+	node.isEnd = true
+}
+
+func (tree *DomainTree) Contains(domain string) bool {
+	domain = normalizeDomain(domain)
+	if domain == "" {
+		return false
+	}
+
+	labels := dns.SplitDomainName(domain)
+	if len(labels) == 0 {
+		return false
+	}
+
+	node := tree.root
+	for i := len(labels) - 1; i >= 0; i-- {
+		if node.isEnd {
+			return true
+		}
+
+		if node.leaf == nil {
+			return false
+		}
+
+		label := labels[i]
+		if node = node.leaf[label]; node == nil {
+			return false
+		}
+	}
+	return node.isEnd
+}
+
+func normalizeDomain(domain string) string {
+	domain = strings.ToLower(strings.TrimSpace(domain))
+	if domain == "" {
+		return ""
+	}
+
+	// Remove trailing dot if present, dns.Fqdn will add it back properly
+	domain = strings.TrimSuffix(domain, ".")
+
+	if domain == "" {
+		return ""
+	}
+
+	return dns.Fqdn(domain)
+}