fix: better DomainSet

component/sniffer/dispatcher.go

@@ -28,8 +28,8 @@ var Dispatcher *SnifferDispatcher
 type SnifferDispatcher struct {
 	enable          bool
 	sniffers        map[sniffer.Sniffer]SnifferConfig
-	forceDomain     *trie.Set
+	forceDomain     *trie.DomainSet
-	skipSNI         *trie.Set
+	skipSNI         *trie.DomainSet
 	skipList        *cache.LruCache[string, uint8]
 	rwMux           sync.RWMutex
 	forceDnsMapping bool
@@ -167,7 +167,7 @@ func NewCloseSnifferDispatcher() (*SnifferDispatcher, error) {
 }
 
 func NewSnifferDispatcher(snifferConfig map[sniffer.Type]SnifferConfig,
-	forceDomain *trie.Set, skipSNI *trie.Set,
+	forceDomain *trie.DomainSet, skipSNI *trie.DomainSet,
 	forceDnsMapping bool, parsePureIp bool) (*SnifferDispatcher, error) {
 	dispatcher := SnifferDispatcher{
 		enable:          true,

component/trie/sskv.go

@@ -1,6 +1,7 @@
+package trie
+
 // Package succinct provides several succinct data types.
 // Modify from https://github.com/openacid/succinct/sskv.go
-package trie
 
 import (
 	"sort"
@@ -16,23 +17,23 @@ const (
 	domainStepByte      = byte('.')
 )
 
-type Set struct {
+type DomainSet struct {
 	leaves, labelBitmap []uint64
 	labels              []byte
 	ranks, selects      []int32
-	isEmpty             bool
 }
 
-// NewSet creates a new *Set struct, from a slice of sorted strings.
+// NewDomainSet creates a new *DomainSet struct, from a slice of sorted strings.
-func NewDomainTrieSet(keys []string) *Set {
+func NewDomainSet(keys []string) *DomainSet {
 	filter := make(map[string]struct{}, len(keys))
 	reserveDomains := make([]string, 0, len(keys))
-	filterFunc := func(reserveDomain string) bool {
+	insert := func(reserveDomain string) {
-		_, ok := filter[reserveDomain]
+		reserveDomain = utils.Reverse(reserveDomain)
-		if !ok {
+		reserveDomain = strings.ToLower(reserveDomain)
+		if _, ok := filter[reserveDomain]; !ok {
 			filter[reserveDomain] = struct{}{}
+			reserveDomains = append(reserveDomains, reserveDomain)
 		}
-		return ok
 	}
 	for _, key := range keys {
 		items, ok := ValidAndSplitDomain(key)
@@ -41,27 +42,20 @@ func NewDomainTrieSet(keys []string) *Set {
 		}
 		if items[0] == complexWildcard {
 			domain := strings.Join(items[1:], domainStep)
-			reserveDomain := utils.Reverse(domain)
+			insert(domain)
-			if !filterFunc(reserveDomain) {
-				reserveDomains = append(reserveDomains, reserveDomain)
-			}
 		}
 
 		domain := strings.Join(items, domainStep)
-		reserveDomain := utils.Reverse(domain)
+		insert(domain)
-		if !filterFunc(reserveDomain) {
-			reserveDomains = append(reserveDomains, reserveDomain)
-		}
 	}
 	sort.Slice(reserveDomains, func(i, j int) bool {
 		return len(reserveDomains[i]) < len(reserveDomains[j])
 	})
 	keys = reserveDomains
-	ss := &Set{}
 	if len(keys) == 0 {
-		ss.isEmpty = true
+		return nil
-		return ss
 	}
+	ss := &DomainSet{}
 	lIdx := 0
 
 	type qElt struct{ s, e, col int }
@@ -93,12 +87,13 @@ func NewDomainTrieSet(keys []string) *Set {
 	return ss
 }
 
-// Has query for a key and return whether it presents in the Set.
+// Has query for a key and return whether it presents in the DomainSet.
-func (ss *Set) Has(key string) bool {
+func (ss *DomainSet) Has(key string) bool {
-	if ss.isEmpty {
+	if ss == nil {
 		return false
 	}
 	key = utils.Reverse(key)
+	key = strings.ToLower(key)
 	// no more labels in this node
 	// skip character matching
 	// go to next level
@@ -150,7 +145,7 @@ func getBit(bm []uint64, i int) uint64 {
 }
 
 // init builds pre-calculated cache to speed up rank() and select()
-func (ss *Set) init() {
+func (ss *DomainSet) init() {
 	ss.selects, ss.ranks = bitmap.IndexSelect32R64(ss.labelBitmap)
 }
 

config/config.go

@@ -136,8 +136,8 @@ type IPTables struct {
 type Sniffer struct {
 	Enable          bool
 	Sniffers        map[snifferTypes.Type]SNIFF.SnifferConfig
-	ForceDomain     *trie.Set
+	ForceDomain     *trie.DomainSet
-	SkipDomain      *trie.Set
+	SkipDomain      *trie.DomainSet
 	ForceDnsMapping bool
 	ParsePureIp     bool
 }
@@ -1344,8 +1344,8 @@ func parseSniffer(snifferRaw RawSniffer) (*Sniffer, error) {
 	}
 
 	sniffer.Sniffers = loadSniffer
-	sniffer.ForceDomain = trie.NewDomainTrieSet(snifferRaw.ForceDomain)
+	sniffer.ForceDomain = trie.NewDomainSet(snifferRaw.ForceDomain)
-	sniffer.SkipDomain = trie.NewDomainTrieSet(snifferRaw.SkipDomain)
+	sniffer.SkipDomain = trie.NewDomainSet(snifferRaw.SkipDomain)
 	return sniffer, nil
 }
 

rules/provider/domain_strategy.go

@@ -7,7 +7,7 @@ import (
 
 type domainStrategy struct {
 	count       int
-	domainRules *trie.Set
+	domainRules *trie.DomainSet
 }
 
 func (d *domainStrategy) ShouldFindProcess() bool {
@@ -27,7 +27,7 @@ func (d *domainStrategy) ShouldResolveIP() bool {
 }
 
 func (d *domainStrategy) OnUpdate(rules []string) {
-	domainTrie := trie.NewDomainTrieSet(rules)
+	domainTrie := trie.NewDomainSet(rules)
 	d.domainRules = domainTrie
 	d.count = len(rules)
 }