fix: add connection statistics tracking in HTTP server

2025-12-16 16:57:08 +00:00 · 2025-11-10 13:26:54 +08:00 · 2025-11-10 13:26:54 +08:00 · 60a5c0801f
commit 60a5c0801f
parent fe2427ddbe
5 changed files with 63 additions and 70 deletions
--- a/README.md
+++ b/README.md
@ -1,6 +1,6 @@
 # UA3F

-UA3F 是下一代 HTTP User-Agent 重写工具，作为一个 HTTP、SOCKS5、TPROXY、REDIRECT、NFQUEUE 服务对 HTTP 进行 User-Agent 透明重写。
+UA3F 是一个 HTTP Header 重写工具，作为一个 HTTP、SOCKS5、TPROXY、REDIRECT、NFQUEUE 服务对 HTTP 请求 Header (例如 User-Agent) 进行透明重写。

 ## 特性

@ -103,7 +103,7 @@ sudo -u shellcrash /usr/bin/ua3f

 UA3F 支持 5 种不同的服务模式，各模式的特点和使用场景如下：

-| 服务模式     | 工作原理           | 是否依赖 Clash 等 | 兼容性 | 性能 | 能否与 Clash 等伴生运行 |
+| 服务模式     | 工作原理           | 是否依赖 Clash 等 | 兼容性 | 效能 | 能否与 Clash 等伴生运行 |
 | ------------ | ------------------ | ----------------- | ------ | ---- | ----------------------- |
 | **HTTP**     | HTTP 代理          | 是                | 高     | 低   | 能                      |
 | **SOCKS5**   | SOCKS5 代理        | 是                | 高     | 低   | 能                      |
@ -111,6 +111,16 @@ UA3F 支持 5 种不同的服务模式，各模式的特点和使用场景如下
 | **REDIRECT** | netfilter REDIRECT | 否                | 中     | 中   | 能                      |
 | **NFQUEUE**  | netfilter NFQUEUE  | 否                | 中     | 高   | 能                      |

+### 重写策略说明
+
+UA3F 支持 3 种不同的重写策略：
+
+| 重写策略   | 重写行为             | 适用服务模式                       |
+| ---------- | -------------------- | ---------------------------------- |
+| **GLOBAL** | 所有请求均进行重写   | 适用于所有服务模式                 |
+| **DIRECT** | 不进行重写，纯转发   | 适用于所有服务模式                 |
+| **RULES**  | 根据重写规则进行重写 | 适用于 HTTP/SOCKS5/TPROXY/REDIRECT |
+
 ## Clash 配置

 Clash 与 UA3F 的配置部署教程详见：[UA3F 与 Clash 从零开始的部署教程](https://sunbk201public.notion.site/UA3F-Clash-16d60a7b5f0e457a9ee97a3be7cbf557?pvs=4)
@ -130,7 +140,7 @@ rules:
 ```

 > [!IMPORTANT]
-> 如果 Clash 使用 Fake-IP 模式，确保 OpenClash 本地 DNS 劫持选择「使用防火墙转发」，不要使用「Dnsmasq 转发」。
+> HTTP/SOCKS5 模式下，如果 Clash 使用 Fake-IP 模式，确保 OpenClash 本地 DNS 劫持选择「使用防火墙转发」，不要使用「Dnsmasq 转发」。

 ### Clash 参考配置

--- a/src/internal/config/config.go
+++ b/src/internal/config/config.go
@ -3,6 +3,7 @@ package config
 import (
 	"flag"
 	"fmt"
+	"strings"
 )

 type ServerMode string
@ -63,7 +64,7 @@ func Parse() (*Config, bool) {
 	flag.Parse()

 	cfg := &Config{
-		ServerMode:     ServerMode(serverMode),
+		ServerMode:     ServerMode(strings.ToUpper(serverMode)),
 		BindAddr:       bindAddr,
 		Port:           port,
 		ListenAddr:     fmt.Sprintf("%s:%d", bindAddr, port),
@ -71,7 +72,7 @@ func Parse() (*Config, bool) {
 		PayloadUA:      payloadUA,
 		UARegex:        uaRegx,
 		PartialReplace: partial,
-		RewriteMode:    RewriteMode(rewriteMode),
+		RewriteMode:    RewriteMode(strings.ToUpper(rewriteMode)),
 		Rules:          rules,
 	}
 	if cfg.ServerMode == ServerModeRedirect || cfg.ServerMode == ServerModeTProxy {
--- a/src/internal/rewrite/rewriter.go
+++ b/src/internal/rewrite/rewriter.go
@ -35,13 +35,11 @@ type Rewriter struct {
 	Cache      *expirable.LRU[string, struct{}]
 }

-// RewriteDecision 重写决策结果
 type RewriteDecision struct {
 	Action      rule.Action
 	MatchedRule *rule.Rule
 }

-// ShouldRewrite 判断是否需要重写
 func (d *RewriteDecision) ShouldRewrite() bool {
 	return d.Action == rule.ActionReplace ||
 		d.Action == rule.ActionReplacePart ||
@ -57,7 +55,6 @@ func New(cfg *config.Config) (*Rewriter, error) {
 		return nil, err
 	}

-	// 创建规则引擎
 	var ruleEngine *rule.Engine
 	if cfg.RewriteMode == config.RewriteModeRules {
 		ruleEngine, err = rule.NewEngine(cfg.Rules)
@ -94,7 +91,6 @@ func (r *Rewriter) inWhitelist(ua string) bool {
 	return false
 }

-// GetRuleEngine 获取规则引擎
 func (r *Rewriter) GetRuleEngine() *rule.Engine {
 	return r.ruleEngine
 }
@ -119,7 +115,7 @@ func (r *Rewriter) EvaluateRewriteDecision(req *http.Request, srcAddr, destAddr
 		req.Header.Set("User-Agent", "")
 	}

-	// 「直接转发」模式：不进行任何重写
+	// DIRECT
 	if r.rewriteMode == config.RewriteModeDirect {
 		log.LogDebugWithAddr(srcAddr, destAddr, "Direct forward mode, skip rewriting")
 		statistics.AddPassThroughRecord(&statistics.PassThroughRecord{
@ -132,11 +128,11 @@ func (r *Rewriter) EvaluateRewriteDecision(req *http.Request, srcAddr, destAddr
 		}
 	}

-	// 「规则判定」模式：使用规则引擎（只匹配一次）
+	// RULES
 	if r.rewriteMode == config.RewriteModeRules && r.ruleEngine != nil {
 		matchedRule := r.ruleEngine.MatchWithRule(req, srcAddr, destAddr)

-		// 没有匹配到任何规则，默认直接转发
+		// no match rule, direct forward
 		if matchedRule == nil {
 			log.LogDebugWithAddr(srcAddr, destAddr, "No rule matched, direct forward")
 			statistics.AddPassThroughRecord(&statistics.PassThroughRecord{
@ -149,7 +145,7 @@ func (r *Rewriter) EvaluateRewriteDecision(req *http.Request, srcAddr, destAddr
 			}
 		}

-		// DROP 动作：丢弃请求
+		// DROP
 		if matchedRule.Action == rule.ActionDrop {
 			log.LogInfoWithAddr(srcAddr, destAddr, "Rule matched: DROP action, request will be dropped")
 			return &RewriteDecision{
@ -158,7 +154,7 @@ func (r *Rewriter) EvaluateRewriteDecision(req *http.Request, srcAddr, destAddr
 			}
 		}

-		// DIRECT 动作：直接转发
+		// DIRECT
 		if matchedRule.Action == rule.ActionDirect {
 			log.LogDebugWithAddr(srcAddr, destAddr, "Rule matched: DIRECT action, skip rewriting")
 			statistics.AddPassThroughRecord(&statistics.PassThroughRecord{
@ -172,14 +168,14 @@ func (r *Rewriter) EvaluateRewriteDecision(req *http.Request, srcAddr, destAddr
 			}
 		}

-		// REPLACE、REPLACE-PART、DELETE 动作：需要重写
+		// REPLACE、REPLACE-PART、DELETE, Rewrite
 		return &RewriteDecision{
 			Action:      matchedRule.Action,
 			MatchedRule: matchedRule,
 		}
 	}

-	// 「全局重写」模式：使用原有逻辑
+	// GLOBAL
 	var err error
 	matches := false
 	isWhitelist := r.inWhitelist(originalUA)
@ -226,11 +222,11 @@ func (r *Rewriter) Rewrite(req *http.Request, srcAddr string, destAddr string, d
 	action := decision.Action
 	var rewritedUA string

-	// 「规则判定」模式：根据规则动作决定如何重写
+	// RULES
 	if r.rewriteMode == config.RewriteModeRules && r.ruleEngine != nil {
 		rewritedUA = r.ruleEngine.ApplyAction(action, rewriteValue, originalUA, decision.MatchedRule)
 	} else {
-		// 「全局重写」模式：使用原有逻辑
+		// GLOBAL
 		rewritedUA = r.buildUserAgent(originalUA)
 	}

@ -328,14 +324,13 @@ func (r *Rewriter) Process(dst net.Conn, src net.Conn, destAddr string, srcAddr
 			return
 		}

-		// 获取重写决策（只匹配一次规则）
 		decision := r.EvaluateRewriteDecision(req, srcAddr, destAddr)
-		// 处理 DROP 动作
+
 		if decision.Action == rule.ActionDrop {
 			log.LogInfoWithAddr(srcAddr, destAddr, "Request dropped by rule")
 			continue
 		}
-		// 如果需要重写，执行重写操作
+
 		if decision.ShouldRewrite() {
 			req = r.Rewrite(req, srcAddr, destAddr, decision)
 		}
--- a/src/internal/rule/rule.go
+++ b/src/internal/rule/rule.go
@ -12,50 +12,43 @@ import (
 	"github.com/sirupsen/logrus"
 )

-// RuleType 规则类型
 type RuleType string

 const (
-	RuleTypeKeyword  RuleType = "KEYWORD"   // 关键字匹配
-	RuleTypeRegex    RuleType = "REGEX"     // 正则表达式匹配
-	RuleTypeIPCIDR   RuleType = "IP-CIDR"   // IP地址段匹配
-	RuleTypeSrcIP    RuleType = "SRC-IP"    // 源IP地址匹配
-	RuleTypeDestPort RuleType = "DEST-PORT" // 目标端口匹配
-	RuleTypeFinal    RuleType = "FINAL"     // 兜底规则
+	RuleTypeKeyword  RuleType = "KEYWORD"
+	RuleTypeRegex    RuleType = "REGEX"
+	RuleTypeIPCIDR   RuleType = "IP-CIDR"
+	RuleTypeSrcIP    RuleType = "SRC-IP"
+	RuleTypeDestPort RuleType = "DEST-PORT"
+	RuleTypeFinal    RuleType = "FINAL"
 )

-// Action 重写策略
 type Action string

 const (
-	ActionReplace     Action = "REPLACE"      // 替换整个 User-Agent
-	ActionReplacePart Action = "REPLACE-PART" // 部分替换
-	ActionDelete      Action = "DELETE"       // 删除 User-Agent
-	ActionDirect      Action = "DIRECT"       // 直接转发
-	ActionDrop        Action = "DROP"         // 丢弃请求
+	ActionReplace     Action = "REPLACE"
+	ActionReplacePart Action = "REPLACE-PART"
+	ActionDelete      Action = "DELETE"
+	ActionDirect      Action = "DIRECT"
+	ActionDrop        Action = "DROP"
 )

-// Rule 重写规则
 type Rule struct {
-	Enabled      bool     `json:"enabled"`       // 是否启用
-	Type         RuleType `json:"type"`          // 规则类型
-	Action       Action   `json:"action"`        // 重写策略
-	MatchValue   string   `json:"match_value"`   // 匹配值
-	RewriteValue string   `json:"rewrite_value"` // 重写值
-	Description  string   `json:"description"`   // 描述
+	Enabled      bool     `json:"enabled"`
+	Type         RuleType `json:"type"`
+	Action       Action   `json:"action"`
+	MatchValue   string   `json:"match_value"`
+	RewriteValue string   `json:"rewrite_value"`
+	Description  string   `json:"description"`

-	// 编译后的正则表达式（仅用于 REGEX 类型）
 	regex *regexp2.Regexp
-	// 解析后的 IP 网络（仅用于 IP-CIDR 和 SRC-IP 类型）
 	ipNet *net.IPNet
 }

-// Engine 规则引擎
 type Engine struct {
 	rules []*Rule
 }

-// NewEngine 创建规则引擎
 func NewEngine(rulesJSON string) (*Engine, error) {
 	if rulesJSON == "" {
 		return &Engine{rules: []*Rule{}}, nil
@ -66,7 +59,6 @@ func NewEngine(rulesJSON string) (*Engine, error) {
 		return nil, fmt.Errorf("failed to parse rules JSON: %w", err)
 	}

-	// 编译正则表达式和解析 IP 网络
 	for _, rule := range rules {
 		if !rule.Enabled {
 			continue
@ -103,7 +95,6 @@ func NewEngine(rulesJSON string) (*Engine, error) {
 	return &Engine{rules: rules}, nil
 }

-// MatchWithRule 匹配规则并返回匹配的规则
 func (e *Engine) MatchWithRule(req *http.Request, srcAddr, destAddr string) *Rule {
 	for _, rule := range e.rules {
 		if !rule.Enabled {
@ -139,13 +130,11 @@ func (e *Engine) MatchWithRule(req *http.Request, srcAddr, destAddr string) *Rul
 	return nil
 }

-// matchKeyword 关键字匹配
 func (e *Engine) matchKeyword(req *http.Request, rule *Rule) bool {
 	ua := req.Header.Get("User-Agent")
 	return strings.Contains(strings.ToLower(ua), strings.ToLower(rule.MatchValue))
 }

-// matchRegex 正则表达式匹配
 func (e *Engine) matchRegex(req *http.Request, rule *Rule) (bool, error) {
 	if rule.regex == nil {
 		return false, nil
@ -154,7 +143,6 @@ func (e *Engine) matchRegex(req *http.Request, rule *Rule) (bool, error) {
 	return rule.regex.MatchString(ua)
 }

-// matchIPCIDR 目标IP地址段匹配
 func (e *Engine) matchIPCIDR(destAddr string, rule *Rule) bool {
 	if rule.ipNet == nil {
 		return false
@ -170,7 +158,6 @@ func (e *Engine) matchIPCIDR(destAddr string, rule *Rule) bool {
 	return rule.ipNet.Contains(ip)
 }

-// matchSrcIP 源IP地址匹配
 func (e *Engine) matchSrcIP(srcAddr string, rule *Rule) bool {
 	if rule.ipNet == nil {
 		return false
@ -186,7 +173,6 @@ func (e *Engine) matchSrcIP(srcAddr string, rule *Rule) bool {
 	return rule.ipNet.Contains(ip)
 }

-// matchDestPort 目标端口匹配
 func (e *Engine) matchDestPort(destAddr string, rule *Rule) bool {
 	_, portStr, err := net.SplitHostPort(destAddr)
 	if err != nil {
@ -203,7 +189,6 @@ func (e *Engine) matchDestPort(destAddr string, rule *Rule) bool {
 	return port == matchPort
 }

-// ApplyAction 应用规则动作
 func (e *Engine) ApplyAction(action Action, rewriteValue string, originalUA string, rule *Rule) string {
 	switch action {
 	case ActionReplace:
@ -217,20 +202,19 @@ func (e *Engine) ApplyAction(action Action, rewriteValue string, originalUA stri
 			}
 			return newUA
 		}
-		// 如果不是正则规则，使用简单字符串替换
+		// if not regex, do simple string replacement
 		return strings.ReplaceAll(originalUA, rule.MatchValue, rewriteValue)
 	case ActionDelete:
 		return ""
 	case ActionDirect:
 		return originalUA
 	case ActionDrop:
-		return originalUA // DROP 由上层处理
+		return originalUA // DROP action handled elsewhere
 	default:
 		return originalUA
 	}
 }

-// HasRules 是否有启用的规则
 func (e *Engine) HasRules() bool {
 	for _, rule := range e.rules {
 		if rule.Enabled {
@ -240,7 +224,6 @@ func (e *Engine) HasRules() bool {
 	return false
 }

-// GetRules 获取所有规则
 func (e *Engine) GetRules() []*Rule {
 	return e.rules
 }
--- a/src/internal/server/http/http.go
+++ b/src/internal/server/http/http.go
@ -6,6 +6,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"time"

 	"github.com/sirupsen/logrus"
 	"github.com/sunbk201/ua3f/internal/config"
@ -13,6 +14,8 @@ import (
 	"github.com/sunbk201/ua3f/internal/rewrite"
 	"github.com/sunbk201/ua3f/internal/rule"
 	"github.com/sunbk201/ua3f/internal/server/utils"
+	"github.com/sunbk201/ua3f/internal/sniff"
+	"github.com/sunbk201/ua3f/internal/statistics"
 )

 type Server struct {
@ -28,6 +31,8 @@ func New(cfg *config.Config, rw *rewrite.Rewriter) *Server {
 }

 func (s *Server) Start() (err error) {
+	go statistics.StartRecorder()
+
 	server := &http.Server{
 		Addr: s.cfg.ListenAddr,
 		Handler: http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
@ -50,6 +55,12 @@ func (s *Server) handleHTTP(w http.ResponseWriter, req *http.Request) {
 		destPort = "80"
 	}
 	destAddr := fmt.Sprintf("%s:%s", req.URL.Hostname(), destPort)
+	statistics.AddConnection(&statistics.ConnectionRecord{
+		Protocol:  sniff.HTTP,
+		SrcAddr:   req.RemoteAddr,
+		DestAddr:  destAddr,
+		StartTime: time.Now(),
+	})

 	logrus.Infof("HTTP request for %s", destAddr)

@ -79,6 +90,7 @@ func (s *Server) handleHTTP(w http.ResponseWriter, req *http.Request) {
 	}
 	w.WriteHeader(resp.StatusCode)
 	io.Copy(w, resp.Body)
+	statistics.RemoveConnection(req.RemoteAddr, destAddr)
 }

 func (s *Server) handleTunneling(w http.ResponseWriter, req *http.Request) {
@ -104,24 +116,16 @@ func (s *Server) handleTunneling(w http.ResponseWriter, req *http.Request) {
 }

 func (s *Server) rewriteAndForward(target net.Conn, req *http.Request, dstAddr, srcAddr string) (err error) {
-	rw := s.rw
-
-	// 获取重写决策（只匹配一次规则）
-	decision := rw.EvaluateRewriteDecision(req, srcAddr, dstAddr)
-
-	// Handle DROP action
+	decision := s.rw.EvaluateRewriteDecision(req, srcAddr, dstAddr)
 	if decision.Action == rule.ActionDrop {
 		log.LogInfoWithAddr(srcAddr, dstAddr, "Request dropped by rule")
 		return fmt.Errorf("request dropped by rule")
 	}
-
-	// 如果需要重写，执行重写操作
 	if decision.ShouldRewrite() {
-		req = rw.Rewrite(req, srcAddr, dstAddr, decision)
+		req = s.rw.Rewrite(req, srcAddr, dstAddr, decision)
 	}
-
-	if err = rw.Forward(target, req); err != nil {
-		err = fmt.Errorf("r.forward: %w", err)
+	if err = s.rw.Forward(target, req); err != nil {
+		err = fmt.Errorf("s.rw.Forward: %w", err)
 		return
 	}
 	return nil
@ -137,7 +141,7 @@ func (s *Server) ForwardTCP(client, target net.Conn, destAddr string) {
 	// Server -> Client (raw)
 	go utils.CopyHalf(client, target)

-	if s.cfg.RewriteMode == "direct" {
+	if s.cfg.RewriteMode == config.RewriteModeDirect {
 		// Client -> Server (raw)
 		go utils.CopyHalf(target, client)
 		return