Fix: domain wildcard behavior

This commit is contained in:
Dreamacro 2020-04-24 23:49:19 +08:00
parent 8eddcd77bf
commit 3fc6d55003
2 changed files with 13 additions and 4 deletions

View file

@ -81,6 +81,7 @@ func (t *Trie) Search(domain string) *Node {
n := t.root n := t.root
var dotWildcardNode *Node var dotWildcardNode *Node
var wildcardNode *Node
for i := len(parts) - 1; i >= 0; i-- { for i := len(parts) - 1; i >= 0; i-- {
part := parts[i] part := parts[i]
@ -88,10 +89,16 @@ func (t *Trie) Search(domain string) *Node {
dotWildcardNode = node dotWildcardNode = node
} }
if n.hasChild(part) { child := n.getChild(part)
n = n.getChild(part) if child == nil && wildcardNode != nil {
} else { child = wildcardNode.getChild(part)
n = n.getChild(wildcard) }
wildcardNode = n.getChild(wildcard)
n = child
if n == nil {
n = wildcardNode
wildcardNode = nil
} }
if n == nil { if n == nil {

View file

@ -34,6 +34,7 @@ func TestTrie_Wildcard(t *testing.T) {
"*.dev", "*.dev",
".org", ".org",
".example.net", ".example.net",
".apple.*",
} }
for _, domain := range domains { for _, domain := range domains {
@ -44,6 +45,7 @@ func TestTrie_Wildcard(t *testing.T) {
assert.NotNil(t, tree.Search("sub.foo.example.com")) assert.NotNil(t, tree.Search("sub.foo.example.com"))
assert.NotNil(t, tree.Search("test.org")) assert.NotNil(t, tree.Search("test.org"))
assert.NotNil(t, tree.Search("test.example.net")) assert.NotNil(t, tree.Search("test.example.net"))
assert.NotNil(t, tree.Search("test.apple.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"))