style: code tidy

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) {
+func (r *Rewriter) ProxyHTTPOrRaw(dst net.Conn, src net.Conn, destAddr string, srcAddr 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
-	}
-
 	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,

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")

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) {