Feature: domain trie support wildcard alias

component/trie/domain.go

@@ -0,0 +1,139 @@
+package trie
+
+import (
+	"errors"
+	"strings"
+)
+
+const (
+	wildcard        = "*"
+	dotWildcard     = ""
+	complexWildcard = "+"
+	domainStep      = "."
+)
+
+var (
+	// ErrInvalidDomain means insert domain is invalid
+	ErrInvalidDomain = errors.New("invalid domain")
+)
+
+// DomainTrie contains the main logic for adding and searching nodes for domain segments.
+// support wildcard domain (e.g *.google.com)
+type DomainTrie struct {
+	root *Node
+}
+
+func validAndSplitDomain(domain string) ([]string, bool) {
+	if domain != "" && domain[len(domain)-1] == '.' {
+		return nil, false
+	}
+
+	parts := strings.Split(domain, domainStep)
+	if len(parts) == 1 {
+		return nil, false
+	}
+
+	for _, part := range parts[1:] {
+		if part == "" {
+			return nil, false
+		}
+	}
+
+	return parts, true
+}
+
+// Insert adds a node to the trie.
+// Support
+// 1. www.example.com
+// 2. *.example.com
+// 3. subdomain.*.example.com
+// 4. .example.com
+// 5. +.example.com
+func (t *DomainTrie) Insert(domain string, data interface{}) error {
+	parts, valid := validAndSplitDomain(domain)
+	if !valid {
+		return ErrInvalidDomain
+	}
+
+	if parts[0] == complexWildcard {
+		t.insert(parts[1:], data)
+		parts[0] = dotWildcard
+		t.insert(parts, data)
+	} else {
+		t.insert(parts, data)
+	}
+
+	return nil
+}
+
+func (t *DomainTrie) insert(parts []string, data interface{}) {
+	node := t.root
+	// reverse storage domain part to save space
+	for i := len(parts) - 1; i >= 0; i-- {
+		part := parts[i]
+		if !node.hasChild(part) {
+			node.addChild(part, newNode(nil))
+		}
+
+		node = node.getChild(part)
+	}
+
+	node.Data = data
+}
+
+// Search is the most important part of the Trie.
+// Priority as:
+// 1. static part
+// 2. wildcard domain
+// 2. dot wildcard domain
+func (t *DomainTrie) Search(domain string) *Node {
+	parts, valid := validAndSplitDomain(domain)
+	if !valid || parts[0] == "" {
+		return nil
+	}
+
+	n := t.root
+	var dotWildcardNode *Node
+	var wildcardNode *Node
+	for i := len(parts) - 1; i >= 0; i-- {
+		part := parts[i]
+
+		if node := n.getChild(dotWildcard); node != nil {
+			dotWildcardNode = node
+		}
+
+		child := n.getChild(part)
+		if child == nil && wildcardNode != nil {
+			child = wildcardNode.getChild(part)
+		}
+		wildcardNode = n.getChild(wildcard)
+
+		n = child
+		if n == nil {
+			n = wildcardNode
+			wildcardNode = nil
+		}
+
+		if n == nil {
+			break
+		}
+	}
+
+	if n == nil {
+		if dotWildcardNode != nil {
+			return dotWildcardNode
+		}
+		return nil
+	}
+
+	if n.Data == nil {
+		return nil
+	}
+
+	return n
+}
+
+// New returns a new, empty Trie.
+func New() *DomainTrie {
+	return &DomainTrie{root: newNode(nil)}
+}

component/trie/domain_test.go

@@ -0,0 +1,90 @@
+package trie
+
+import (
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var localIP = net.IP{127, 0, 0, 1}
+
+func TestTrie_Basic(t *testing.T) {
+	tree := New()
+	domains := []string{
+		"example.com",
+		"google.com",
+	}
+
+	for _, domain := range domains {
+		tree.Insert(domain, localIP)
+	}
+
+	node := tree.Search("example.com")
+	assert.NotNil(t, node)
+	assert.True(t, node.Data.(net.IP).Equal(localIP))
+	assert.NotNil(t, tree.Insert("", localIP))
+}
+
+func TestTrie_Wildcard(t *testing.T) {
+	tree := New()
+	domains := []string{
+		"*.example.com",
+		"sub.*.example.com",
+		"*.dev",
+		".org",
+		".example.net",
+		".apple.*",
+		"+.foo.com",
+	}
+
+	for _, domain := range domains {
+		tree.Insert(domain, localIP)
+	}
+
+	assert.NotNil(t, tree.Search("sub.example.com"))
+	assert.NotNil(t, tree.Search("sub.foo.example.com"))
+	assert.NotNil(t, tree.Search("test.org"))
+	assert.NotNil(t, tree.Search("test.example.net"))
+	assert.NotNil(t, tree.Search("test.apple.com"))
+	assert.NotNil(t, tree.Search("test.foo.com"))
+	assert.NotNil(t, tree.Search("foo.com"))
+	assert.Nil(t, tree.Search("foo.sub.example.com"))
+	assert.Nil(t, tree.Search("foo.example.dev"))
+	assert.Nil(t, tree.Search("example.com"))
+}
+
+func TestTrie_Priority(t *testing.T) {
+	tree := New()
+	domains := []string{
+		".dev",
+		"example.dev",
+		"*.example.dev",
+		"test.example.dev",
+	}
+
+	assertFn := func(domain string, data int) {
+		node := tree.Search(domain)
+		assert.NotNil(t, node)
+		assert.Equal(t, data, node.Data)
+	}
+
+	for idx, domain := range domains {
+		tree.Insert(domain, idx)
+	}
+
+	assertFn("test.dev", 0)
+	assertFn("foo.bar.dev", 0)
+	assertFn("example.dev", 1)
+	assertFn("foo.example.dev", 2)
+	assertFn("test.example.dev", 3)
+}
+
+func TestTrie_Boundary(t *testing.T) {
+	tree := New()
+	tree.Insert("*.dev", localIP)
+
+	assert.NotNil(t, tree.Insert(".", localIP))
+	assert.NotNil(t, tree.Insert("..dev", localIP))
+	assert.Nil(t, tree.Search("dev"))
+}

component/trie/node.go

@@ -0,0 +1,26 @@
+package trie
+
+// Node is the trie's node
+type Node struct {
+	Data     interface{}
+	children map[string]*Node
+}
+
+func (n *Node) getChild(s string) *Node {
+	return n.children[s]
+}
+
+func (n *Node) hasChild(s string) bool {
+	return n.getChild(s) != nil
+}
+
+func (n *Node) addChild(s string, child *Node) {
+	n.children[s] = child
+}
+
+func newNode(data interface{}) *Node {
+	return &Node{
+		Data:     data,
+		children: map[string]*Node{},
+	}
+}