From 4948f3f2135dc1744dceea3041c07d9c4f38c765 Mon Sep 17 00:00:00 2001
From: Skyxim <noreply@skyxim.com>
Date: Tue, 11 Oct 2022 21:35:26 +0800
Subject: [PATCH] chore: Cache and skip multiple failed addresses

---
 component/sniffer/dispatcher.go | 49 +++++++++++++++++++++++++++------
 1 file changed, 41 insertions(+), 8 deletions(-)

diff --git a/component/sniffer/dispatcher.go b/component/sniffer/dispatcher.go
index 82a400e6..a860914d 100644
--- a/component/sniffer/dispatcher.go
+++ b/component/sniffer/dispatcher.go
@@ -2,10 +2,13 @@ package sniffer
 
 import (
 	"errors"
+	"fmt"
+	"github.com/Dreamacro/clash/common/cache"
 	"github.com/Dreamacro/clash/constant/sniffer"
 	"net"
 	"net/netip"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/Dreamacro/clash/component/trie"
@@ -33,6 +36,8 @@ type (
 		foreDomain *trie.DomainTrie[bool]
 		skipSNI    *trie.DomainTrie[bool]
 		portRanges *[]utils.Range[uint16]
+		skipList   *cache.LruCache[string, uint8]
+		rwMux      sync.RWMutex
 	}
 )
 
@@ -43,6 +48,7 @@ func (sd *SnifferDispatcher) TCPSniff(conn net.Conn, metadata *C.Metadata) {
 	}
 
 	if metadata.Host == "" || sd.foreDomain.Search(metadata.Host) != nil {
+
 		port, err := strconv.ParseUint(metadata.DstPort, 10, 16)
 		if err != nil {
 			log.Debugln("[Sniffer] Dst port is error")
@@ -61,7 +67,17 @@ func (sd *SnifferDispatcher) TCPSniff(conn net.Conn, metadata *C.Metadata) {
 			return
 		}
 
+		sd.rwMux.RLock()
+		dst := fmt.Sprintf("%s:%s", metadata.DstIP, metadata.DstPort)
+		if count, ok := sd.skipList.Get(dst); ok && count > 5 {
+			log.Debugln("[Sniffer] Skip sniffing[%s] due to multiple failures", dst)
+			defer sd.rwMux.RUnlock()
+			return
+		}
+		sd.rwMux.RUnlock()
+
 		if host, err := sd.sniffDomain(bufConn, metadata); err != nil {
+			sd.cacheSniffFailed(metadata)
 			log.Debugln("[Sniffer] All sniffing sniff failed with from [%s:%s] to [%s:%s]", metadata.SrcIP, metadata.SrcPort, metadata.String(), metadata.DstPort)
 			return
 		} else {
@@ -70,6 +86,10 @@ func (sd *SnifferDispatcher) TCPSniff(conn net.Conn, metadata *C.Metadata) {
 				return
 			}
 
+			sd.rwMux.RLock()
+			sd.skipList.Delete(dst)
+			sd.rwMux.RUnlock()
+
 			sd.replaceDomain(metadata, host)
 		}
 	}
@@ -95,14 +115,15 @@ func (sd *SnifferDispatcher) Enable() bool {
 }
 
 func (sd *SnifferDispatcher) sniffDomain(conn *CN.BufferedConn, metadata *C.Metadata) (string, error) {
-	for _, sniffer := range sd.sniffers {
-		if sniffer.SupportNetwork() == C.TCP {
-			_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
+	for _, s := range sd.sniffers {
+		if s.SupportNetwork() == C.TCP {
+			_ = conn.SetReadDeadline(time.Now().Add(1 * time.Second))
 			_, err := conn.Peek(1)
 			_ = conn.SetReadDeadline(time.Time{})
 			if err != nil {
 				_, ok := err.(*net.OpError)
 				if ok {
+					sd.cacheSniffFailed(metadata)
 					log.Errorln("[Sniffer] [%s] may not have any sent data, Consider adding skip", metadata.DstIP.String())
 					_ = conn.Close()
 				}
@@ -117,15 +138,15 @@ func (sd *SnifferDispatcher) sniffDomain(conn *CN.BufferedConn, metadata *C.Meta
 				continue
 			}
 
-			host, err := sniffer.SniffTCP(bytes)
+			host, err := s.SniffTCP(bytes)
 			if err != nil {
-				//log.Debugln("[Sniffer] [%s] Sniff data failed %s", sniffer.Protocol(), metadata.DstIP)
+				//log.Debugln("[Sniffer] [%s] Sniff data failed %s", s.Protocol(), metadata.DstIP)
 				continue
 			}
 
 			_, err = netip.ParseAddr(host)
 			if err == nil {
-				//log.Debugln("[Sniffer] [%s] Sniff data failed %s", sniffer.Protocol(), metadata.DstIP)
+				//log.Debugln("[Sniffer] [%s] Sniff data failed %s", s.Protocol(), metadata.DstIP)
 				continue
 			}
 
@@ -136,6 +157,17 @@ func (sd *SnifferDispatcher) sniffDomain(conn *CN.BufferedConn, metadata *C.Meta
 	return "", ErrorSniffFailed
 }
 
+func (sd *SnifferDispatcher) cacheSniffFailed(metadata *C.Metadata) {
+	sd.rwMux.Lock()
+	dst := fmt.Sprintf("%s:%s", metadata.DstIP, metadata.DstPort)
+	count, _ := sd.skipList.Get(dst)
+	if count <= 5 {
+		count++
+	}
+	sd.skipList.Set(dst, count)
+	sd.rwMux.Unlock()
+}
+
 func NewCloseSnifferDispatcher() (*SnifferDispatcher, error) {
 	dispatcher := SnifferDispatcher{
 		enable: false,
@@ -151,16 +183,17 @@ func NewSnifferDispatcher(needSniffer []sniffer.Type, forceDomain *trie.DomainTr
 		foreDomain: forceDomain,
 		skipSNI:    skipSNI,
 		portRanges: ports,
+		skipList:   cache.NewLRUCache[string, uint8](cache.WithSize[string, uint8](128), cache.WithAge[string, uint8](600)),
 	}
 
 	for _, snifferName := range needSniffer {
-		sniffer, err := NewSniffer(snifferName)
+		s, err := NewSniffer(snifferName)
 		if err != nil {
 			log.Errorln("Sniffer name[%s] is error", snifferName)
 			return &SnifferDispatcher{enable: false}, err
 		}
 
-		dispatcher.sniffers = append(dispatcher.sniffers, sniffer)
+		dispatcher.sniffers = append(dispatcher.sniffers, s)
 	}
 
 	return &dispatcher, nil