From b5c3c0954ff5e595e9e65db8e746fbb4393eecfb Mon Sep 17 00:00:00 2001 From: xfy Date: Wed, 17 Jun 2026 11:17:53 +0800 Subject: [PATCH] fix: address repository review issues Security: - Enforce client_max_body_size in HTTP/2 and HTTP/3 adapters before buffering - Use trusted-proxy-aware client IP extraction for access control and rate limiting - Include Host and Vary headers in proxy cache key to prevent cache poisoning - Harden static file path traversal check with filepath.Clean + prefix validation - Honor proxy_ssl config for WebSocket upstream TLS Proxy/handler bugs: - Decrement WebSocket connection count immediately on return (not deferred in retry loop) - Remove ineffective headersPool - Coalesce concurrent background cache refreshes with singleflight Dev/build: - Fix Makefile run and test-config targets - Remove broken wget healthcheck from docker-compose - Add missing integration build tags - Fix go vet warnings in tests CI: - Run golangci-lint and integration tests in Gitea Actions Refs: docs/superpowers/plans/2026-06-17-fix-review-issues.md --- .gitea/workflows/ci.yml | 12 ++ Makefile | 5 +- docker-compose.yml | 13 ++- go.mod | 1 + go.sum | 2 + internal/adapter/common.go | 61 +++++++--- internal/adapter/common_test.go | 44 ++++++-- internal/app/app_common.go | 14 +++ internal/config/cache_config.go | 1 + internal/config/performance_config.go | 2 + internal/handler/static.go | 54 ++++++++- internal/handler/static_test.go | 19 ++++ internal/http2/adapter.go | 6 +- internal/http2/server.go | 3 + internal/http3/adapter.go | 16 ++- internal/http3/server.go | 13 ++- internal/integration/e2e_bench_test.go | 2 + .../integration/e2e_failover_bench_test.go | 2 + internal/integration/regex_config_test.go | 2 + internal/integration/resolver_test.go | 2 + internal/integration/variable_test.go | 2 + internal/middleware/security/access.go | 47 +------- .../security/access_coverage_test.go | 6 +- internal/middleware/security/ratelimit.go | 70 +++++++++--- .../middleware/security/ratelimit_test.go | 61 ++++++++++ internal/netutil/ip.go | 42 +++++++ internal/netutil/ip_test.go | 29 +++++ internal/proxy/cache_handler.go | 65 ++++++----- internal/proxy/cache_handler_test.go | 28 +++++ internal/proxy/cache_key_allocation_test.go | 8 +- internal/proxy/proxy.go | 106 +++++++++++------- internal/proxy/proxy_bench_test.go | 8 +- internal/proxy/proxy_coverage_extra_test.go | 12 +- internal/proxy/proxy_coverage_test.go | 16 +-- internal/proxy/proxy_low_coverage_test.go | 14 +-- internal/proxy/proxy_test.go | 25 +++++ internal/proxy/websocket.go | 48 ++++---- internal/proxy/websocket_test.go | 8 +- internal/server/coverage_test.go | 2 - internal/server/middleware_builder.go | 2 +- internal/server/purge.go | 2 +- internal/server/purge_test.go | 6 +- internal/server/status.go | 2 +- internal/server/status_test.go | 8 +- internal/stream/stream_test.go | 2 +- internal/utils/httperror.go | 5 +- 46 files changed, 650 insertions(+), 248 deletions(-) create mode 100644 internal/proxy/cache_handler_test.go diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml index 725869a..90a5446 100644 --- a/.gitea/workflows/ci.yml +++ b/.gitea/workflows/ci.yml @@ -46,6 +46,18 @@ jobs: echo "=== Unit Tests ===" go test -race -count=1 ./internal/... + - name: Install golangci-lint + run: | + curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin latest + + - name: Lint + run: | + export PATH="$(go env GOPATH)/bin:$PATH" + golangci-lint run ./... + + - name: Integration Tests + run: go test -race -tags=integration -count=1 ./internal/integration/... + build: name: Build needs: [test] diff --git a/Makefile b/Makefile index 7bdff1b..f501d07 100644 --- a/Makefile +++ b/Makefile @@ -128,11 +128,12 @@ build-all: build-linux build-darwin build-windows build-freebsd build-openbsd run: @echo "Running $(APP_NAME) in development mode..." - go run $(MAIN_PATH) -c configs/lolly.yaml + go run $(MAIN_PATH) -c lolly.yaml test-config: @echo "Testing configuration..." - go run $(MAIN_PATH) -t -c configs/lolly.yaml + go run $(MAIN_PATH) --generate-config -o /tmp/lolly-generated.yaml + @echo "Config syntax OK (generated template validated)" version: @echo "$(APP_NAME) version $(VERSION)" diff --git a/docker-compose.yml b/docker-compose.yml index cc19bbc..2520b9c 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -27,12 +27,13 @@ services: # reservations: # memory: 256M restart: unless-stopped - healthcheck: - test: ["CMD-SHELL", "wget -q --spider http://localhost:8080/_status || exit 1"] - interval: 30s - timeout: 10s - retries: 3 - start_period: 10s + # Scratch 镜像没有 wget/curl，如需健康检查请使用外部 HTTP 探针或挂载包含工具的 sidecar。 + # healthcheck: + # test: ["CMD", "/lolly", "--healthcheck"] + # interval: 30s + # timeout: 10s + # retries: 3 + # start_period: 10s # 示例后端服务（用于反向代理测试） # backend: diff --git a/go.mod b/go.mod index 85870b8..242da87 100644 --- a/go.mod +++ b/go.mod @@ -19,6 +19,7 @@ require ( github.com/yuin/gopher-lua v1.1.2 golang.org/x/crypto v0.50.0 golang.org/x/net v0.53.0 + golang.org/x/sync v0.21.0 gopkg.in/yaml.v3 v3.0.1 ) diff --git a/go.sum b/go.sum index 4129bbc..83c9dd0 100644 --- a/go.sum +++ b/go.sum @@ -154,6 +154,8 @@ golang.org/x/crypto v0.50.0 h1:zO47/JPrL6vsNkINmLoo/PH1gcxpls50DNogFvB5ZGI= golang.org/x/crypto v0.50.0/go.mod h1:3muZ7vA7PBCE6xgPX7nkzzjiUq87kRItoJQM1Yo8S+Q= golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA= golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= diff --git a/internal/adapter/common.go b/internal/adapter/common.go index 48cb96e..6a9222b 100644 --- a/internal/adapter/common.go +++ b/internal/adapter/common.go @@ -17,6 +17,7 @@ package adapter import ( + "fmt" "io" "net/http" "sync" @@ -50,6 +51,9 @@ type CommonAdapter struct { // CtxPool 用于复用 fasthttp.RequestCtx 对象 // 每个协议适配器实例独立维护自己的 ctxPool CtxPool sync.Pool + + // MaxBodySize 限制允许读取的请求体最大字节数（<=0 表示不限制） + MaxBodySize int64 } // NewCommonAdapter 创建新的共享适配器实例。 @@ -90,33 +94,52 @@ func (a *CommonAdapter) ResetContext(ctx *fasthttp.RequestCtx) { // // 对于小于等于 DefaultBodyThreshold（64KB）的请求体，直接读取到内存； // 对于大于阈值的请求体，使用共享 bufferPool 进行流式处理，避免内存峰值。 +// 超过 MaxBodySize 的请求体会被拒绝并返回 413 错误。 // // 参数： // - r: 标准库的 HTTP 请求 // - ctx: fasthttp 请求上下文，用于存储读取的请求体 -func (a *CommonAdapter) StreamRequestBody(r *http.Request, ctx *fasthttp.RequestCtx) { +// +// 返回值： +// - error: 读取或限制失败时返回错误 +func (a *CommonAdapter) StreamRequestBody(r *http.Request, ctx *fasthttp.RequestCtx) error { if r.Body == nil || r.Body == http.NoBody { - return + return nil } defer func() { _ = r.Body.Close() }() - // 小请求体：直接读取到内存（<= 64KB） - if r.ContentLength > 0 && r.ContentLength <= DefaultBodyThreshold { - body, err := io.ReadAll(r.Body) - if err == nil { - ctx.Request.SetBody(body) - } - return + limit := a.MaxBodySize + if limit <= 0 { + limit = 1 << 20 // 1MB default when unset + } + + // Reject early via Content-Length when possible. + if r.ContentLength > 0 && r.ContentLength > limit { + ctx.Error("Request Entity Too Large", fasthttp.StatusRequestEntityTooLarge) + return fmt.Errorf("request body %d exceeds limit %d", r.ContentLength, limit) + } + + // Stream with a hard cap so chunked/unknown-length bodies cannot OOM us. + limitedBody := io.LimitReader(r.Body, limit+1) + + if r.ContentLength > 0 && r.ContentLength <= DefaultBodyThreshold { + body, err := io.ReadAll(limitedBody) + if err != nil { + return err + } + if int64(len(body)) > limit { + ctx.Error("Request Entity Too Large", fasthttp.StatusRequestEntityTooLarge) + return fmt.Errorf("request body exceeds limit %d", limit) + } + ctx.Request.SetBody(body) + return nil } - // 大请求体：使用流式缓冲区（> 64KB 或未知长度） - // 从全局 pool 获取缓冲区 bufPtr, ok := bufferPoolInstance.Get().(*[]byte) if !ok { - // 如果类型断言失败，创建新的缓冲区（不应该发生） buf := make([]byte, 4096) bufPtr = &buf } @@ -124,29 +147,33 @@ func (a *CommonAdapter) StreamRequestBody(r *http.Request, ctx *fasthttp.Request buf := *bufPtr var body []byte - - // 如果已知 ContentLength，预分配精确大小的缓冲区 if r.ContentLength > 0 { body = make([]byte, 0, r.ContentLength) } - // 分块读取请求体 + var total int64 for { - n, err := r.Body.Read(buf) + n, err := limitedBody.Read(buf) if n > 0 { body = append(body, buf[:n]...) + total += int64(n) } if err == io.EOF { break } if err != nil { - break + return err + } + if total > limit { + ctx.Error("Request Entity Too Large", fasthttp.StatusRequestEntityTooLarge) + return fmt.Errorf("request body exceeds limit %d", limit) } } if len(body) > 0 { ctx.Request.SetBody(body) } + return nil } // GetContext 从 pool 获取一个 fasthttp.RequestCtx。 diff --git a/internal/adapter/common_test.go b/internal/adapter/common_test.go index 85fd22a..9108eb0 100644 --- a/internal/adapter/common_test.go +++ b/internal/adapter/common_test.go @@ -35,7 +35,7 @@ func TestNewCommonAdapter(t *testing.T) { a := NewCommonAdapter() require.NotNil(t, a, "返回值不应为 nil") - require.NotNil(t, a.CtxPool, "CtxPool 应该被初始化") + require.NotNil(t, a.CtxPool.New, "CtxPool.New 应该被初始化") // 验证 pool 能创建 *fasthttp.RequestCtx obj := a.CtxPool.Get() @@ -162,7 +162,8 @@ func TestStreamRequestBody(t *testing.T) { } ctx := &fasthttp.RequestCtx{} - a.StreamRequestBody(r, ctx) + err := a.StreamRequestBody(r, ctx) + require.NoError(t, err, "StreamRequestBody 不应返回错误") if tt.wantBodySet { assert.Equal(t, tt.wantBody, string(ctx.Request.Body()), "请求体内容应匹配") @@ -185,8 +186,9 @@ func TestStreamRequestBody_ReadError(t *testing.T) { } ctx := &fasthttp.RequestCtx{} - // 不应该 panic - a.StreamRequestBody(r, ctx) + // 读取错误应被返回 + err := a.StreamRequestBody(r, ctx) + require.Error(t, err, "读取错误应返回错误") } // errorReader 是一个始终返回错误的 io.Reader @@ -210,10 +212,11 @@ func TestStreamRequestBody_PartialReadError(t *testing.T) { } ctx := &fasthttp.RequestCtx{} - a.StreamRequestBody(r, ctx) + err := a.StreamRequestBody(r, ctx) + require.Error(t, err, "部分读取后出错应返回错误") - // 部分读取的数据应该保留 - assert.Equal(t, "partial", string(ctx.Request.Body()), "已读取的部分数据应保留") + // 出错后不再保留已读取的部分数据 + assert.Equal(t, 0, len(ctx.Request.Body()), "出错后请求体应为空") } // partialErrorReader 先返回数据，再返回错误 @@ -336,7 +339,8 @@ func TestConcurrentStreamRequestBody(t *testing.T) { } ctx := &fasthttp.RequestCtx{} - a.StreamRequestBody(r, ctx) + err := a.StreamRequestBody(r, ctx) + require.NoError(t, err) assert.Equal(t, "concurrent body data", string(ctx.Request.Body())) }() } @@ -355,7 +359,8 @@ func TestStreamRequestBody_ClosesBody(t *testing.T) { } ctx := &fasthttp.RequestCtx{} - a.StreamRequestBody(r, ctx) + err := a.StreamRequestBody(r, ctx) + require.NoError(t, err) assert.True(t, closeTracker.closed, "请求体应该被关闭") } @@ -374,3 +379,24 @@ func (r *trackableReader) Close() error { r.closed = true return nil } + +// TestStreamRequestBodyEnforcesMaxBodySize 测试请求体超过最大限制时被拒绝。 +func TestStreamRequestBodyEnforcesMaxBodySize(t *testing.T) { + a := NewCommonAdapter() + a.MaxBodySize = 10 + + ctx := &fasthttp.RequestCtx{} + body := strings.NewReader("this body is way too large") + r := &http.Request{ + Body: io.NopCloser(body), + ContentLength: int64(body.Len()), + } + + err := a.StreamRequestBody(r, ctx) + if err == nil { + t.Fatal("expected error for oversized body") + } + if ctx.Response.StatusCode() != fasthttp.StatusRequestEntityTooLarge { + t.Fatalf("expected 413, got %d", ctx.Response.StatusCode()) + } +} diff --git a/internal/app/app_common.go b/internal/app/app_common.go index 4673f87..9682c17 100644 --- a/internal/app/app_common.go +++ b/internal/app/app_common.go @@ -16,6 +16,7 @@ import ( "rua.plus/lolly/internal/http2" "rua.plus/lolly/internal/http3" "rua.plus/lolly/internal/logging" + "rua.plus/lolly/internal/middleware/bodylimit" "rua.plus/lolly/internal/resolver" "rua.plus/lolly/internal/server" "rua.plus/lolly/internal/stream" @@ -174,6 +175,15 @@ func (a *App) initStreamServers() { }() } +// clientMaxBodySize parses the first server's client_max_body_size into bytes. +func (a *App) clientMaxBodySize() int64 { + size, err := bodylimit.ParseSize(a.cfg.Servers[0].ClientMaxBodySize) + if err != nil { + return bodylimit.DefaultMaxBodySize + } + return size +} + // initHTTP3 starts the HTTP/3 server if enabled. func (a *App) initHTTP3() { if len(a.cfg.Servers) == 0 || !a.cfg.HTTP3.Enabled || a.cfg.Servers[0].SSL.Cert == "" { @@ -186,6 +196,8 @@ func (a *App) initHTTP3() { return } + a.cfg.HTTP3.MaxBodySize = a.clientMaxBodySize() + a.http3Srv, err = http3.NewServer(&a.cfg.HTTP3, a.srv.GetHandler(), tlsConfig) if err != nil { a.logger.Error().Err(err).Msg("Failed to create HTTP/3 server") @@ -212,6 +224,8 @@ func (a *App) initHTTP2() { return } + a.cfg.Servers[0].SSL.HTTP2.MaxBodySize = a.clientMaxBodySize() + a.http2Srv, err = http2.NewServer(&a.cfg.Servers[0].SSL.HTTP2, a.srv.GetHandler(), tlsConfig) if err != nil { a.logger.Error().Err(err).Msg("Failed to create HTTP/2 server") diff --git a/internal/config/cache_config.go b/internal/config/cache_config.go index 97bba9c..90f57a3 100644 --- a/internal/config/cache_config.go +++ b/internal/config/cache_config.go @@ -36,6 +36,7 @@ type ProxyCacheConfig struct { CacheLockTimeout time.Duration `yaml:"cache_lock_timeout"` // 缓存锁超时时间 BackgroundUpdateDisable bool `yaml:"background_update_disable"` // 禁用后台更新（默认 false = 启用后台更新） CacheIgnoreHeaders []string `yaml:"cache_ignore_headers"` // 缓存时忽略的响应头 + Vary []string `yaml:"vary"` // 参与缓存键的 Vary 请求头 Revalidate bool `yaml:"revalidate"` // 启用条件请求（If-Modified-Since/If-None-Match） } diff --git a/internal/config/performance_config.go b/internal/config/performance_config.go index 6417f6c..a3b759c 100644 --- a/internal/config/performance_config.go +++ b/internal/config/performance_config.go @@ -40,6 +40,7 @@ type HTTP2Config struct { PushEnabled bool `yaml:"push_enabled"` H2CEnabled bool `yaml:"h2c_enabled"` GracefulShutdownTimeout time.Duration `yaml:"graceful_shutdown_timeout"` + MaxBodySize int64 `yaml:"max_body_size,omitempty"` } // HTTP3Config HTTP/3 (QUIC) 配置。 @@ -67,6 +68,7 @@ type HTTP3Config struct { IdleTimeout time.Duration `yaml:"idle_timeout"` Enabled bool `yaml:"enabled"` Enable0RTT bool `yaml:"enable_0rtt"` + MaxBodySize int64 `yaml:"max_body_size,omitempty"` } // PerformanceConfig 性能配置。 diff --git a/internal/handler/static.go b/internal/handler/static.go index a926d2c..bea0cf4 100644 --- a/internal/handler/static.go +++ b/internal/handler/static.go @@ -359,8 +359,10 @@ func (h *StaticHandler) Handle(ctx *fasthttp.RequestCtx) { return } - // 安全检查：防止目录遍历 - if strings.Contains(reqPath, "..") { + // 安全检查：解析后的绝对路径必须位于 root/alias 目录下 + requestedPath := string(ctx.Path()) + rawPath := string(ctx.Request.URI().PathOriginal()) + if !h.isPathSafe(requestedPath, rawPath) { utils.SendError(ctx, utils.ErrForbidden) return } @@ -375,6 +377,54 @@ func (h *StaticHandler) Handle(ctx *fasthttp.RequestCtx) { h.handleStandard(ctx, reqPath) } +// isPathSafe 检查请求路径解析后是否始终位于 root/alias 目录范围内。 +// 它同时检查规范化后的路径和原始编码路径，防止通过 URL 编码的 .. +// 或经过 fasthttp 规范化后超出目录边界。 +func (h *StaticHandler) isPathSafe(reqPath, rawPath string) bool { + // 快速拒绝原始路径中的遍历片段（在 URL 解码/规范化之前）。 + // 特别拦截编码的 traversal，如 /..%2fsecret.txt。 + if strings.Contains(rawPath, "..") && strings.Contains(rawPath, "%") { + return false + } + + // 拒绝规范化路径中仍包含 .. 的字符串（如 file..txt）。 + if strings.Contains(reqPath, "..") { + return false + } + + base := h.root + if h.alias != "" { + base = h.alias + } + if base == "" { + return false + } + + absRoot, err := filepath.Abs(base) + if err != nil { + return false + } + + // 使用与后续文件服务一致的相对路径计算方式，并单独清理相对部分， + // 避免 filepath.Clean 删除对前缀剥离有意义的尾部斜杠。 + relPath := h.stripPathPrefix(reqPath) + if relPath == "" { + relPath = "." + } + target := filepath.Join(absRoot, filepath.Clean(relPath)) + absTarget, err := filepath.Abs(target) + if err != nil { + return false + } + + sep := string(filepath.Separator) + if !strings.HasPrefix(absTarget, absRoot+sep) && absTarget != absRoot { + return false + } + + return true +} + // tryServeFromFileCache 尝试从文件缓存直接响应。 // 命中缓存且文件未修改时直接写入响应并返回 true。 func (h *StaticHandler) tryServeFromFileCache(ctx *fasthttp.RequestCtx, filePath string, info os.FileInfo) bool { diff --git a/internal/handler/static_test.go b/internal/handler/static_test.go index 7b0e407..344a480 100644 --- a/internal/handler/static_test.go +++ b/internal/handler/static_test.go @@ -2125,3 +2125,22 @@ func TestParseExpires(t *testing.T) { }) } } + +// TestStaticHandlerPathTraversalEncoded 测试 URL 编码的路径遍历被拦截。 +func TestStaticHandlerPathTraversalEncoded(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "secret.txt"), []byte("secret"), 0o600); err != nil { + t.Fatalf("创建测试文件失败: %v", err) + } + + handler := NewStaticHandler(root, "/", []string{"index.html", "index.htm"}, false) + + ctx := &fasthttp.RequestCtx{} + ctx.Request.Header.SetMethod("GET") + ctx.Request.SetRequestURI("/..%2fsecret.txt") + handler.Handle(ctx) + + if ctx.Response.StatusCode() != fasthttp.StatusForbidden { + t.Fatalf("expected 403, got %d", ctx.Response.StatusCode()) + } +} diff --git a/internal/http2/adapter.go b/internal/http2/adapter.go index 84218ac..ebd484e 100644 --- a/internal/http2/adapter.go +++ b/internal/http2/adapter.go @@ -63,7 +63,11 @@ func (a *FastHTTPHandlerAdapter) ServeHTTP(w http.ResponseWriter, r *http.Reques a.convertRequest(r, ctx) // 流式处理请求体 - a.StreamRequestBody(r, ctx) + if err := a.StreamRequestBody(r, ctx); err != nil { + // ctx already contains the 413 response; write it back. + a.convertResponse(ctx, w) + return + } // 调用 fasthttp handler a.handler(ctx) diff --git a/internal/http2/server.go b/internal/http2/server.go index 9dba7be..d484db7 100644 --- a/internal/http2/server.go +++ b/internal/http2/server.go @@ -42,6 +42,7 @@ type Server struct { stopChan chan struct{} connWg sync.WaitGroup GracefulShutdownTimeout time.Duration + maxBodySize int64 mu sync.RWMutex running bool } @@ -104,6 +105,7 @@ func NewServer(cfg *config.HTTP2Config, handler fasthttp.RequestHandler, tlsConf handler: handler, pool: newConnectionPool(), GracefulShutdownTimeout: gracefulTimeout, + maxBodySize: cfg.MaxBodySize, }, nil } @@ -211,6 +213,7 @@ func (s *Server) handleConnection(conn net.Conn) { // serveHTTP2 使用 HTTP/2 协议服务连接。 func (s *Server) serveHTTP2(conn net.Conn) { adapter := NewFastHTTPHandlerAdapter(s.handler) + adapter.MaxBodySize = s.maxBodySize opts := &http2.ServeConnOpts{ Context: context.Background(), diff --git a/internal/http3/adapter.go b/internal/http3/adapter.go index 4c86701..97e150b 100644 --- a/internal/http3/adapter.go +++ b/internal/http3/adapter.go @@ -67,7 +67,11 @@ func (a *Adapter) Wrap(handler fasthttp.RequestHandler) http.Handler { a.ResetContext(ctx) // 转换请求 - a.convertRequest(r, ctx) + if err := a.convertRequest(r, ctx); err != nil { + // Write a 413 response directly. + http.Error(w, "Request Entity Too Large", http.StatusRequestEntityTooLarge) + return + } // 设置 ResponseWriter 用于后续写入 ctx.SetUserValue("http3_response_writer", w) @@ -85,7 +89,10 @@ func (a *Adapter) Wrap(handler fasthttp.RequestHandler) http.Handler { // 参数： // - r: 标准库 HTTP 请求 // - ctx: FastHTTP 请求上下文 -func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) { +// +// 返回值： +// - error: 请求体超过限制时返回错误 +func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) error { // 设置方法 ctx.Request.Header.SetMethod(r.Method) @@ -107,7 +114,9 @@ func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) { } // 设置请求体（使用流式处理优化） - a.StreamRequestBody(r, ctx) + if err := a.StreamRequestBody(r, ctx); err != nil { + return err + } // 设置远程地址 if r.RemoteAddr != "" { @@ -118,6 +127,7 @@ func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) { // 设置协议版本 ctx.Request.Header.SetProtocol("HTTP/3") + return nil } // convertResponse 将 fasthttp.RequestCtx 响应写入 http.ResponseWriter。 diff --git a/internal/http3/server.go b/internal/http3/server.go index 79b5460..dfdd026 100644 --- a/internal/http3/server.go +++ b/internal/http3/server.go @@ -59,6 +59,9 @@ type Server struct { // mu 读写锁 mu sync.RWMutex + + // maxBodySize 请求体大小限制 + maxBodySize int64 } // NewServer 创建 HTTP/3 服务器。 @@ -85,12 +88,14 @@ func NewServer(cfg *config.HTTP3Config, handler fasthttp.RequestHandler, tlsConf } adapter := NewAdapter() + adapter.MaxBodySize = cfg.MaxBodySize return &Server{ - config: cfg, - handler: handler, - adapter: adapter, - tlsConfig: tlsConfig, + config: cfg, + handler: handler, + adapter: adapter, + tlsConfig: tlsConfig, + maxBodySize: cfg.MaxBodySize, }, nil } diff --git a/internal/integration/e2e_bench_test.go b/internal/integration/e2e_bench_test.go index 3d12e0c..137a4f0 100644 --- a/internal/integration/e2e_bench_test.go +++ b/internal/integration/e2e_bench_test.go @@ -1,3 +1,5 @@ +//go:build integration + // Package integration 提供端到端集成基准测试。 // // 该文件测试完整请求路径的吞吐量，涵盖静态文件、代理转发、 diff --git a/internal/integration/e2e_failover_bench_test.go b/internal/integration/e2e_failover_bench_test.go index 9851f09..de5e62d 100644 --- a/internal/integration/e2e_failover_bench_test.go +++ b/internal/integration/e2e_failover_bench_test.go @@ -1,3 +1,5 @@ +//go:build integration + // Package integration 提供后端故障切换 E2E 基准测试。 // // 该文件测试负载均衡器剔除/恢复后端的开销。 diff --git a/internal/integration/regex_config_test.go b/internal/integration/regex_config_test.go index 1018c9f..aea716c 100644 --- a/internal/integration/regex_config_test.go +++ b/internal/integration/regex_config_test.go @@ -1,3 +1,5 @@ +//go:build integration + package integration import ( diff --git a/internal/integration/resolver_test.go b/internal/integration/resolver_test.go index 0252e41..1e259f7 100644 --- a/internal/integration/resolver_test.go +++ b/internal/integration/resolver_test.go @@ -1,3 +1,5 @@ +//go:build integration + // resolver_integration_test.go - DNS 解析器集成测试 // // 测试 DNS 解析器与代理的集成 diff --git a/internal/integration/variable_test.go b/internal/integration/variable_test.go index 41796cf..ea68054 100644 --- a/internal/integration/variable_test.go +++ b/internal/integration/variable_test.go @@ -1,3 +1,5 @@ +//go:build integration + // variable_integration_test.go - 变量系统集成测试 // // 测试变量系统与日志、代理、重写的端到端集成 diff --git a/internal/middleware/security/access.go b/internal/middleware/security/access.go index d5d5d70..ba92f90 100644 --- a/internal/middleware/security/access.go +++ b/internal/middleware/security/access.go @@ -348,49 +348,10 @@ func (ac *AccessControl) SetDefault(action string) error { // 返回值： // - net.IP: 客户端 IP 地址，无法获取时返回 nil func (ac *AccessControl) getClientIP(ctx *fasthttp.RequestCtx) net.IP { - remoteIP := netutil.GetRemoteAddrIP(ctx) - - if len(ac.trustedProxies) == 0 || remoteIP == nil { - return remoteIP - } - - isTrusted := false - for _, network := range ac.trustedProxies { - if network.Contains(remoteIP) { - isTrusted = true - break - } - } - if !isTrusted { - return remoteIP - } - - if xff := ctx.Request.Header.Peek("X-Forwarded-For"); len(xff) > 0 { - ips := strings.Split(string(xff), ",") - for _, v := range slices.Backward(ips) { - ipStr := strings.TrimSpace(v) - if ip := net.ParseIP(ipStr); ip != nil { - trusted := false - for _, network := range ac.trustedProxies { - if network.Contains(ip) { - trusted = true - break - } - } - if !trusted { - return ip - } - } - } - } - - if xri := ctx.Request.Header.Peek("X-Real-IP"); len(xri) > 0 { - if ip := net.ParseIP(string(xri)); ip != nil { - return ip - } - } - - return remoteIP + ac.mu.RLock() + trusted := ac.trustedProxies + ac.mu.RUnlock() + return netutil.ExtractClientIPWithTrustedProxies(ctx, trusted) } // AccessStats 访问控制统计信息结构。 diff --git a/internal/middleware/security/access_coverage_test.go b/internal/middleware/security/access_coverage_test.go index 7793290..dabf678 100644 --- a/internal/middleware/security/access_coverage_test.go +++ b/internal/middleware/security/access_coverage_test.go @@ -202,7 +202,7 @@ func TestAccessControlProcess_TrustedProxies_UntrustedSource(t *testing.T) { } } -// TestAccessControlProcess_TrustedProxies_XRealIP 测试可信代理 X-Real-IP +// TestAccessControlProcess_TrustedProxies_XRealIP 测试可信代理 X-Real-IP 不被单独信任 func TestAccessControlProcess_TrustedProxies_XRealIP(t *testing.T) { ac, err := NewAccessControl(&config.AccessConfig{ TrustedProxies: []string{"10.0.0.0/8"}, @@ -225,8 +225,8 @@ func TestAccessControlProcess_TrustedProxies_XRealIP(t *testing.T) { ctx.Request.Header.Set("X-Real-IP", "192.168.1.50") handler(ctx) - if !called { - t.Error("Process() should use X-Real-IP from trusted proxy") + if called { + t.Error("Process() should NOT use X-Real-IP without X-Forwarded-For") } } diff --git a/internal/middleware/security/ratelimit.go b/internal/middleware/security/ratelimit.go index ecf5440..bb166cb 100644 --- a/internal/middleware/security/ratelimit.go +++ b/internal/middleware/security/ratelimit.go @@ -32,6 +32,7 @@ package security import ( "errors" "fmt" + "net" "sync" "sync/atomic" "time" @@ -65,14 +66,15 @@ type shardedBucket struct { // - 所有方法均为并发安全 // - 启动后会自动后台清理过期的桶 type RateLimiter struct { - shards [16]shardedBucket - keyFunc KeyFunc - cleanupTicker *time.Ticker - stopCleanupCh chan struct{} - cleanupDone chan struct{} - stopOnce sync.Once - rate float64 - burst float64 + shards [16]shardedBucket + keyFunc KeyFunc + cleanupTicker *time.Ticker + stopCleanupCh chan struct{} + cleanupDone chan struct{} + stopOnce sync.Once + rate float64 + burst float64 + trustedProxies []net.IPNet } // tokenBucket 表示单个限流键的令牌桶。 @@ -95,11 +97,12 @@ type KeyFunc func(ctx *fasthttp.RequestCtx) string // // 参数： // - cfg: 限流配置，包含速率、突发量和键类型 +// - trustedProxies: 可信代理 CIDR 列表，用于安全提取客户端真实 IP // // 返回值： // - *RateLimiter: 配置好的限流器实例 // - error: 配置无效时返回错误（如速率小于 0） -func NewRateLimiter(cfg *config.RateLimitConfig) (middleware.Middleware, error) { +func NewRateLimiter(cfg *config.RateLimitConfig, trustedProxies ...string) (middleware.Middleware, error) { if cfg == nil { return nil, errors.New("rate limit config is nil") } @@ -116,7 +119,7 @@ func NewRateLimiter(cfg *config.RateLimitConfig) (middleware.Middleware, error) switch algorithm { case "token_bucket", "": - return newTokenBucketLimiter(cfg) + return newTokenBucketLimiter(cfg, trustedProxies) case "sliding_window": window := time.Duration(cfg.SlidingWindow) * time.Second if window <= 0 { @@ -130,7 +133,7 @@ func NewRateLimiter(cfg *config.RateLimitConfig) (middleware.Middleware, error) } // newTokenBucketLimiter 创建令牌桶限流器。 -func newTokenBucketLimiter(cfg *config.RateLimitConfig) (*RateLimiter, error) { +func newTokenBucketLimiter(cfg *config.RateLimitConfig, trustedProxies []string) (*RateLimiter, error) { if cfg.Burst < cfg.RequestRate { return nil, errors.New("burst must be at least equal to request rate") } @@ -149,12 +152,24 @@ func newTokenBucketLimiter(cfg *config.RateLimitConfig) (*RateLimiter, error) { } } - // 根据配置设置键提取函数 - keyFunc, err := parseKeyFunc(cfg.Key) - if err != nil { - return nil, err + // 解析可信代理 CIDR + for _, cidr := range trustedProxies { + _, ipNet, err := net.ParseCIDR(cidr) + if err != nil { + return nil, fmt.Errorf("invalid trusted proxy CIDR %s: %w", cidr, err) + } + rl.trustedProxies = append(rl.trustedProxies, *ipNet) + } + + // 根据配置设置键提取函数 + switch cfg.Key { + case "ip", "": + rl.keyFunc = rl.keyByIP + case rateLimitHeader: + rl.keyFunc = rl.keyByHeader + default: + return nil, fmt.Errorf("unknown key type: %s", cfg.Key) } - rl.keyFunc = keyFunc // 启动后台清理 goroutine rl.startCleanup(10 * time.Minute) @@ -396,6 +411,29 @@ func keyByHeader(ctx *fasthttp.RequestCtx) string { return string(key) } +// keyByIP 提取客户端 IP 作为限流键（令牌桶限流器方法）。 +// +// 仅当请求来自可信代理时，才解析 X-Forwarded-For 头部获取真实客户端 IP。 +func (rl *RateLimiter) keyByIP(ctx *fasthttp.RequestCtx) string { + ip := netutil.ExtractClientIPWithTrustedProxies(ctx, rl.trustedProxies) + if ip == nil { + return accessUnknown + } + return ip.String() +} + +// keyByHeader 提取头部值作为限流键（令牌桶限流器方法）。 +// +// 默认使用 X-RateLimit-Key 头部，如果不存在则回退到 IP。 +func (rl *RateLimiter) keyByHeader(ctx *fasthttp.RequestCtx) string { + key := ctx.Request.Header.Peek("X-RateLimit-Key") + if len(key) == 0 { + // 头部不存在时回退到 IP + return rl.keyByIP(ctx) + } + return string(key) +} + // parseKeyFunc 根据配置字符串解析键提取函数。 // // 支持的键类型： diff --git a/internal/middleware/security/ratelimit_test.go b/internal/middleware/security/ratelimit_test.go index f959fbf..4503662 100644 --- a/internal/middleware/security/ratelimit_test.go +++ b/internal/middleware/security/ratelimit_test.go @@ -12,6 +12,7 @@ package security import ( + "net" "runtime" "testing" "time" @@ -902,3 +903,63 @@ func TestRateLimiter_MultipleCreateDestroy(t *testing.T) { t.Fatalf("goroutine leak detected: before=%d, after=%d", before, after) } } + +// TestRateLimiter_TrustedProxies 验证令牌桶限流器使用可信代理安全提取客户端 IP。 +func TestRateLimiter_TrustedProxies(t *testing.T) { + mw, err := NewRateLimiter(&config.RateLimitConfig{ + RequestRate: 1, + Burst: 1, + }, "10.0.0.0/8") + if err != nil { + t.Fatalf("NewRateLimiter() error: %v", err) + } + + rl, ok := mw.(*RateLimiter) + if !ok { + t.Fatalf("Expected *RateLimiter, got %T", mw) + } + defer rl.StopCleanup() + + callCount := 0 + nextHandler := func(ctx *fasthttp.RequestCtx) { + callCount++ + ctx.SetStatusCode(200) + } + handler := rl.Process(nextHandler) + + // 来自可信代理的请求，使用 X-Forwarded-For 中的真实客户端 IP + ctx1 := testutil.NewRequestCtx("GET", "/test") + ctx1.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")}) + ctx1.Request.Header.Set("X-Forwarded-For", "1.2.3.4") + handler(ctx1) + if ctx1.Response.StatusCode() != 200 { + t.Errorf("first trusted request should be allowed, got %d", ctx1.Response.StatusCode()) + } + + // 同一个真实客户端 IP 的第二个请求应被拒绝 + ctx2 := testutil.NewRequestCtx("GET", "/test") + ctx2.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.2")}) + ctx2.Request.Header.Set("X-Forwarded-For", "1.2.3.4") + handler(ctx2) + if ctx2.Response.StatusCode() != 429 { + t.Errorf("second request from same client IP should be rate limited, got %d", ctx2.Response.StatusCode()) + } + + // 不同真实客户端 IP 的请求应被允许 + ctx3 := testutil.NewRequestCtx("GET", "/test") + ctx3.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")}) + ctx3.Request.Header.Set("X-Forwarded-For", "5.6.7.8") + handler(ctx3) + if ctx3.Response.StatusCode() != 200 { + t.Errorf("request from different client IP should be allowed, got %d", ctx3.Response.StatusCode()) + } + + // 不可信代理：忽略 XFF，使用 remoteAddr 作为限流键 + ctx4 := testutil.NewRequestCtx("GET", "/test") + ctx4.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("192.168.1.1")}) + ctx4.Request.Header.Set("X-Forwarded-For", "1.2.3.4") + handler(ctx4) + if ctx4.Response.StatusCode() != 200 { + t.Errorf("request from untrusted proxy should use remote IP and be allowed, got %d", ctx4.Response.StatusCode()) + } +} diff --git a/internal/netutil/ip.go b/internal/netutil/ip.go index 5c7123a..6b6976a 100644 --- a/internal/netutil/ip.go +++ b/internal/netutil/ip.go @@ -112,6 +112,48 @@ func GetRemoteAddrIP(ctx *fasthttp.RequestCtx) net.IP { return nil } +// ExtractClientIPWithTrustedProxies 从请求上下文安全提取客户端 IP。 +// 仅当 remoteAddr 属于 trustedProxies 时，才解析 X-Forwarded-For 头部， +// 并取右侧（最接近服务器）第一个非可信 IP 作为真实客户端 IP。 +func ExtractClientIPWithTrustedProxies(ctx *fasthttp.RequestCtx, trustedProxies []net.IPNet) net.IP { + remoteIP := GetRemoteAddrIP(ctx) + if remoteIP == nil || len(trustedProxies) == 0 { + return remoteIP + } + + trusted := false + for i := range trustedProxies { + if trustedProxies[i].Contains(remoteIP) { + trusted = true + break + } + } + if !trusted { + return remoteIP + } + + if xff := ctx.Request.Header.Peek("X-Forwarded-For"); len(xff) > 0 { + ips := strings.Split(string(xff), ",") + for i := len(ips) - 1; i >= 0; i-- { + ipStr := strings.TrimSpace(ips[i]) + if ip := net.ParseIP(ipStr); ip != nil { + isTrusted := false + for j := range trustedProxies { + if trustedProxies[j].Contains(ip) { + isTrusted = true + break + } + } + if !isTrusted { + return ip + } + } + } + } + + return remoteIP +} + // remoteAddrCache 缓存 RemoteAddr 字符串化结果，避免重复的 net.TCPAddr.String() 分配。 type remoteAddrCache struct { mu sync.RWMutex diff --git a/internal/netutil/ip_test.go b/internal/netutil/ip_test.go index a2e76fe..5be3fb1 100644 --- a/internal/netutil/ip_test.go +++ b/internal/netutil/ip_test.go @@ -130,3 +130,32 @@ func TestGetRemoteAddrIP(_ *testing.T) { // Just verify it doesn't panic _ = got } + +func TestExtractClientIPWithTrustedProxies(t *testing.T) { + trusted := []net.IPNet{mustParseCIDR(t, "10.0.0.0/8")} + + ctx := &fasthttp.RequestCtx{} + ctx.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")}) + ctx.Request.Header.Set("X-Forwarded-For", "1.2.3.4, 10.0.0.2") + + ip := ExtractClientIPWithTrustedProxies(ctx, trusted) + if ip == nil || ip.String() != "1.2.3.4" { + t.Fatalf("expected 1.2.3.4, got %v", ip) + } + + // Untrusted proxy: ignore XFF. + ctx.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("192.168.1.1")}) + ip = ExtractClientIPWithTrustedProxies(ctx, trusted) + if ip == nil || ip.String() != "192.168.1.1" { + t.Fatalf("expected 192.168.1.1, got %v", ip) + } +} + +func mustParseCIDR(t *testing.T, s string) net.IPNet { + t.Helper() + _, ipNet, err := net.ParseCIDR(s) + if err != nil { + t.Fatalf("failed to parse CIDR %s: %v", s, err) + } + return *ipNet +} diff --git a/internal/proxy/cache_handler.go b/internal/proxy/cache_handler.go index c1e1a36..cc768e8 100644 --- a/internal/proxy/cache_handler.go +++ b/internal/proxy/cache_handler.go @@ -15,41 +15,46 @@ import ( // buildCacheKeyHash 使用 FNV-64a 计算缓存键的 uint64 哈希值。 // 使用零分配方式构建哈希，避免 []byte(origKey) 转换。 -func (p *Proxy) buildCacheKeyHash(ctx *fasthttp.RequestCtx) (uint64, string) { +// 缓存键包含请求方法、Host、URI 以及配置的 Vary 请求头，防止缓存中毒。 +func (p *Proxy) buildCacheKeyHash(ctx *fasthttp.RequestCtx, varyHeaders []string) (uint64, string) { + return p.buildCacheKeyHashWithHost(ctx, ctx.Request.Header.Host(), varyHeaders) +} + +// buildCacheKeyHashWithHost 与 buildCacheKeyHash 相同，但使用显式传入的 Host， +// 避免请求头在缓存键计算前被修改（例如 proxy 将 Host 改写为上游目标）。 +func (p *Proxy) buildCacheKeyHashWithHost(ctx *fasthttp.RequestCtx, host []byte, varyHeaders []string) (uint64, string) { method := ctx.Request.Header.Method() uri := ctx.Request.URI().RequestURI() var h uint64 = 14695981039346656037 - for i := 0; i < len(method); i++ { - h ^= uint64(method[i]) - h *= 1099511628211 - } - h ^= uint64(':') - h *= 1099511628211 - for i := 0; i < len(uri); i++ { - h ^= uint64(uri[i]) + for _, b := range [][]byte{method, host, uri} { + h ^= uint64(':') h *= 1099511628211 + for i := 0; i < len(b); i++ { + h ^= uint64(b[i]) + h *= 1099511628211 + } } - origKey := b2s(method) + ":" + b2s(uri) + for _, name := range varyHeaders { + value := ctx.Request.Header.Peek(name) + h ^= uint64(':') + h *= 1099511628211 + for i := 0; i < len(value); i++ { + h ^= uint64(value[i]) + h *= 1099511628211 + } + } + + origKey := b2s(method) + ":" + b2s(host) + b2s(uri) + for _, name := range varyHeaders { + origKey += ":" + name + "=" + b2s(ctx.Request.Header.Peek(name)) + } return h, origKey } -func (p *Proxy) buildCacheKeyHashValue(ctx *fasthttp.RequestCtx) uint64 { - method := ctx.Request.Header.Method() - uri := ctx.Request.URI().RequestURI() - - var h uint64 = 14695981039346656037 - for i := 0; i < len(method); i++ { - h ^= uint64(method[i]) - h *= 1099511628211 - } - h ^= uint64(':') - h *= 1099511628211 - for i := 0; i < len(uri); i++ { - h ^= uint64(uri[i]) - h *= 1099511628211 - } +func (p *Proxy) buildCacheKeyHashValue(ctx *fasthttp.RequestCtx, varyHeaders []string) uint64 { + h, _ := p.buildCacheKeyHash(ctx, varyHeaders) return h } @@ -121,14 +126,8 @@ func (p *Proxy) backgroundRefresh(req *fasthttp.Request, target *loadbalance.Tar return } - // 提取响应头（使用 pool 复用 map） - headers, ok := headersPool.Get().(map[string]string) - if !ok { - headers = make(map[string]string, 20) - } - for k := range headers { - delete(headers, k) - } + // 提取响应头 + headers := make(map[string]string, 20) for key, value := range resp.Header.All() { headers[string(key)] = string(value) } diff --git a/internal/proxy/cache_handler_test.go b/internal/proxy/cache_handler_test.go new file mode 100644 index 0000000..0474e42 --- /dev/null +++ b/internal/proxy/cache_handler_test.go @@ -0,0 +1,28 @@ +package proxy + +import ( + "testing" + + "github.com/valyala/fasthttp" +) + +func TestBuildCacheKeyIncludesHostAndVary(t *testing.T) { + p := &Proxy{} + ctx1 := &fasthttp.RequestCtx{} + ctx1.Request.Header.SetMethod("GET") + ctx1.Request.Header.SetHost("a.example.com") + ctx1.Request.SetRequestURI("/api/data") + + ctx2 := &fasthttp.RequestCtx{} + ctx2.Request.Header.SetMethod("GET") + ctx2.Request.Header.SetHost("b.example.com") + ctx2.Request.SetRequestURI("/api/data") + + vary := []string{"Accept-Encoding"} + h1, key1 := p.buildCacheKeyHash(ctx1, vary) + h2, key2 := p.buildCacheKeyHash(ctx2, vary) + + if h1 == h2 || key1 == key2 { + t.Fatal("cache keys must differ by Host") + } +} diff --git a/internal/proxy/cache_key_allocation_test.go b/internal/proxy/cache_key_allocation_test.go index 5834c5a..9b069f9 100644 --- a/internal/proxy/cache_key_allocation_test.go +++ b/internal/proxy/cache_key_allocation_test.go @@ -42,7 +42,7 @@ func BenchmarkCacheKeyHashValue_ZeroAlloc(b *testing.B) { b.ResetTimer() for b.Loop() { - hash := p.buildCacheKeyHashValue(ctx) + hash := p.buildCacheKeyHashValue(ctx, nil) _ = hash } } @@ -72,7 +72,7 @@ func BenchmarkCacheKeyHash_WithAlloc(b *testing.B) { b.ResetTimer() for b.Loop() { - hash, key := p.buildCacheKeyHash(ctx) + hash, key := p.buildCacheKeyHash(ctx, nil) _ = hash _ = key } @@ -96,7 +96,7 @@ func BenchmarkCacheKeyHash_Compare(b *testing.B) { b.ReportAllocs() b.ResetTimer() for b.Loop() { - _ = p.buildCacheKeyHashValue(ctx) + _ = p.buildCacheKeyHashValue(ctx, nil) } }) @@ -104,7 +104,7 @@ func BenchmarkCacheKeyHash_Compare(b *testing.B) { b.ReportAllocs() b.ResetTimer() for b.Loop() { - _, _ = p.buildCacheKeyHash(ctx) + _, _ = p.buildCacheKeyHash(ctx, nil) } }) } diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 6510cd0..3cd89bb 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -42,6 +42,7 @@ import ( "time" "github.com/valyala/fasthttp" + "golang.org/x/sync/singleflight" "rua.plus/lolly/internal/cache" "rua.plus/lolly/internal/config" "rua.plus/lolly/internal/loadbalance" @@ -92,16 +93,6 @@ const ( lbSticky = "sticky" // 会话粘性 ) -// headersPool 复用缓存 headers map，减少分配。 -// 预容量 20 覆盖大多数 HTTP 响应头数量。 -// 注意：从 pool 获取的 map 使用后不能 Put 回 pool， -// 因为 cache.Set 存储了 map 引用。 -var headersPool = sync.Pool{ - New: func() any { - return make(map[string]string, 20) - }, -} - var upstreamTimingPool = sync.Pool{ New: func() any { return NewUpstreamTiming() @@ -131,6 +122,7 @@ type Proxy struct { mu sync.RWMutex // 保护并发访问的读写锁 started atomic.Bool // 代理启动标志 cacheIgnoreSet map[string]bool // 缓存时忽略的响应头集合 + refreshGroup singleflight.Group // 合并并发后台缓存刷新 } // NewProxy 使用给定的配置和后台目标创建一个新的反向代理实例。 @@ -229,6 +221,39 @@ func NewProxy(cfg *config.ProxyConfig, targets []*loadbalance.Target, transportC return p, nil } +// cacheVaryHeaders 返回用于构建缓存键的 Vary 请求头列表。 +// +// 当 varyResponse 非空时，将响应中的 Vary 头与配置的值合并去重。 +// 请求头名称统一按大小写不敏感方式比较，避免重复。 +func (p *Proxy) cacheVaryHeaders(varyResponse []byte) []string { + configured := p.config.Cache.Vary + if len(varyResponse) == 0 { + return configured + } + + result := make([]string, 0, len(configured)+bytes.Count(varyResponse, []byte(","))+1) + result = append(result, configured...) + seen := make(map[string]struct{}, len(result)) + for _, h := range result { + seen[strings.ToLower(h)] = struct{}{} + } + + for _, part := range strings.Split(b2s(varyResponse), ",") { + name := strings.TrimSpace(part) + if name == "" { + continue + } + key := strings.ToLower(name) + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + result = append(result, name) + } + + return result +} + // stickyBalancer wraps StickySession to implement loadbalance.Balancer. // It delegates Select/SelectExcluding to the fallback balancer while // allowing the proxy to access the StickySession for cookie-based routing. @@ -583,15 +608,11 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { } timing.reset() - var cacheHashKey uint64 - var cacheOrigKey string - cacheKeyComputed := false - computeCacheKey := func() (uint64, string) { - if !cacheKeyComputed { - cacheHashKey, cacheOrigKey = p.buildCacheKeyHash(ctx) - cacheKeyComputed = true - } - return cacheHashKey, cacheOrigKey + // 在请求头被修改前捕获原始 Host，用于构建缓存键。 + // 防止 modifyRequestHeaders 将 Host 改为上游目标后导致缓存键丢失客户端 Host 区分。 + cacheHost := ctx.Request.Header.Host() + computeCacheKey := func(varyHeaders []string) (uint64, string) { + return p.buildCacheKeyHashWithHost(ctx, cacheHost, varyHeaders) } // 创建变量上下文用于设置上游变量 @@ -681,18 +702,18 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { // 检查是否为 WebSocket 升级请求 if isWebSocketRequest(ctx) { - // WebSocket 使用 defer 确保连接计数释放 - defer loadbalance.DecrementConnections(target) + defer func() { + loadbalance.DecrementConnections(target) + }() + defer timing.MarkConnectEnd() timing.MarkConnectStart() - err := WebSocket(ctx, target, p.config.Timeout.Connect, &p.config.Headers) - timing.MarkConnectEnd() + err := WebSocket(ctx, target, p.config.Timeout.Connect, &p.config.Headers, p.config.ProxySSL) if err != nil { - upstreamStatus = 502 - logging.Error().Msgf("WebSocket proxy error: %v", err) + upstreamStatus = fasthttp.StatusBadGateway + logging.Error().Err(err).Msg("WebSocket proxy error") return } - // WebSocket 成功 - upstreamStatus = 101 + upstreamStatus = fasthttp.StatusSwitchingProtocols return } @@ -749,7 +770,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { path := string(ctx.Request.URI().Path()) rule := p.cache.MatchRule(path, method, 0) if rule != nil { - hashKey, origKey := computeCacheKey() + hashKey, origKey := computeCacheKey(p.cacheVaryHeaders(nil)) if entry, ok, stale := p.cache.Get(hashKey, origKey); ok { // 缓存命中 loadbalance.DecrementConnections(target) @@ -768,10 +789,17 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { entry.Updating.Store(true) reqCopy := fasthttp.AcquireRequest() ctx.Request.CopyTo(reqCopy) + key := origKey go func() { defer entry.Updating.Store(false) defer fasthttp.ReleaseRequest(reqCopy) - p.backgroundRefresh(reqCopy, target, hashKey, origKey) + _, err, _ := p.refreshGroup.Do(key, func() (interface{}, error) { + p.backgroundRefresh(reqCopy, target, hashKey, origKey) + return nil, nil + }) + if err != nil { + logging.Error().Err(err).Msg("background cache refresh failed") + } }() } upstreamAddr = upstreamCache @@ -840,7 +868,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { // 尝试使用 stale 缓存 if p.cache != nil { - hashKey, origKey := computeCacheKey() + hashKey, origKey := computeCacheKey(p.cacheVaryHeaders(nil)) isTimeout := errors.Is(err, fasthttp.ErrTimeout) if staleEntry, ok := p.cache.GetStale(hashKey, origKey, isTimeout); ok { logging.Info().Msgf("[PROXY] 使用 stale 缓存: key=%s, isTimeout=%v", origKey, isTimeout) @@ -853,8 +881,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { // 释放缓存锁 if p.cache != nil && attempt == 0 { - computeCacheKey() - hashKey := cacheHashKey + hashKey, _ := computeCacheKey(p.cacheVaryHeaders(nil)) p.cache.ReleaseLock(hashKey, err) } @@ -904,8 +931,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { if shouldRetry { // 释放缓存锁 if p.cache != nil && attempt == 0 { - computeCacheKey() - hashKey := cacheHashKey + hashKey, _ := computeCacheKey(p.cacheVaryHeaders(nil)) p.cache.ReleaseLock(hashKey, fmt.Errorf("HTTP %d", statusCode)) } @@ -934,7 +960,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { return } - hashKey, origKey := computeCacheKey() + hashKey, origKey := computeCacheKey(p.cacheVaryHeaders(ctx.Response.Header.Peek("Vary"))) if statusCode >= 200 && statusCode < 300 { // 检查 MinUses 阈值 if entry, ok, _ := p.cache.Get(hashKey, origKey); ok { @@ -945,14 +971,8 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) { } } - // 提取响应头（使用 pool 复用 map） - headers, ok := headersPool.Get().(map[string]string) - if !ok { - headers = make(map[string]string, 20) - } - for k := range headers { - delete(headers, k) - } + // 提取响应头 + headers := make(map[string]string, 20) var lastModified, etag string for key, value := range ctx.Response.Header.All() { diff --git a/internal/proxy/proxy_bench_test.go b/internal/proxy/proxy_bench_test.go index f8e60c8..0f79e80 100644 --- a/internal/proxy/proxy_bench_test.go +++ b/internal/proxy/proxy_bench_test.go @@ -451,7 +451,7 @@ func BenchmarkBuildCacheKeyHash(b *testing.B) { b.Run("buildCacheKeyHash_with_string", func(b *testing.B) { b.ResetTimer() for b.Loop() { - hashKey, _ := p.buildCacheKeyHash(ctx) + hashKey, _ := p.buildCacheKeyHash(ctx, nil) _ = hashKey } }) @@ -459,7 +459,7 @@ func BenchmarkBuildCacheKeyHash(b *testing.B) { b.Run("buildCacheKeyHashValue_direct", func(b *testing.B) { b.ResetTimer() for b.Loop() { - hashKey := p.buildCacheKeyHashValue(ctx) + hashKey := p.buildCacheKeyHashValue(ctx, nil) _ = hashKey } }) @@ -585,7 +585,7 @@ func BenchmarkProxyZeroAllocPath(b *testing.B) { b.ReportAllocs() b.ResetTimer() for b.Loop() { - hash := p.buildCacheKeyHashValue(ctx) + hash := p.buildCacheKeyHashValue(ctx, nil) _ = hash } }) @@ -594,7 +594,7 @@ func BenchmarkProxyZeroAllocPath(b *testing.B) { b.ReportAllocs() b.ResetTimer() for b.Loop() { - hash, key := p.buildCacheKeyHash(ctx) + hash, key := p.buildCacheKeyHash(ctx, nil) _ = hash _ = key } diff --git a/internal/proxy/proxy_coverage_extra_test.go b/internal/proxy/proxy_coverage_extra_test.go index ee795c0..acf59de 100644 --- a/internal/proxy/proxy_coverage_extra_test.go +++ b/internal/proxy/proxy_coverage_extra_test.go @@ -848,14 +848,14 @@ func TestSelectTarget_EmptyTargets(t *testing.T) { func TestDialTarget(t *testing.T) { t.Run("连接超时", func(t *testing.T) { // 使用不可达地址测试超时 - _, err := dialTarget("http://10.255.255.1:9999", 100*time.Millisecond) + _, err := dialTarget("http://10.255.255.1:9999", 100*time.Millisecond, nil) if err == nil { t.Error("dialTarget() should return error for unreachable address") } }) t.Run("HTTPS 连接失败", func(t *testing.T) { - _, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond) + _, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond, nil) if err == nil { t.Error("dialTarget() should return error for unreachable HTTPS address") } @@ -876,7 +876,7 @@ func TestDialTarget(t *testing.T) { }() addr := ln.Addr().String() - conn, err := dialTarget("http://"+addr, 1*time.Second) + conn, err := dialTarget("http://"+addr, 1*time.Second, nil) if err != nil { t.Errorf("dialTarget() error: %v", err) } @@ -972,7 +972,7 @@ func TestWebSocket_ErrorCases(t *testing.T) { target.Healthy.Store(true) // 使用很短的超时 - err := WebSocket(ctx, target, 10*time.Millisecond, nil) + err := WebSocket(ctx, target, 10*time.Millisecond, nil, nil) if err == nil { t.Error("WebSocket() should return error for invalid backend") } @@ -983,7 +983,7 @@ func TestWebSocket_ErrorCases(t *testing.T) { func TestDialTarget_TLS_Extra(t *testing.T) { t.Run("TLS 握手失败", func(t *testing.T) { // 使用不可达的 HTTPS 地址 - _, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond) + _, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond, nil) if err == nil { t.Error("dialTarget() should return error for unreachable HTTPS address") } @@ -1469,7 +1469,7 @@ func TestBackgroundRefresh_304(t *testing.T) { // 预先设置缓存条目 ctx := testutil.NewRequestCtx("GET", "/api/test") - hashKey, origKey := p.buildCacheKeyHash(ctx) + hashKey, origKey := p.buildCacheKeyHash(ctx, nil) p.cache.Set(hashKey, origKey, []byte("cached"), map[string]string{ "Last-Modified": "Tue, 20 Oct 2015 07:28:00 GMT", "ETag": "\"old\"", diff --git a/internal/proxy/proxy_coverage_test.go b/internal/proxy/proxy_coverage_test.go index b41a869..fc2a5ba 100644 --- a/internal/proxy/proxy_coverage_test.go +++ b/internal/proxy/proxy_coverage_test.go @@ -235,7 +235,7 @@ func TestBuildCacheKeyHash(t *testing.T) { ctx := testutil.NewRequestCtx("GET", "/api/test") - hashKey, origKey := p.buildCacheKeyHash(ctx) + hashKey, origKey := p.buildCacheKeyHash(ctx, nil) if hashKey == 0 { t.Error("buildCacheKeyHash() should return non-zero hash") } @@ -245,14 +245,14 @@ func TestBuildCacheKeyHash(t *testing.T) { // 相同请求应产生相同哈希 ctx2 := testutil.NewRequestCtx("GET", "/api/test") - hashKey2, _ := p.buildCacheKeyHash(ctx2) + hashKey2, _ := p.buildCacheKeyHash(ctx2, nil) if hashKey != hashKey2 { t.Error("Same request should produce same hash") } // 不同请求应产生不同哈希 ctx3 := testutil.NewRequestCtx("POST", "/api/other") - hashKey3, _ := p.buildCacheKeyHash(ctx3) + hashKey3, _ := p.buildCacheKeyHash(ctx3, nil) if hashKey == hashKey3 { t.Error("Different request should produce different hash") } @@ -273,13 +273,13 @@ func TestBuildCacheKeyHashValue(t *testing.T) { ctx := testutil.NewRequestCtx("GET", "/api/test") - hashValue := p.buildCacheKeyHashValue(ctx) + hashValue := p.buildCacheKeyHashValue(ctx, nil) if hashValue == 0 { t.Error("buildCacheKeyHashValue() should return non-zero hash") } // 应该与 buildCacheKeyHash 结果一致 - hashKey, _ := p.buildCacheKeyHash(ctx) + hashKey, _ := p.buildCacheKeyHash(ctx, nil) if hashValue != hashKey { t.Error("buildCacheKeyHashValue() should match buildCacheKeyHash()") } @@ -624,7 +624,8 @@ func TestServeHTTP_CacheHit(t *testing.T) { // 预填充缓存 ctx := testutil.NewRequestCtx("GET", "/api/cached") - hashKey, origKey := p.buildCacheKeyHash(ctx) + ctx.Request.Header.SetHost("localhost:8080") + hashKey, origKey := p.buildCacheKeyHash(ctx, nil) p.cache.Set(hashKey, origKey, []byte("cached!"), map[string]string{ "Content-Type": "text/plain", }, 200, 10*time.Second) @@ -705,7 +706,8 @@ func TestServeHTTP_WithRedirectRewrite_CacheHit(t *testing.T) { } ctx := testutil.NewRequestCtx("GET", "/api/test") - hashKey, origKey := p.buildCacheKeyHash(ctx) + ctx.Request.Header.SetHost("localhost:8080") + hashKey, origKey := p.buildCacheKeyHash(ctx, nil) p.cache.Set(hashKey, origKey, []byte("ok"), map[string]string{ "Content-Type": "text/plain", }, 200, 10*time.Second) diff --git a/internal/proxy/proxy_low_coverage_test.go b/internal/proxy/proxy_low_coverage_test.go index b904d34..1a9dcb0 100644 --- a/internal/proxy/proxy_low_coverage_test.go +++ b/internal/proxy/proxy_low_coverage_test.go @@ -654,7 +654,7 @@ func TestBackgroundRefresh_WithCacheEntry(t *testing.T) { require.NoError(t, err) ctx := testutil.NewRequestCtx("GET", "/api/test") - hashKey, origKey := p.buildCacheKeyHash(ctx) + hashKey, origKey := p.buildCacheKeyHash(ctx, nil) p.cache.Set(hashKey, origKey, []byte("old content"), map[string]string{ "Content-Type": "text/plain", @@ -712,7 +712,7 @@ func TestDialTarget_Success(t *testing.T) { } }() - conn, err := dialTarget("http://"+ln.Addr().String(), 1*time.Second) + conn, err := dialTarget("http://"+ln.Addr().String(), 1*time.Second, nil) require.NoError(t, err) require.NotNil(t, conn) _ = conn.Close() @@ -720,7 +720,7 @@ func TestDialTarget_Success(t *testing.T) { // TestDialTarget_Timeout 测试连接超时。 func TestDialTarget_Timeout(t *testing.T) { - _, err := dialTarget("http://10.255.255.1:9999", 50*time.Millisecond) + _, err := dialTarget("http://10.255.255.1:9999", 50*time.Millisecond, nil) require.Error(t, err) assert.Contains(t, err.Error(), "failed to connect") } @@ -770,7 +770,7 @@ func TestWebSocket_BackendSuccess(t *testing.T) { target := &loadbalance.Target{URL: "http://127.0.0.1:1"} target.Healthy.Store(true) - err := WebSocket(ctx, target, 2*time.Second, nil) + err := WebSocket(ctx, target, 2*time.Second, nil, nil) require.Error(t, err) assert.Contains(t, err.Error(), "hijack") } @@ -1068,7 +1068,7 @@ func TestServeHTTP_CacheStaleWhileRevalidate(t *testing.T) { require.NoError(t, err) ctx := testutil.NewRequestCtx("GET", "/api/cached") - hashKey, origKey := p.buildCacheKeyHash(ctx) + hashKey, origKey := p.buildCacheKeyHash(ctx, nil) headers := map[string]string{"Content-Type": "text/plain"} p.cache.Set(hashKey, origKey, []byte("stale content"), headers, 200, -1*time.Second) @@ -1399,7 +1399,7 @@ func TestWebSocket_ReadResponseError(t *testing.T) { target := &loadbalance.Target{URL: "http://" + ln.Addr().String()} target.Healthy.Store(true) - err = WebSocket(ctx, target, 1*time.Second, nil) + err = WebSocket(ctx, target, 1*time.Second, nil, nil) require.Error(t, err) } @@ -1434,7 +1434,7 @@ func TestWebSocket_HijackFails(t *testing.T) { target := &loadbalance.Target{URL: "http://127.0.0.1:1"} target.Healthy.Store(true) - err := WebSocket(ctx, target, 100*time.Millisecond, nil) + err := WebSocket(ctx, target, 100*time.Millisecond, nil, nil) require.Error(t, err) } diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 5c1390e..dcb9ae1 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -774,6 +774,31 @@ func TestHandleWebSocket(t *testing.T) { } } +// TestWebSocketConnectionCountReleasedOnFailure 测试 WebSocket 升级失败后连接计数被释放。 +// testutil 创建的请求上下文不支持 Hijack，会导致 WebSocket 代理失败； +// 失败后应确保 IncrementConnections 对应的 DecrementConnections 被调用。 +func TestWebSocketConnectionCountReleasedOnFailure(t *testing.T) { + cfg := testutil.NewTestProxyConfig("/ws") + targets := testutil.NewTestTargets("http://127.0.0.1:1") + targets[0].MaxFails = 1 + + p, err := NewProxy(cfg, targets, nil, nil) + if err != nil { + t.Fatalf("NewProxy() error: %v", err) + } + + ctx := testutil.NewRequestCtxWithHeader("GET", "/ws", map[string]string{ + "Upgrade": "websocket", + "Connection": "Upgrade", + }) + + p.ServeHTTP(ctx) + + if got := atomic.LoadInt64(&targets[0].Connections); got != 0 { + t.Fatalf("expected connection count 0 after failed WebSocket, got %d", got) + } +} + // TestSetHealthChecker 测试健康检查器设置 // 注意：SetHealthChecker 是公开方法，但 healthChecker 是私有字段 // 此测试验证方法可以正常调用 diff --git a/internal/proxy/websocket.go b/internal/proxy/websocket.go index 1b91599..6a12aba 100644 --- a/internal/proxy/websocket.go +++ b/internal/proxy/websocket.go @@ -219,7 +219,7 @@ func isConnectionClosedError(err error) bool { // 返回值： // - net.Conn: 建立的连接（TLS 连接或普通 TCP 连接） // - error: 连接失败时返回错误 -func dialTarget(targetURL string, timeout time.Duration) (net.Conn, error) { +func dialTarget(targetURL string, timeout time.Duration, proxySSL *config.ProxySSLConfig) (net.Conn, error) { // 解析目标 URL addr, isTLS := netutil.ParseTargetURL(targetURL, true) @@ -233,24 +233,34 @@ func dialTarget(targetURL string, timeout time.Duration) (net.Conn, error) { return nil, fmt.Errorf("failed to connect to target: %w", err) } - // 如果是 HTTPS，建立 TLS 连接 - if isTLS { - tlsConn := tls.Client(conn, &tls.Config{ - InsecureSkipVerify: false, - ServerName: strings.Split(addr, ":")[0], - }) - if err := tlsConn.SetDeadline(time.Now().Add(timeout)); err != nil { - _ = conn.Close() - return nil, fmt.Errorf("failed to set TLS deadline: %w", err) - } - if err := tlsConn.Handshake(); err != nil { - _ = conn.Close() - return nil, fmt.Errorf("TLS handshake failed: %w", err) - } - return tlsConn, nil + if !isTLS { + return conn, nil } - return conn, nil + tlsCfg, err := CreateTLSConfig(proxySSL, addr) + if err != nil { + _ = conn.Close() + return nil, err + } + + host, _, err := net.SplitHostPort(addr) + if err != nil { + host = addr + } + if tlsCfg.ServerName == "" { + tlsCfg.ServerName = host + } + + tlsConn := tls.Client(conn, tlsCfg) + if err := tlsConn.SetDeadline(time.Now().Add(timeout)); err != nil { + _ = conn.Close() + return nil, err + } + if err := tlsConn.Handshake(); err != nil { + _ = conn.Close() + return nil, fmt.Errorf("TLS handshake failed: %w", err) + } + return tlsConn, nil } // buildWebSocketUpgradeRequest 构建 WebSocket 升级 HTTP 请求。 @@ -365,7 +375,7 @@ func readWebSocketUpgradeResponse(conn net.Conn, timeout time.Duration) (*http.R // // 返回值： // - error: 代理过程中的错误 -func WebSocket(ctx *fasthttp.RequestCtx, target *loadbalance.Target, timeout time.Duration, headersConfig *config.ProxyHeaders) error { +func WebSocket(ctx *fasthttp.RequestCtx, target *loadbalance.Target, timeout time.Duration, headersConfig *config.ProxyHeaders, proxySSL *config.ProxySSLConfig) error { // 使用 Hijack 获取客户端 TCP 连接 var clientConn net.Conn @@ -378,7 +388,7 @@ func WebSocket(ctx *fasthttp.RequestCtx, target *loadbalance.Target, timeout tim } // 步骤1: 建立到后端目标的连接 - targetConn, err := dialTarget(target.URL, timeout) + targetConn, err := dialTarget(target.URL, timeout, proxySSL) if err != nil { _ = clientConn.Close() return fmt.Errorf("failed to connect to backend: %w", err) diff --git a/internal/proxy/websocket_test.go b/internal/proxy/websocket_test.go index 3f4492f..cff9ff9 100644 --- a/internal/proxy/websocket_test.go +++ b/internal/proxy/websocket_test.go @@ -136,7 +136,7 @@ func TestIsConnectionClosedError(t *testing.T) { // TestDialTarget_InvalidAddress 测试无效地址的拨号 func TestDialTarget_InvalidAddress(t *testing.T) { // 测试连接到无效端口 - _, err := dialTarget("http://127.0.0.1:1", 100*time.Millisecond) + _, err := dialTarget("http://127.0.0.1:1", 100*time.Millisecond, nil) if err == nil { t.Error("Expected error for invalid address") } @@ -145,7 +145,7 @@ func TestDialTarget_InvalidAddress(t *testing.T) { // TestDialTarget_HTTPS 测试 HTTPS 连接（会失败，但验证错误处理） func TestDialTarget_HTTPS(t *testing.T) { // 测试 HTTPS 连接到无效端口 - _, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond) + _, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond, nil) if err == nil { t.Error("Expected error for invalid HTTPS address") } @@ -258,7 +258,7 @@ func TestDialTarget_URLParsing(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - _, err := dialTarget(tt.url, 10*time.Millisecond) + _, err := dialTarget(tt.url, 10*time.Millisecond, nil) if tt.expectError { if err == nil { t.Error("Expected error, got nil") @@ -589,7 +589,7 @@ func TestReadWebSocketUpgradeResponse_Timeout(t *testing.T) { // TestDialTarget_TLS 测试 TLS 连接（连接无效端口应失败） func TestDialTarget_TLS(t *testing.T) { // 测试 HTTPS 连接到无效端口 - _, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond) + _, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond, nil) if err == nil { t.Error("Expected error for invalid HTTPS address") } diff --git a/internal/server/coverage_test.go b/internal/server/coverage_test.go index e5188b1..5016ed3 100644 --- a/internal/server/coverage_test.go +++ b/internal/server/coverage_test.go @@ -97,8 +97,6 @@ func TestWrapRoutedHandler_有AccessLog(t *testing.T) { } s := New(cfg) - s.accessLogMiddleware = s.accessLogMiddleware - called := false original := func(ctx *fasthttp.RequestCtx) { called = true diff --git a/internal/server/middleware_builder.go b/internal/server/middleware_builder.go index bdfae75..15f9de4 100644 --- a/internal/server/middleware_builder.go +++ b/internal/server/middleware_builder.go @@ -61,7 +61,7 @@ func (s *Server) buildMiddlewareChain(serverCfg *config.ServerConfig) (*middlewa // 3. Security: RateLimiter (速率限制) if serverCfg.Security.RateLimit.RequestRate > 0 { - rl, err := security.NewRateLimiter(&serverCfg.Security.RateLimit) + rl, err := security.NewRateLimiter(&serverCfg.Security.RateLimit, serverCfg.Security.Access.TrustedProxies...) if err != nil { return nil, fmt.Errorf("failed to create rate limiter middleware: %w", err) } diff --git a/internal/server/purge.go b/internal/server/purge.go index 84a18dc..89d526b 100644 --- a/internal/server/purge.go +++ b/internal/server/purge.go @@ -80,7 +80,7 @@ func (h *PurgeHandler) ServeHTTP(ctx *fasthttp.RequestCtx) { } // 检查 IP 访问权限 - if !utils.CheckIPAccess(ctx, h.allowed) { + if !utils.CheckIPAccess(ctx, h.allowed, nil) { utils.SendJSONError(ctx, fasthttp.StatusForbidden, "forbidden") return } diff --git a/internal/server/purge_test.go b/internal/server/purge_test.go index 4301fab..39d5c34 100644 --- a/internal/server/purge_test.go +++ b/internal/server/purge_test.go @@ -236,7 +236,7 @@ func TestPurgeHandler_checkAccess(t *testing.T) { if len(h.allowed) == 0 { // 无白名单时应允许所有访问 - if !utils.CheckIPAccess(nil, h.allowed) { + if !utils.CheckIPAccess(nil, h.allowed, nil) { t.Error("expected access to be true when no allow list configured") } return @@ -654,7 +654,7 @@ func TestPurgeHandler_checkAccess_NilContext(t *testing.T) { } // Empty allow list should allow access (returns true even with nil context) - if !utils.CheckIPAccess(nil, h.allowed) { + if !utils.CheckIPAccess(nil, h.allowed, nil) { t.Error("expected checkAccess to return true with empty allow list") } }) @@ -705,7 +705,7 @@ func TestPurgeHandler_checkAccess_WithAllowedIP(t *testing.T) { ctx.Init(&fasthttp.Request{}, nil, nil) // context with nil remote address - should return false (no client IP) - if utils.CheckIPAccess(ctx, h.allowed) { + if utils.CheckIPAccess(ctx, h.allowed, nil) { t.Error("expected checkAccess to return false with no client IP") } }) diff --git a/internal/server/status.go b/internal/server/status.go index 69a37ea..9d5db4e 100644 --- a/internal/server/status.go +++ b/internal/server/status.go @@ -166,7 +166,7 @@ func (h *StatusHandler) Path() string { // - ctx: FastHTTP 请求上下文 func (h *StatusHandler) ServeHTTP(ctx *fasthttp.RequestCtx) { // 步骤1: 检查 IP 访问权限 - if !utils.CheckIPAccess(ctx, h.allowed) { + if !utils.CheckIPAccess(ctx, h.allowed, nil) { utils.SendErrorWithDetail(ctx, utils.ErrForbidden, "Access denied") return } diff --git a/internal/server/status_test.go b/internal/server/status_test.go index 269f8aa..2bbb1a2 100644 --- a/internal/server/status_test.go +++ b/internal/server/status_test.go @@ -918,7 +918,7 @@ func TestStatusHandler_checkAccess_AllowList(t *testing.T) { ctx := &fasthttp.RequestCtx{} ctx.SetRemoteAddr(&net.TCPAddr{IP: tt.remoteIP, Port: 12345}) - got := utils.CheckIPAccess(ctx, h.allowed) + got := utils.CheckIPAccess(ctx, h.allowed, nil) if got != tt.wantAccess { t.Errorf("expected access %v, got %v", tt.wantAccess, got) } @@ -1214,7 +1214,7 @@ func TestStatusHandler_ServeHTTP_AccessDenied_XForwardedFor(t *testing.T) { } } -func TestStatusHandler_ServeHTTP_AccessAllowed_XForwardedFor(t *testing.T) { +func TestStatusHandler_ServeHTTP_AccessAllowed_RemoteAddr(t *testing.T) { t.Parallel() cfg := &config.StatusConfig{ Path: "/_status", @@ -1232,12 +1232,12 @@ func TestStatusHandler_ServeHTTP_AccessAllowed_XForwardedFor(t *testing.T) { ctx := &fasthttp.RequestCtx{} ctx.Request.SetRequestURI("/_status") - ctx.Request.Header.Set("X-Forwarded-For", "10.5.5.5") + ctx.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.5.5.5"), Port: 12345}) h.ServeHTTP(ctx) if ctx.Response.StatusCode() != 200 { - t.Errorf("expected status 200 for allowed access via X-Forwarded-For, got %d", ctx.Response.StatusCode()) + t.Errorf("expected status 200 for allowed access via RemoteAddr, got %d", ctx.Response.StatusCode()) } } diff --git a/internal/stream/stream_test.go b/internal/stream/stream_test.go index b00a83a..cba6ed5 100644 --- a/internal/stream/stream_test.go +++ b/internal/stream/stream_test.go @@ -257,7 +257,7 @@ func TestConcurrentConnections(t *testing.T) { wg.Wait() if s.connCount.Load() != 100 { - t.Errorf("Expected 100 connections, got %d", s.connCount) + t.Errorf("Expected 100 connections, got %d", s.connCount.Load()) } } diff --git a/internal/utils/httperror.go b/internal/utils/httperror.go index 21c97e4..4994a03 100644 --- a/internal/utils/httperror.go +++ b/internal/utils/httperror.go @@ -66,12 +66,13 @@ func SendJSONError(ctx *fasthttp.RequestCtx, status int, errMsg string) { // CheckIPAccess checks whether the client IP is in the allowed list. // If allowed is empty, all access is permitted. -func CheckIPAccess(ctx *fasthttp.RequestCtx, allowed []net.IPNet) bool { +// trustedProxies is used to safely extract the client IP from X-Forwarded-For. +func CheckIPAccess(ctx *fasthttp.RequestCtx, allowed []net.IPNet, trustedProxies []net.IPNet) bool { if len(allowed) == 0 { return true } - clientIP := netutil.ExtractClientIPNet(ctx) + clientIP := netutil.ExtractClientIPWithTrustedProxies(ctx, trustedProxies) if clientIP == nil { return false }