fix: domain tree match failed

This commit is contained in:
Skyxim 2022-04-16 11:55:49 +08:00
parent 71a1f5dfbd
commit 1cf9321aa0
2 changed files with 15 additions and 15 deletions

View File

@ -19,17 +19,17 @@ var (
var Dispatcher SnifferDispatcher var Dispatcher SnifferDispatcher
type SnifferDispatcher struct { type SnifferDispatcher struct {
enable bool enable bool
force bool force bool
sniffers []C.Sniffer sniffers []C.Sniffer
reverseDomainTree *trie.DomainTrie[struct{}] domains *trie.DomainTrie[bool]
tcpHandler func(conn *CN.BufferedConn, metadata *C.Metadata) tcpHandler func(conn *CN.BufferedConn, metadata *C.Metadata)
} }
func (sd *SnifferDispatcher) forceReplace(conn *CN.BufferedConn, metadata *C.Metadata) { func (sd *SnifferDispatcher) forceReplace(conn *CN.BufferedConn, metadata *C.Metadata) {
host, err := sd.sniffDomain(conn, metadata) host, err := sd.sniffDomain(conn, metadata)
if err != nil { if err != nil {
log.Debugln("[Sniffer] All sniffing sniff failed with from [%s:%s] to [%s:%s]", metadata.SrcIP, metadata.SrcPort, metadata.DstIP, metadata.DstPort) log.Debugln("[Sniffer] All sniffing sniff failed with from [%s:%s] to [%s:%s]", metadata.SrcIP, metadata.SrcPort, metadata.String(), metadata.DstPort)
return return
} else { } else {
if sd.inReverse(host) { if sd.inReverse(host) {
@ -49,7 +49,7 @@ func (sd *SnifferDispatcher) replace(conn *CN.BufferedConn, metadata *C.Metadata
host, err := sd.sniffDomain(conn, metadata) host, err := sd.sniffDomain(conn, metadata)
if err != nil { if err != nil {
log.Debugln("[Sniffer] All sniffing sniff failed with from [%s:%s] to [%s:%s]", metadata.SrcIP, metadata.SrcPort, metadata.DstIP, metadata.DstPort) log.Debugln("[Sniffer] All sniffing sniff failed with from [%s:%s] to [%s:%s]", metadata.SrcIP, metadata.SrcPort, metadata.String(), metadata.DstPort)
return return
} }
@ -66,7 +66,7 @@ func (sd *SnifferDispatcher) TCPSniff(conn net.Conn, metadata *C.Metadata) {
} }
func (sd *SnifferDispatcher) inReverse(host string) bool { func (sd *SnifferDispatcher) inReverse(host string) bool {
return sd.reverseDomainTree != nil && sd.reverseDomainTree.Search(host) != nil return sd.domains != nil && sd.domains.Search(host) != nil
} }
func (sd *SnifferDispatcher) replaceDomain(metadata *C.Metadata, host string) { func (sd *SnifferDispatcher) replaceDomain(metadata *C.Metadata, host string) {
@ -122,11 +122,11 @@ func NewCloseSnifferDispatcher() (*SnifferDispatcher, error) {
return &dispatcher, nil return &dispatcher, nil
} }
func NewSnifferDispatcher(needSniffer []C.SnifferType, force bool, reverses *trie.DomainTrie[struct{}]) (*SnifferDispatcher, error) { func NewSnifferDispatcher(needSniffer []C.SnifferType, force bool, reverses *trie.DomainTrie[bool]) (*SnifferDispatcher, error) {
dispatcher := SnifferDispatcher{ dispatcher := SnifferDispatcher{
enable: true, enable: true,
force: force, force: force,
reverseDomainTree: reverses, domains: reverses,
} }
for _, snifferName := range needSniffer { for _, snifferName := range needSniffer {

View File

@ -124,7 +124,7 @@ type Sniffer struct {
Enable bool Enable bool
Force bool Force bool
Sniffers []C.SnifferType Sniffers []C.SnifferType
Reverses *trie.DomainTrie[struct{}] Reverses *trie.DomainTrie[bool]
} }
// Experimental config // Experimental config
@ -928,9 +928,9 @@ func parseSniffer(snifferRaw SnifferRaw) (*Sniffer, error) {
sniffer.Sniffers = append(sniffer.Sniffers, st) sniffer.Sniffers = append(sniffer.Sniffers, st)
} }
sniffer.Reverses = trie.New[struct{}]() sniffer.Reverses = trie.New[bool]()
for _, domain := range snifferRaw.Reverse { for _, domain := range snifferRaw.Reverse {
err := sniffer.Reverses.Insert(domain, struct{}{}) err := sniffer.Reverses.Insert(domain, true)
if err != nil { if err != nil {
return nil, fmt.Errorf("error domian[%s], error:%v", domain, err) return nil, fmt.Errorf("error domian[%s], error:%v", domain, err)
} }