him411/UA3F
1
0
Fork 0
mirror of https://github.com/SunBK201/UA3F.git synced 2025-12-16 08:44:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

style: code tidy

Browse Source
This commit is contained in:
SunBK201 2025-11-01 00:30:57 +08:00
parent 872e519c3d
commit 754997a33d
3 changed files with 22 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
src/internal/rewrite/rewriter.go
View File

@ -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,

4
src/internal/rewrite/rewriter_test.go
View File

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

13
src/internal/server/utils/tcp.go
View File

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