Fix: domain trie search

alvin-up · Jan 26, 2022 · b1a639f · b1a639f
1 parent cfe7354
commit b1a639f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/component/trie/domain.go b/component/trie/domain.go
@@ -109,13 +109,13 @@ func (t *DomainTrie) search(node *Node, parts []string) *Node {
 	}
 
 	if c := node.getChild(parts[len(parts)-1]); c != nil {
-		if n := t.search(c, parts[:len(parts)-1]); n != nil {
+		if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != nil {
 			return n
 		}
 	}
 
 	if c := node.getChild(wildcard); c != nil {
-		if n := t.search(c, parts[:len(parts)-1]); n != nil {
+		if n := t.search(c, parts[:len(parts)-1]); n != nil && n.Data != nil {
 			return n
 		}
 	}

diff --git a/component/trie/domain_test.go b/component/trie/domain_test.go
@@ -97,3 +97,11 @@ func TestTrie_Boundary(t *testing.T) {
 	assert.NotNil(t, tree.Insert("..dev", localIP))
 	assert.Nil(t, tree.Search("dev"))
 }
+
+func TestTrie_WildcardBoundary(t *testing.T) {
+	tree := New()
+	tree.Insert("+.*", localIP)
+	tree.Insert("stun.*.*.*", localIP)
+
+	assert.NotNil(t, tree.Search("example.com"))
+}
diff --git a/component/trie/node.go b/component/trie/node.go
@@ -2,8 +2,8 @@ package trie
 
 // Node is the trie's node
 type Node struct {
-	Data     interface{}
 	children map[string]*Node
+	Data     interface{}
 }
 
 func (n *Node) getChild(s string) *Node {