Fix: domain trie should backtrack to parent if match fail (#758)

This commit is contained in:
Kr328 2020-06-24 18:41:23 +08:00 committed by GitHub
parent 5f3db72422
commit 3dfff84cc3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 30 additions and 33 deletions

View file

@ -96,39 +96,7 @@ func (t *DomainTrie) Search(domain string) *Node {
return nil return nil
} }
n := t.root n := t.search(t.root, parts)
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 { if n.Data == nil {
return nil return nil
@ -137,6 +105,30 @@ func (t *DomainTrie) Search(domain string) *Node {
return n return n
} }
func (t *DomainTrie) search(node *Node, parts []string) *Node {
if len(parts) == 0 {
return node
}
if c := node.getChild(parts[len(parts)-1]); c != nil {
if n := t.search(c, parts[:len(parts)-1]); n != nil {
return n
}
}
if c := node.getChild(wildcard); c != nil {
if n := t.search(c, parts[:len(parts)-1]); n != nil {
return n
}
}
if c := node.getChild(dotWildcard); c != nil {
return c
}
return nil
}
// New returns a new, empty Trie. // New returns a new, empty Trie.
func New() *DomainTrie { func New() *DomainTrie {
return &DomainTrie{root: newNode(nil)} return &DomainTrie{root: newNode(nil)}

View file

@ -39,6 +39,10 @@ func TestTrie_Wildcard(t *testing.T) {
".example.net", ".example.net",
".apple.*", ".apple.*",
"+.foo.com", "+.foo.com",
"+.stun.*.*",
"+.stun.*.*.*",
"+.stun.*.*.*.*",
"stun.l.google.com",
} }
for _, domain := range domains { for _, domain := range domains {
@ -52,6 +56,7 @@ func TestTrie_Wildcard(t *testing.T) {
assert.NotNil(t, tree.Search("test.apple.com")) assert.NotNil(t, tree.Search("test.apple.com"))
assert.NotNil(t, tree.Search("test.foo.com")) assert.NotNil(t, tree.Search("test.foo.com"))
assert.NotNil(t, tree.Search("foo.com")) assert.NotNil(t, tree.Search("foo.com"))
assert.NotNil(t, tree.Search("global.stun.website.com"))
assert.Nil(t, tree.Search("foo.sub.example.com")) assert.Nil(t, tree.Search("foo.sub.example.com"))
assert.Nil(t, tree.Search("foo.example.dev")) assert.Nil(t, tree.Search("foo.example.dev"))
assert.Nil(t, tree.Search("example.com")) assert.Nil(t, tree.Search("example.com"))