From ae8913587de0c63329f80f2f035f237179368cca Mon Sep 17 00:00:00 2001 From: wwqgtxx Date: Wed, 29 Mar 2023 18:47:56 +0800 Subject: [PATCH] fix: better DomainSet --- component/sniffer/dispatcher.go | 6 ++--- component/trie/sskv.go | 43 ++++++++++++++----------------- config/config.go | 8 +++--- rules/provider/domain_strategy.go | 4 +-- 4 files changed, 28 insertions(+), 33 deletions(-) diff --git a/component/sniffer/dispatcher.go b/component/sniffer/dispatcher.go index 3fe40bff..bf0b1bb3 100644 --- a/component/sniffer/dispatcher.go +++ b/component/sniffer/dispatcher.go @@ -28,8 +28,8 @@ var Dispatcher *SnifferDispatcher type SnifferDispatcher struct { enable bool sniffers map[sniffer.Sniffer]SnifferConfig - forceDomain *trie.Set - skipSNI *trie.Set + forceDomain *trie.DomainSet + 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, diff --git a/component/trie/sskv.go b/component/trie/sskv.go index 8784d013..16fc6990 100644 --- a/component/trie/sskv.go +++ b/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. -func NewDomainTrieSet(keys []string) *Set { +// NewDomainSet creates a new *DomainSet struct, from a slice of sorted strings. +func NewDomainSet(keys []string) *DomainSet { filter := make(map[string]struct{}, len(keys)) reserveDomains := make([]string, 0, len(keys)) - filterFunc := func(reserveDomain string) bool { - _, ok := filter[reserveDomain] - if !ok { + insert := func(reserveDomain string) { + reserveDomain = utils.Reverse(reserveDomain) + 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) - if !filterFunc(reserveDomain) { - reserveDomains = append(reserveDomains, reserveDomain) - } + insert(domain) } domain := strings.Join(items, domainStep) - reserveDomain := utils.Reverse(domain) - if !filterFunc(reserveDomain) { - reserveDomains = append(reserveDomains, reserveDomain) - } + insert(domain) } 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 ss + return nil } + 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. -func (ss *Set) Has(key string) bool { - if ss.isEmpty { +// Has query for a key and return whether it presents in the DomainSet. +func (ss *DomainSet) Has(key string) bool { + 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) } diff --git a/config/config.go b/config/config.go index 71050ab5..1083839a 100644 --- a/config/config.go +++ b/config/config.go @@ -136,8 +136,8 @@ type IPTables struct { type Sniffer struct { Enable bool Sniffers map[snifferTypes.Type]SNIFF.SnifferConfig - ForceDomain *trie.Set - SkipDomain *trie.Set + ForceDomain *trie.DomainSet + 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.SkipDomain = trie.NewDomainTrieSet(snifferRaw.SkipDomain) + sniffer.ForceDomain = trie.NewDomainSet(snifferRaw.ForceDomain) + sniffer.SkipDomain = trie.NewDomainSet(snifferRaw.SkipDomain) return sniffer, nil } diff --git a/rules/provider/domain_strategy.go b/rules/provider/domain_strategy.go index 9f4ab0d8..0b2a5d3c 100644 --- a/rules/provider/domain_strategy.go +++ b/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) }