fix: better DomainSet

This commit is contained in:
wwqgtxx 2023-03-29 18:47:56 +08:00
parent c486d7c67e
commit ae8913587d
4 changed files with 28 additions and 33 deletions

View File

@ -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,

View File

@ -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)
}

View File

@ -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
}

View File

@ -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)
}