Fix(domain-trie): domain could without dot
This commit is contained in:
parent
3497fdaf45
commit
9e77c650d9
2 changed files with 17 additions and 5 deletions
|
@ -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-- {
|
||||||
|
|
|
@ -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")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue