Fix(domain-trie): domain could without dot

This commit is contained in:
Dreamacro 2019-07-16 00:57:08 +08:00
parent 3497fdaf45
commit 9e77c650d9
2 changed files with 17 additions and 5 deletions

View file

@ -21,17 +21,21 @@ type Trie struct {
root *Node root *Node
} }
func isValidDomain(domain string) bool {
return domain[0] != '.' && domain[len(domain)-1] != '.'
}
// Insert adds a node to the trie. // Insert adds a node to the trie.
// Support // Support
// 1. www.example.com // 1. www.example.com
// 2. *.example.com // 2. *.example.com
// 3. subdomain.*.example.com // 3. subdomain.*.example.com
func (t *Trie) Insert(domain string, data interface{}) error { func (t *Trie) Insert(domain string, data interface{}) error {
parts := strings.Split(domain, domainStep) if !isValidDomain(domain) {
if len(parts) < 2 {
return ErrInvalidDomain return ErrInvalidDomain
} }
parts := strings.Split(domain, domainStep)
node := t.root node := t.root
// reverse storage domain part to save space // reverse storage domain part to save space
for i := len(parts) - 1; i >= 0; i-- { for i := len(parts) - 1; i >= 0; i-- {
@ -52,10 +56,10 @@ func (t *Trie) Insert(domain string, data interface{}) error {
// 1. static part // 1. static part
// 2. wildcard domain // 2. wildcard domain
func (t *Trie) Search(domain string) *Node { func (t *Trie) Search(domain string) *Node {
parts := strings.Split(domain, domainStep) if !isValidDomain(domain) {
if len(parts) < 2 {
return nil return nil
} }
parts := strings.Split(domain, domainStep)
n := t.root n := t.root
for i := len(parts) - 1; i >= 0; i-- { for i := len(parts) - 1; i >= 0; i-- {

View file

@ -65,11 +65,19 @@ func TestTrie_Boundary(t *testing.T) {
tree := New() tree := New()
tree.Insert("*.dev", localIP) tree.Insert("*.dev", localIP)
if err := tree.Insert("com", localIP); err == nil { if err := tree.Insert(".", localIP); err == nil {
t.Error("should recv err")
}
if err := tree.Insert(".com", localIP); err == nil {
t.Error("should recv err") t.Error("should recv err")
} }
if tree.Search("dev") != nil { if tree.Search("dev") != nil {
t.Error("should recv nil") t.Error("should recv nil")
} }
if tree.Search(".dev") != nil {
t.Error("should recv nil")
}
} }