diff --git a/component/trie/domain.go b/component/trie/domain.go index d9463c6e..86e5245a 100644 --- a/component/trie/domain.go +++ b/component/trie/domain.go @@ -25,7 +25,7 @@ func ValidAndSplitDomain(domain string) ([]string, bool) { if domain != "" && domain[len(domain)-1] == '.' { return nil, false } - + domain=strings.ToLower(domain) parts := strings.Split(domain, domainStep) if len(parts) == 1 { if parts[0] == "" { @@ -123,6 +123,30 @@ func (t *DomainTrie[T]) Optimize() { t.root.optimize() } +func (t *DomainTrie[T]) Foreach(print func(domain string, data T)) { + for key, data := range t.root.getChildren() { + recursion([]string{key}, data, print) + } +} + +func recursion[T any](items []string, node *Node[T], fn func(domain string, data T)) { + for key, data := range node.getChildren() { + newItems := append([]string{key}, items...) + if data != nil && data.inited { + domain := joinDomain(newItems) + if domain[0] == domainStepByte { + domain = complexWildcard + domain + } + fn(domain, data.Data()) + } + recursion(newItems, data, fn) + } +} + +func joinDomain(items []string) string { + return strings.Join(items, domainStep) +} + // New returns a new, empty Trie. func New[T any]() *DomainTrie[T] { return &DomainTrie[T]{root: newNode[T]()} diff --git a/component/trie/domain_test.go b/component/trie/domain_test.go index c54b3d3b..2dfd1c34 100644 --- a/component/trie/domain_test.go +++ b/component/trie/domain_test.go @@ -105,3 +105,23 @@ func TestTrie_WildcardBoundary(t *testing.T) { assert.NotNil(t, tree.Search("example.com")) } + +func TestTrie_Foreach(t *testing.T) { + tree := New[netip.Addr]() + domainList := []string{ + "google.com", + "stun.*.*.*", + "test.*.google.com", + "+.baidu.com", + "*.baidu.com", + "*.*.baidu.com", + } + for _, domain := range domainList { + tree.Insert(domain, localIP) + } + count := 0 + tree.Foreach(func(domain string, data netip.Addr) { + count++ + }) + assert.Equal(t, 7, count) +} diff --git a/component/trie/node.go b/component/trie/node.go index e19b40ac..3aa2bc7d 100644 --- a/component/trie/node.go +++ b/component/trie/node.go @@ -116,6 +116,18 @@ func (n *Node[T]) setData(data T) { n.inited = true } +func (n *Node[T]) getChildren() map[string]*Node[T] { + if n.childMap == nil { + if n.childNode != nil { + m := make(map[string]*Node[T]) + m[n.childStr] = n.childNode + return m + } + } else { + return n.childMap + } + return nil +} func (n *Node[T]) Data() T { return n.data } diff --git a/component/trie/sskv.go b/component/trie/sskv.go index f6008659..7e86f539 100644 --- a/component/trie/sskv.go +++ b/component/trie/sskv.go @@ -25,38 +25,14 @@ type DomainSet struct { // NewDomainSet creates a new *DomainSet struct, from a slice of sorted strings. func NewDomainSet(keys []string) *DomainSet { - filter := make(map[string]struct{}, len(keys)) + domainTrie := New[struct{}]() + for _, domain := range keys { + domainTrie.Insert(domain, struct{}{}) + } reserveDomains := make([]string, 0, len(keys)) - insert := func(domain string) { - reserveDomain := utils.Reverse(domain) - reserveDomain = strings.ToLower(reserveDomain) - if _, ok := filter[reserveDomain]; !ok { - filter[reserveDomain] = struct{}{} - domains := make([]string, 0, len(reserveDomains)) - if strings.HasSuffix(reserveDomain, ".+") { - for _, domain := range reserveDomains { - if !strings.HasPrefix(domain, reserveDomain[0:len(reserveDomain)-2]) { - domains = append(domains, domain) - } - } - reserveDomains = domains - } - reserveDomains = append(reserveDomains, reserveDomain) - } - } - for _, key := range keys { - items, ok := ValidAndSplitDomain(key) - if !ok { - continue - } - if items[0] == complexWildcard { - domain := strings.Join(items[1:], domainStep) - insert(domain) - } - - domain := strings.Join(items, domainStep) - insert(domain) - } + domainTrie.Foreach(func(domain string, data struct{}) { + reserveDomains = append(reserveDomains, utils.Reverse(domain)) + }) sort.Slice(reserveDomains, func(i, j int) bool { return len(reserveDomains[i]) < len(reserveDomains[j]) })