From f7085fd2df5d6eb76c27a981118a4dcb147cbb90 Mon Sep 17 00:00:00 2001 From: SunBK201 Date: Sat, 1 Nov 2025 00:30:57 +0800 Subject: [PATCH] style: code tidy --- src/internal/rewrite/rewriter.go | 26 +++++++++----------------- src/internal/rewrite/rewriter_test.go | 4 ++-- src/internal/server/utils/tcp.go | 13 +++++++++++-- 3 files changed, 22 insertions(+), 21 deletions(-) diff --git a/src/internal/rewrite/rewriter.go b/src/internal/rewrite/rewriter.go index 3041358..f8ed9be 100644 --- a/src/internal/rewrite/rewriter.go +++ b/src/internal/rewrite/rewriter.go @@ -53,7 +53,7 @@ type Rewriter struct { enablePartialReplace bool uaRegex *regexp2.Regexp - cache *expirable.LRU[string, string] + Cache *expirable.LRU[string, string] whitelist map[string]struct{} } @@ -78,7 +78,7 @@ func New(cfg *config.Config) (*Rewriter, error) { pattern: cfg.UAPattern, enablePartialReplace: cfg.EnablePartialReplace, uaRegex: uaRegex, - cache: cache, + Cache: cache, whitelist: whitelist, }, nil } @@ -87,17 +87,9 @@ func New(cfg *config.Config) (*Rewriter, error) { // - If target in LRU cache: pass-through (raw). // - Else if HTTP: rewrite UA (unless whitelisted or pattern not matched). // - Else: mark target in LRU and pass-through. -func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string) (err error) { - srcAddr := src.RemoteAddr().String() - - // Fast path: known pass-through - if r.cache.Contains(destAddr) { - log.LogDebugWithAddr(src.RemoteAddr().String(), destAddr, "LRU Relay Cache Hit, pass-through") - io.Copy(dst, src) - return nil - } - +func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string, srcAddr string) (err error) { reader := bufio.NewReader(src) + defer func() { if err != nil { log.LogDebugWithAddr(srcAddr, destAddr, fmt.Sprintf("ProxyHTTPOrRaw: %s", err.Error())) @@ -106,7 +98,7 @@ func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string) ( }() if strings.HasSuffix(destAddr, "443") && isTLSClientHello(reader) { - r.cache.Add(destAddr, destAddr) + r.Cache.Add(destAddr, destAddr) log.LogDebugWithAddr(srcAddr, destAddr, "TLS ClientHello detected") return } @@ -116,7 +108,7 @@ func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string) ( return } if !isHTTP { - r.cache.Add(destAddr, destAddr) + r.Cache.Add(destAddr, destAddr) log.LogDebugWithAddr(srcAddr, destAddr, "Not HTTP, added to LRU Relay Cache") return } @@ -135,7 +127,7 @@ func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string) ( if isWebSocket(h2) { log.LogDebugWithAddr(srcAddr, destAddr, "WebSocket detected, pass-through") } else { - r.cache.Add(destAddr, destAddr) + r.Cache.Add(destAddr, destAddr) log.LogDebugWithAddr(srcAddr, destAddr, "Not HTTP, added to LRU Relay Cache") } return @@ -150,7 +142,7 @@ func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string) ( // No UA header: pass-through after writing this first request if originalUA == "" { - r.cache.Add(destAddr, destAddr) + r.Cache.Add(destAddr, destAddr) log.LogDebugWithAddr(srcAddr, destAddr, "Not found User-Agent, Add LRU Relay Cache") if err = req.Write(dst); err != nil { err = fmt.Errorf("req.Write: %w", err) @@ -175,7 +167,7 @@ func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string) ( } if isWhitelist { log.LogDebugWithAddr(srcAddr, destAddr, fmt.Sprintf("Hit User-Agent Whitelist: %s", originalUA)) - r.cache.Add(destAddr, destAddr) + r.Cache.Add(destAddr, destAddr) } statistics.AddPassThroughRecord(&statistics.PassThroughRecord{ Host: destAddr, diff --git a/src/internal/rewrite/rewriter_test.go b/src/internal/rewrite/rewriter_test.go index a1d5a69..9d19d1e 100644 --- a/src/internal/rewrite/rewriter_test.go +++ b/src/internal/rewrite/rewriter_test.go @@ -48,7 +48,7 @@ func TestNewRewriter(t *testing.T) { assert.Equal(t, cfg.UAPattern, rewriter.pattern) assert.Equal(t, cfg.EnablePartialReplace, rewriter.enablePartialReplace) assert.NotNil(t, rewriter.uaRegex) - assert.NotNil(t, rewriter.cache) + assert.NotNil(t, rewriter.Cache) } func TestIsHTTP(t *testing.T) { @@ -83,7 +83,7 @@ func TestProxyHTTPOrRaw_HTTPRewrite(t *testing.T) { dstBuf := &bytes.Buffer{} dst := &mockConn{Reader: nil, Writer: dstBuf} - r.ProxyHTTPOrRaw(dst, src, "example.com:80") + r.ProxyHTTPOrRaw(dst, src, "example.com:80", "srcAddr") out := dstBuf.String() assert.Contains(t, out, "User-Agent: MockUA/1.0") diff --git a/src/internal/server/utils/tcp.go b/src/internal/server/utils/tcp.go index 425aedc..0ad90fd 100644 --- a/src/internal/server/utils/tcp.go +++ b/src/internal/server/utils/tcp.go @@ -6,6 +6,7 @@ import ( "net" "github.com/sirupsen/logrus" + "github.com/sunbk201/ua3f/internal/log" "github.com/sunbk201/ua3f/internal/rewrite" ) @@ -38,7 +39,7 @@ func CopyHalf(dst, src net.Conn) { } // ProxyHalf runs the rewriter proxy on src->dst and then half-closes both sides. -func ProxyHalf(dst, src net.Conn, rw *rewrite.Rewriter, destAddrPort string) { +func ProxyHalf(dst, src net.Conn, rw *rewrite.Rewriter, destAddr string) { defer func() { if tc, ok := dst.(*net.TCPConn); ok { _ = tc.CloseWrite() @@ -51,7 +52,15 @@ func ProxyHalf(dst, src net.Conn, rw *rewrite.Rewriter, destAddrPort string) { _ = src.Close() } }() - _ = rw.ProxyHTTPOrRaw(dst, src, destAddrPort) + + // Fast path: known pass-through + srcAddr := src.RemoteAddr().String() + if rw.Cache.Contains(destAddr) { + log.LogDebugWithAddr(srcAddr, destAddr, "LRU Relay Cache Hit, pass-through") + io.Copy(dst, src) + return + } + _ = rw.ProxyHTTPOrRaw(dst, src, destAddr, srcAddr) } func GetConnFD(conn net.Conn) (fd int, err error) {