Improves flow handling and adds runtime stats APIs

Refactors TCP and UDP flow managers to enhance analyzer selection and flow binding accuracy, including O(1) UDP stream rebinding by 5-tuple. Introduces runtime stats tracking for engine and ruleset operations, exposing new APIs for granular performance and error metrics. Optimizes GeoMatcher with result caching and supports efficient geosite set matching, reducing redundant computation in ruleset expressions.
2026-05-13 06:10:38 +05:30
parent 3f895adb43
commit 7a3f6e945d
23 changed files with 1440 additions and 152 deletions
@@ -0,0 +1,235 @@
+package engine
+
+import (
+	"bytes"
+	"strings"
+
+	"git.difuse.io/Difuse/Mellaris/analyzer"
+)
+
+type analyzerSelector struct {
+	mode  AnalyzerSelectionMode
+	stats *statsCounters
+}
+
+func newAnalyzerSelector(mode AnalyzerSelectionMode, stats *statsCounters) *analyzerSelector {
+	if mode == "" {
+		mode = AnalyzerSelectionModeSignature
+	}
+	return &analyzerSelector{mode: mode, stats: stats}
+}
+
+func (s *analyzerSelector) SelectTCP(ans []analyzer.Analyzer, payload []byte) []analyzer.Analyzer {
+	if s == nil || s.mode == AnalyzerSelectionModeAlways || len(ans) <= 1 {
+		return ans
+	}
+	allowed := tcpAllowedAnalyzers(payload)
+	if len(allowed) == 0 {
+		return ans
+	}
+	out := make([]analyzer.Analyzer, 0, len(ans))
+	for _, a := range ans {
+		name := strings.ToLower(a.Name())
+		if _, known := knownTCPAnalyzers[name]; !known {
+			out = append(out, a)
+			continue
+		}
+		if allowed[name] {
+			out = append(out, a)
+		}
+	}
+	s.recordSelection(len(ans), len(out))
+	if len(out) == 0 {
+		return ans
+	}
+	return out
+}
+
+func (s *analyzerSelector) SelectUDP(ans []analyzer.Analyzer, payload []byte) []analyzer.Analyzer {
+	if s == nil || s.mode == AnalyzerSelectionModeAlways || len(ans) <= 1 {
+		return ans
+	}
+	allowed := udpAllowedAnalyzers(payload)
+	if len(allowed) == 0 {
+		return ans
+	}
+	out := make([]analyzer.Analyzer, 0, len(ans))
+	for _, a := range ans {
+		name := strings.ToLower(a.Name())
+		if _, known := knownUDPAnalyzers[name]; !known {
+			out = append(out, a)
+			continue
+		}
+		if allowed[name] {
+			out = append(out, a)
+		}
+	}
+	s.recordSelection(len(ans), len(out))
+	if len(out) == 0 {
+		return ans
+	}
+	return out
+}
+
+func (s *analyzerSelector) recordSelection(total, selected int) {
+	if s == nil || s.stats == nil || total <= 0 {
+		return
+	}
+	s.stats.AnalyzerSelectionsTotal.Add(1)
+	if selected < total {
+		s.stats.AnalyzerSelectionsPruned.Add(1)
+	}
+}
+
+var (
+	knownTCPAnalyzers = map[string]struct{}{
+		"fet":     {},
+		"http":    {},
+		"socks":   {},
+		"ssh":     {},
+		"tls":     {},
+		"trojan":  {},
+		"dns":     {},
+		"openvpn": {},
+	}
+	knownUDPAnalyzers = map[string]struct{}{
+		"dns":       {},
+		"openvpn":   {},
+		"quic":      {},
+		"wireguard": {},
+	}
+)
+
+func tcpAllowedAnalyzers(payload []byte) map[string]bool {
+	allowed := make(map[string]bool, 4)
+	if looksLikeTLS(payload) {
+		allowed["tls"] = true
+		allowed["trojan"] = true
+		allowed["fet"] = true
+	}
+	if looksLikeHTTP(payload) {
+		allowed["http"] = true
+		allowed["fet"] = true
+	}
+	if looksLikeSSH(payload) {
+		allowed["ssh"] = true
+		allowed["fet"] = true
+	}
+	if looksLikeSOCKS(payload) {
+		allowed["socks"] = true
+		allowed["fet"] = true
+	}
+	if looksLikeDNSTCP(payload) {
+		allowed["dns"] = true
+		allowed["fet"] = true
+	}
+	if len(allowed) == 0 {
+		return nil
+	}
+	return allowed
+}
+
+func udpAllowedAnalyzers(payload []byte) map[string]bool {
+	allowed := make(map[string]bool, 4)
+	if looksLikeWireGuard(payload) {
+		allowed["wireguard"] = true
+	}
+	if looksLikeOpenVPN(payload) {
+		allowed["openvpn"] = true
+	}
+	if looksLikeQUIC(payload) {
+		allowed["quic"] = true
+	}
+	if looksLikeDNSUDP(payload) {
+		allowed["dns"] = true
+	}
+	if len(allowed) == 0 {
+		return nil
+	}
+	return allowed
+}
+
+func looksLikeTLS(payload []byte) bool {
+	if len(payload) < 3 {
+		return false
+	}
+	return (payload[0] == 0x16 || payload[0] == 0x17) && payload[1] == 0x03 && payload[2] <= 0x09
+}
+
+func looksLikeHTTP(payload []byte) bool {
+	if len(payload) < 3 {
+		return false
+	}
+	head := strings.ToUpper(string(payload[:3]))
+	switch head {
+	case "GET", "HEA", "POS", "PUT", "DEL", "CON", "OPT", "TRA", "PAT":
+		return true
+	default:
+		return false
+	}
+}
+
+func looksLikeSSH(payload []byte) bool {
+	return len(payload) >= 4 && bytes.HasPrefix(payload, []byte("SSH-"))
+}
+
+func looksLikeSOCKS(payload []byte) bool {
+	if len(payload) < 2 {
+		return false
+	}
+	return payload[0] == 0x04 || payload[0] == 0x05
+}
+
+func looksLikeDNSTCP(payload []byte) bool {
+	if len(payload) < 14 {
+		return false
+	}
+	msgLen := int(payload[0])<<8 | int(payload[1])
+	if msgLen <= 0 || msgLen+2 > len(payload) {
+		return false
+	}
+	qd := int(payload[6])<<8 | int(payload[7])
+	an := int(payload[8])<<8 | int(payload[9])
+	return qd+an > 0
+}
+
+func looksLikeDNSUDP(payload []byte) bool {
+	if len(payload) < 12 {
+		return false
+	}
+	qd := int(payload[4])<<8 | int(payload[5])
+	an := int(payload[6])<<8 | int(payload[7])
+	ns := int(payload[8])<<8 | int(payload[9])
+	ar := int(payload[10])<<8 | int(payload[11])
+	return qd+an+ns+ar > 0
+}
+
+func looksLikeQUIC(payload []byte) bool {
+	if len(payload) < 6 {
+		return false
+	}
+	// Long header with non-zero version.
+	if payload[0]&0x80 == 0 {
+		return false
+	}
+	version := uint32(payload[1])<<24 | uint32(payload[2])<<16 | uint32(payload[3])<<8 | uint32(payload[4])
+	return version != 0
+}
+
+func looksLikeOpenVPN(payload []byte) bool {
+	if len(payload) == 0 {
+		return false
+	}
+	opcode := payload[0] >> 3
+	return opcode >= 1 && opcode <= 11
+}
+
+func looksLikeWireGuard(payload []byte) bool {
+	if len(payload) < 4 {
+		return false
+	}
+	if payload[0] < 1 || payload[0] > 4 {
+		return false
+	}
+	return payload[1] == 0 && payload[2] == 0 && payload[3] == 0
+}