92 lines
1.6 KiB
Go
92 lines
1.6 KiB
Go
package trie
|
|
|
|
import (
|
|
"errors"
|
|
"strings"
|
|
)
|
|
|
|
const (
|
|
wildcard = "*"
|
|
domainStep = "."
|
|
)
|
|
|
|
var (
|
|
// ErrInvalidDomain means insert domain is invalid
|
|
ErrInvalidDomain = errors.New("invalid domain")
|
|
)
|
|
|
|
// Trie contains the main logic for adding and searching nodes for domain segments.
|
|
// support wildcard domain (e.g *.google.com)
|
|
type Trie struct {
|
|
root *Node
|
|
}
|
|
|
|
func isValidDomain(domain string) bool {
|
|
return domain != "" && domain[0] != '.' && domain[len(domain)-1] != '.'
|
|
}
|
|
|
|
// Insert adds a node to the trie.
|
|
// Support
|
|
// 1. www.example.com
|
|
// 2. *.example.com
|
|
// 3. subdomain.*.example.com
|
|
func (t *Trie) Insert(domain string, data interface{}) error {
|
|
if !isValidDomain(domain) {
|
|
return ErrInvalidDomain
|
|
}
|
|
|
|
parts := strings.Split(domain, domainStep)
|
|
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
|
|
return nil
|
|
}
|
|
|
|
// Search is the most important part of the Trie.
|
|
// Priority as:
|
|
// 1. static part
|
|
// 2. wildcard domain
|
|
func (t *Trie) Search(domain string) *Node {
|
|
if !isValidDomain(domain) {
|
|
return nil
|
|
}
|
|
parts := strings.Split(domain, domainStep)
|
|
|
|
n := t.root
|
|
for i := len(parts) - 1; i >= 0; i-- {
|
|
part := parts[i]
|
|
|
|
var child *Node
|
|
if !n.hasChild(part) {
|
|
if !n.hasChild(wildcard) {
|
|
return nil
|
|
}
|
|
|
|
child = n.getChild(wildcard)
|
|
} else {
|
|
child = n.getChild(part)
|
|
}
|
|
|
|
n = child
|
|
}
|
|
|
|
if n.Data == nil {
|
|
return nil
|
|
}
|
|
|
|
return n
|
|
}
|
|
|
|
// New returns a new, empty Trie.
|
|
func New() *Trie {
|
|
return &Trie{root: newNode(nil)}
|
|
}
|