Fix(domain-trie): Incorrect result

This commit is contained in:
Dreamacro 2019-07-15 18:00:51 +08:00
parent 6077e825c5
commit 3497fdaf45
2 changed files with 15 additions and 5 deletions

View file

@ -75,6 +75,10 @@ func (t *Trie) Search(domain string) *Node {
n = child n = child
} }
if n.Data == nil {
return nil
}
return n return n
} }

View file

@ -5,6 +5,8 @@ import (
"testing" "testing"
) )
var localIP = net.IP{127, 0, 0, 1}
func TestTrie_Basic(t *testing.T) { func TestTrie_Basic(t *testing.T) {
tree := New() tree := New()
domains := []string{ domains := []string{
@ -13,7 +15,7 @@ func TestTrie_Basic(t *testing.T) {
} }
for _, domain := range domains { for _, domain := range domains {
tree.Insert(domain, net.ParseIP("127.0.0.1")) tree.Insert(domain, localIP)
} }
node := tree.Search("example.com") node := tree.Search("example.com")
@ -21,7 +23,7 @@ func TestTrie_Basic(t *testing.T) {
t.Error("should not recv nil") t.Error("should not recv nil")
} }
if !node.Data.(net.IP).Equal(net.IP{127, 0, 0, 1}) { if !node.Data.(net.IP).Equal(localIP) {
t.Error("should equal 127.0.0.1") t.Error("should equal 127.0.0.1")
} }
} }
@ -35,7 +37,7 @@ func TestTrie_Wildcard(t *testing.T) {
} }
for _, domain := range domains { for _, domain := range domains {
tree.Insert(domain, nil) tree.Insert(domain, localIP)
} }
if tree.Search("sub.example.com") == nil { if tree.Search("sub.example.com") == nil {
@ -53,13 +55,17 @@ func TestTrie_Wildcard(t *testing.T) {
if tree.Search("foo.example.dev") != nil { if tree.Search("foo.example.dev") != nil {
t.Error("should recv nil") t.Error("should recv nil")
} }
if tree.Search("example.com") != nil {
t.Error("should recv nil")
}
} }
func TestTrie_Boundary(t *testing.T) { func TestTrie_Boundary(t *testing.T) {
tree := New() tree := New()
tree.Insert("*.dev", nil) tree.Insert("*.dev", localIP)
if err := tree.Insert("com", nil); err == nil { if err := tree.Insert("com", localIP); err == nil {
t.Error("should recv err") t.Error("should recv err")
} }