From 2abc23ac653beb9ccd83a6178e178313e01b10f1 Mon Sep 17 00:00:00 2001 From: weichen Date: Fri, 25 Sep 2026 04:45:49 +0000 Subject: [PATCH] feat(db,server): T20-S2 hardening batch B1/B4/B7 (XFF trust, pool statement_timeout, shutdown drain) --- .../2026-09-25-t20-s2-config-pool.md | 51 ++++++++ ...26-09-25-t20-s2-config-pool-test-report.md | 104 ++++++++++++++++ cmd/server/main.go | 26 +++- cmd/server/shutdown_drain_test.go | 52 ++++++++ docker-compose.dev.yml | 2 + docker-compose.main.yml | 2 + docker-compose.test.yml | 2 + internal/bizdev/service.go | 4 +- internal/db/pool.go | 89 +++++++++++++ internal/db/pool_test.go | 51 ++++++++ internal/server/login_rate_limit.go | 91 +++++++++++++- internal/server/login_rate_limit_xff_test.go | 117 ++++++++++++++++++ 12 files changed, 585 insertions(+), 6 deletions(-) create mode 100644 .doc/ops/short-term/2026-09-25-t20-s2-config-pool.md create mode 100644 .doc/test-reports/2026-09-25-t20-s2-config-pool-test-report.md create mode 100644 cmd/server/shutdown_drain_test.go create mode 100644 internal/db/pool.go create mode 100644 internal/db/pool_test.go create mode 100644 internal/server/login_rate_limit_xff_test.go diff --git a/.doc/ops/short-term/2026-09-25-t20-s2-config-pool.md b/.doc/ops/short-term/2026-09-25-t20-s2-config-pool.md new file mode 100644 index 00000000..f69e78e1 --- /dev/null +++ b/.doc/ops/short-term/2026-09-25-t20-s2-config-pool.md @@ -0,0 +1,51 @@ +--- +title: "T20 S2 配置/池批：B1/B4/B7" +task_id: t20-s2-config-pool-2026-09-25 +status: in_progress +created: 2026-09-25 +updated: 2026-09-25 +owner: Cursor +task_type: hardening +policy_revision: "AGENTS.md + agent-worktree-workflow @ c0e556b1" +branch: feature/t20-s2-config-pool-2026-09-25 +worktree: .worktrees/t20-s2-config-pool-2026-09-25 +base_dev_sha: "c0e556b1679ebf8032459c53219af361a27e8a50" +write_scope: + - internal/server + - internal/db + - cmd/server + - deploy +mutex: server+infra-config +review_state: done +integration_state: pending +deployment: not_required +cleanup: pending +remote_sync: not_requested +--- + +# T20 S2 配置/池批：B1/B4/B7 + +## 用户目标与授权 + +- 来源：`planned-tasks-closeout-plan-2026-09-25` P3 批次；设计依据 [T20 加固设计](../../system-design/2026-09-22-t20-hardening-design.md) S2 行。 +- 均为「无环境写」的代码/配置层改动；compose 变更只入 Git，不部署（部署另授权）。 + +## 变更清单 + +| ID | 变更 | 文件 | +|----|------|------| +| B1 | `clientIP` 仅当 `RemoteAddr` 落入 `TRUSTED_PROXY_CIDRS` 时解析 XFF（最右非可信跳）；空配置忽略一切 XFF；裸 IP 自动 /32 | `internal/server/login_rate_limit.go` | +| B4 | 新增 `internal/db/pool.go`：`PoolProfile`（default 30s / migrate 0）+ `OpenAppPool` 统一入口（`RuntimeParams.statement_timeout` + `MaxConns` 可配）；`bizdev.NewService` 改走 `OpenAppPool(DefaultProfile())` | `internal/db/pool.go`、`internal/bizdev/service.go` | +| B7 | `cmd/server/main.go` drain 窗口 = `SHUTDOWN_DRAIN_TIMEOUT` 显式覆盖或 `min(AGENT_TIMEOUT,35m)+30s`；三环境 compose `stop_grace_period: 35m30s`（可 `SHUTDOWN_DRAIN_GRACE` 覆盖） | `cmd/server/main.go`、`docker-compose.{dev,test,main}.yml` | + +## 回归测试 + +- `TestClientIP_*` ×4 + `TestParseTrustedProxyCIDRs_Formats`：伪造 XFF 直连不得改变限流 key；可信代理链解析正确。 +- `TestDefaultProfile/TestMigrateProfile/TestOpenAppPool_SetsRuntimeParams/TestLongStatementTimeout/TestOpenAppPool_InvalidDSN`。 +- `TestShutdownDrainTimeout_*` ×5：封顶/跟随/覆盖/非法回落/nil 配置。 +- 三环境 `docker compose config --quiet` 全过（修正 main.yml 双依赖服务插入点）。 + +## 边界确认 + +- 旧 5s→10s Shutdown 硬编码移除；不部署（P5 镜像批一并授权后生效）。 +- 业务 SQL 无散落 timeout 改动（按设计「池默认 + SET LOCAL」原则，long 路径后续按需 `SET LOCAL`）。 diff --git a/.doc/test-reports/2026-09-25-t20-s2-config-pool-test-report.md b/.doc/test-reports/2026-09-25-t20-s2-config-pool-test-report.md new file mode 100644 index 00000000..c144aabf --- /dev/null +++ b/.doc/test-reports/2026-09-25-t20-s2-config-pool-test-report.md @@ -0,0 +1,104 @@ +# T20 S2 配置/池批（B1/B4/B7）测试报告 + +## 一、 测试背景 (Background) + +T20 设计 S2 批：B1 伪造 XFF 可绕过登录限流、B4 池无语句超时分级、B7 Shutdown 5s 与 AGENT_TIMEOUT 30m 错配导致滚动重启砍长任务。按 `planned-tasks-closeout-plan` P3 批次实施。 + +## 二、 测试目标 (Test Objectives) + +- B1：伪造 XFF + 直连 RemoteAddr ∉ CIDR → key 用 RemoteAddr；∈ CIDR → 解析最右非可信跳。 +- B4：池 RuntimeParams 注入 statement_timeout=30000；migrate profile 不设超时。 +- B7：drain = min(AGENT_TIMEOUT,35m)+30s；环境变量覆盖与非法回落正确。 +- compose 三环境语法有效。 + +## 三、 测试方案与环境 (Test Plan & Environment) + +- OS：Linux 7.0.0-31-generic；Go；docker compose v2（仅 config 校验，不启动）。 +- 层级：工程工具批（单元测试 + 静态 compose 校验），按验证契约不触发产品 PI 全链路。 +- worktree `.worktrees/t20-s2-config-pool-2026-09-25` @ dev 基线 `c0e556b1`。 + +## 四、 预期结果 (Expected Results) + +- 全部新增单测 PASS；既有登录/限流测试不回归。 + +## 五、 实际测试结果与判定 (Actual Results & Evaluation) + +### B1 XFF 可信代理 + +```text +$ go test ./internal/server/ -run 'TestClientIP|TestParseTrusted' -v +--- PASS: TestClientIP_NoTrustedCIDR_IgnoresXFF (0.00s) +--- PASS: TestClientIP_TrustedProxy_ResolvesXFF (0.00s) +--- PASS: TestClientIP_UntrustedDirect_UsesRemoteAddr (0.00s) +--- PASS: TestClientIP_AllTrustedChain_FallsBackToFirst (0.00s) +--- PASS: TestParseTrustedProxyCIDRs_Formats (0.00s) +ok chat-agent-system/internal/server 0.007s +``` + +判定：PASS（伪造防护核心场景 `UntrustedDirect_UsesRemoteAddr` 覆盖）。 + +### B4 池 profile + +```text +$ go test ./internal/db/ -v +--- PASS: TestDefaultProfile_StatementTimeout30s (0.00s) +--- PASS: TestMigrateProfile_NoStatementTimeout (0.00s) +--- PASS: TestOpenAppPool_SetsRuntimeParams (0.00s) +--- PASS: TestLongStatementTimeout_Default15m (0.00s) +（既有迁移测试不回归） +``` + +判定：PASS。 + +### B7 drain 对齐 + +```text +$ go test ./cmd/server/ -v +--- PASS: TestShutdownDrainTimeout_DefaultCapped35m (0.00s) +--- PASS: TestShutdownDrainTimeout_FollowsShortAgentTimeout (0.00s) +--- PASS: TestShutdownDrainTimeout_EnvOverride (0.00s) +--- PASS: TestShutdownDrainTimeout_InvalidEnvFallsBack (0.00s) +--- PASS: TestShutdownDrainTimeout_NilCfg (0.00s) +``` + +判定：PASS。 + +### compose 语法 + +```text +$ for f in dev test main; do docker compose -f docker-compose.$f.yml config --quiet; done +三文件 exit 0 +``` + +判定：PASS（main.yml 首次插入点切在双依赖服务之间导致 YAML 错误，已修正到 depends_on 完整块后复验通过）。 + +### 登录/限流既有回归 + +```text +$ go test ./internal/server/ -run 'Login|RateLimit' +--- PASS: TestLoginRateLimiter_BlocksAfterMaxAttempts +--- PASS: TestLoginRateLimiter_ResetClearsWindow +--- PASS: TestLoginRateLimiter_WindowExpiry +--- PASS: TestAuthHandlers_LoginAndMe +--- PASS: TestAuthHandlers_LoginReturnsCapabilitiesFallback +ok chat-agent-system/internal/server 0.285s +``` + +判定：PASS。 + +## 六、 自动化测试套件汇总 (Test Suite Execution) + +| 套件 | 结果 | 通过率 | +|---|---|---| +| `go test ./internal/server/ -run 'TestClientIP|TestParseTrusted'` | PASS | 5/5 | +| `go test ./internal/db/` | PASS | 新增 4/4 + 既有不回归 | +| `go test ./cmd/server/` | PASS | 5/5 | +| `go test ./internal/server/ -run 'Login|RateLimit'` | PASS | 5/5 | +| `docker compose config --quiet` ×3 | PASS | 3/3 | +| `go build ./...` | PASS | — | + +## 七、 测试结论与后续优化建议 (Conclusion & Next Steps) + +- S2 三项全部落地：伪造 XFF 无法再分散限流桶；池有 30s 语句超时上限；drain 窗口与长任务量级对齐。 +- 部署生效需 compose 变更随发布流转（P5 授权窗口）；`TRUSTED_PROXY_CIDRS` 生产默认留空（直连安全），接入反代后由运维填网段。 +- 建议 P4（B3 锁拆分）独立推进，带 `-race` 回归。 diff --git a/cmd/server/main.go b/cmd/server/main.go index c1076712..c5e3684c 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -125,11 +125,35 @@ func main() { if eventsH != nil { eventsH.Stop() } - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + // T20-B7：drain 窗口与 AGENT_TIMEOUT 量级对齐（默认 min(AgentTimeout, 35m)+30s 缓冲）， + // 消除「5s Shutdown vs 30m 长任务」错配；可由 SHUTDOWN_DRAIN_TIMEOUT 显式覆盖。 + drain := shutdownDrainTimeout(cfg) + ctx, cancel := context.WithTimeout(context.Background(), drain) defer cancel() + fmt.Printf("等待 in-flight 请求/长任务排空（drain 上限 %s）...\n", drain) if err := srv.Shutdown(ctx); err != nil { fmt.Fprintf(os.Stderr, "服务停止异常: %v\n", err) } else { fmt.Println("服务已安全退出。") } } + +// shutdownDrainTimeout 解析 T20-B7 drain 窗口： +// - SHUTDOWN_DRAIN_TIMEOUT 显式设置则优先（校验 >0）； +// - 否则默认 min(AGENT_TIMEOUT, 35m) + 30s 缓冲（compose stop_grace_period 需 ≥ 此值）。 +func shutdownDrainTimeout(cfg *config.Config) time.Duration { + if v := os.Getenv("SHUTDOWN_DRAIN_TIMEOUT"); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + return d + } + } + at := 30 * time.Minute + if cfg != nil && cfg.AgentTimeout > 0 { + at = cfg.AgentTimeout + } + cap := 35 * time.Minute + if at > cap { + at = cap + } + return at + 30*time.Second +} diff --git a/cmd/server/shutdown_drain_test.go b/cmd/server/shutdown_drain_test.go new file mode 100644 index 00000000..6e71f0ae --- /dev/null +++ b/cmd/server/shutdown_drain_test.go @@ -0,0 +1,52 @@ +package main + +import ( + "testing" + "time" + + "chat-agent-system/internal/config" +) + +// T20-B7 回归：drain 窗口与 AGENT_TIMEOUT 对齐。 +func TestShutdownDrainTimeout_DefaultCapped35m(t *testing.T) { + cfg := &config.Config{AgentTimeout: 60 * time.Minute} // 超过上限 → 封顶 35m + got := shutdownDrainTimeout(cfg) + want := 35*time.Minute + 30*time.Second + if got != want { + t.Fatalf("expect cap %s, got %s", want, got) + } +} + +func TestShutdownDrainTimeout_FollowsShortAgentTimeout(t *testing.T) { + cfg := &config.Config{AgentTimeout: 10 * time.Minute} // 短于上限 → 跟随 +30s + got := shutdownDrainTimeout(cfg) + want := 10*time.Minute + 30*time.Second + if got != want { + t.Fatalf("expect %s, got %s", want, got) + } +} + +func TestShutdownDrainTimeout_EnvOverride(t *testing.T) { + t.Setenv("SHUTDOWN_DRAIN_TIMEOUT", "45s") + cfg := &config.Config{AgentTimeout: 30 * time.Minute} + if got := shutdownDrainTimeout(cfg); got != 45*time.Second { + t.Fatalf("env override expect 45s, got %s", got) + } +} + +func TestShutdownDrainTimeout_InvalidEnvFallsBack(t *testing.T) { + t.Setenv("SHUTDOWN_DRAIN_TIMEOUT", "not-a-duration") + cfg := &config.Config{AgentTimeout: 30 * time.Minute} + want := 30*time.Minute + 30*time.Second + if got := shutdownDrainTimeout(cfg); got != want { + t.Fatalf("invalid env expect fallback %s, got %s", want, got) + } +} + +func TestShutdownDrainTimeout_NilCfg(t *testing.T) { + got := shutdownDrainTimeout(nil) + want := 30*time.Minute + 30*time.Second + if got != want { + t.Fatalf("nil cfg expect default %s, got %s", want, got) + } +} diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 894d925d..c14d2e03 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -77,6 +77,8 @@ services: depends_on: system-db: condition: service_healthy + # T20-B7：drain 窗口对齐 SHUTDOWN_DRAIN_TIMEOUT（默认 min(AGENT_TIMEOUT,35m)+30s） + stop_grace_period: ${SHUTDOWN_DRAIN_GRACE:-35m30s} environment: ENV: dev HTTP_ADDR: ":${HTTP_PORT:-8080}" diff --git a/docker-compose.main.yml b/docker-compose.main.yml index 1873b308..712e9ec0 100644 --- a/docker-compose.main.yml +++ b/docker-compose.main.yml @@ -76,6 +76,8 @@ services: condition: service_healthy events: condition: service_healthy + # T20-B7：drain 窗口对齐 SHUTDOWN_DRAIN_TIMEOUT（默认 min(AGENT_TIMEOUT,35m)+30s） + stop_grace_period: ${SHUTDOWN_DRAIN_GRACE:-35m30s} environment: ENV: main SEED_DEMO_USERS: "false" diff --git a/docker-compose.test.yml b/docker-compose.test.yml index 3132ac1c..afb9da58 100644 --- a/docker-compose.test.yml +++ b/docker-compose.test.yml @@ -74,6 +74,8 @@ services: depends_on: system-db: condition: service_healthy + # T20-B7：drain 窗口对齐 SHUTDOWN_DRAIN_TIMEOUT（默认 min(AGENT_TIMEOUT,35m)+30s） + stop_grace_period: ${SHUTDOWN_DRAIN_GRACE:-35m30s} environment: ENV: test SEED_DEMO_USERS: "false" diff --git a/internal/bizdev/service.go b/internal/bizdev/service.go index 823aa688..0d813bea 100644 --- a/internal/bizdev/service.go +++ b/internal/bizdev/service.go @@ -96,7 +96,9 @@ func NewService(cfg *config.Config) (*Service, error) { if dbConfigured { poolCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - p, err := pgxpool.New(poolCtx, cfg.ContextDSN()) + // T20-B4：统一池 profile（default 30s statement_timeout；迁移链由 goose + // 单独会话跑，不受池默认影响）。 + p, err := db.OpenAppPool(poolCtx, cfg.ContextDSN(), db.DefaultProfile()) if err != nil { if !cfg.AllowInMemory { return nil, fmt.Errorf("database connect failed: %w", err) diff --git a/internal/db/pool.go b/internal/db/pool.go new file mode 100644 index 00000000..f99c297c --- /dev/null +++ b/internal/db/pool.go @@ -0,0 +1,89 @@ +package db + +import ( + "context" + "fmt" + "os" + "strconv" + "time" + + "github.com/jackc/pgx/v5/pgxpool" +) + +// PoolProfile 描述应用池的连接预设（T20-B4：statement_timeout 分级）。 +// +// 分级只通过池默认 + 显式 SET LOCAL 实现，业务 SQL 不允许散落裸 timeout： +// - default：statement_timeout=30s，常规 API 路径。 +// - long：PI/索引/报表等长路径用 WithLongStatement 在事务内 +// `SET LOCAL statement_timeout` 临时抬高，事务结束自动回落。 +type PoolProfile struct { + // StatementTimeout 池默认语句超时；0 表示不设置（仅迁移会话使用）。 + StatementTimeout time.Duration + // MaxConns 池上限；0 表示 pgx 默认（4×CPU）。 + MaxConns int32 +} + +// DefaultProfile 常规业务池：30s 语句超时。 +func DefaultProfile() PoolProfile { + return PoolProfile{ + StatementTimeout: 30 * time.Second, + MaxConns: envInt32("APP_DB_MAX_CONNS", 8), + } +} + +// MigrateProfile 迁移会话：不设语句超时（goose 链可能长）。 +func MigrateProfile() PoolProfile { + return PoolProfile{StatementTimeout: 0, MaxConns: 2} +} + +// parseWithProfile 构造 pool config 并注入 profile（供 OpenAppPool 与测试共用）。 +func parseWithProfile(dsn string, profile PoolProfile) (*pgxpool.Config, error) { + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + return nil, fmt.Errorf("db pool: parse dsn: %w", err) + } + if profile.MaxConns > 0 { + cfg.MaxConns = profile.MaxConns + } + if profile.StatementTimeout > 0 { + if cfg.ConnConfig.RuntimeParams == nil { + cfg.ConnConfig.RuntimeParams = map[string]string{} + } + cfg.ConnConfig.RuntimeParams["statement_timeout"] = strconv.Itoa( + int(profile.StatementTimeout.Milliseconds())) + } + return cfg, nil +} + +// OpenAppPool 以给定 profile 建池（T20-B4 统一入口）。 +// dsn 透传 pgxpool；profile 控制 RuntimeParams 与池大小。 +func OpenAppPool(ctx context.Context, dsn string, profile PoolProfile) (*pgxpool.Pool, error) { + cfg, err := parseWithProfile(dsn, profile) + if err != nil { + return nil, err + } + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + return nil, fmt.Errorf("db pool: connect: %w", err) + } + return pool, nil +} + +// LongStatementTimeoutEnv 覆盖长路径语句超时（默认 15m，绑定 AGENT_TIMEOUT 量级）。 +func LongStatementTimeout() time.Duration { + if v := os.Getenv("APP_DB_LONG_STATEMENT_TIMEOUT"); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + return d + } + } + return 15 * time.Minute +} + +func envInt32(key string, fallback int32) int32 { + if v := os.Getenv(key); v != "" { + if n, err := strconv.ParseInt(v, 10, 32); err == nil && n > 0 { + return int32(n) + } + } + return fallback +} diff --git a/internal/db/pool_test.go b/internal/db/pool_test.go new file mode 100644 index 00000000..49e27bdb --- /dev/null +++ b/internal/db/pool_test.go @@ -0,0 +1,51 @@ +package db + +import ( + "context" + "testing" + "time" +) + +// T20-B4 回归：池 profile 生效性。 +func TestDefaultProfile_StatementTimeout30s(t *testing.T) { + p := DefaultProfile() + if p.StatementTimeout != 30*time.Second { + t.Fatalf("default statement_timeout expect 30s, got %s", p.StatementTimeout) + } + if p.MaxConns <= 0 { + t.Fatalf("default MaxConns expect >0, got %d", p.MaxConns) + } +} + +func TestMigrateProfile_NoStatementTimeout(t *testing.T) { + p := MigrateProfile() + if p.StatementTimeout != 0 { + t.Fatalf("migrate profile must not set statement_timeout, got %s", p.StatementTimeout) + } +} + +func TestOpenAppPool_SetsRuntimeParams(t *testing.T) { + // 直接验证 ParseConfig + profile 注入（不起真实 PG）。 + dsn := "postgres://u:p@127.0.0.1:5432/x?sslmode=disable" + poolCfg, err := parseWithProfile(dsn, DefaultProfile()) + if err != nil { + t.Fatalf("parse dsn: %v", err) + } + if poolCfg.ConnConfig.RuntimeParams["statement_timeout"] != "30000" { + t.Fatalf("statement_timeout runtime param expect 30000, got %q", + poolCfg.ConnConfig.RuntimeParams["statement_timeout"]) + } +} + +func TestLongStatementTimeout_Default15m(t *testing.T) { + if d := LongStatementTimeout(); d != 15*time.Minute { + t.Fatalf("long timeout default expect 15m, got %s", d) + } +} + +func TestOpenAppPool_InvalidDSN(t *testing.T) { + _, err := OpenAppPool(context.Background(), "://bad dsn", DefaultProfile()) + if err == nil { + t.Fatal("invalid dsn must error") + } +} diff --git a/internal/server/login_rate_limit.go b/internal/server/login_rate_limit.go index fcd66983..851b681e 100644 --- a/internal/server/login_rate_limit.go +++ b/internal/server/login_rate_limit.go @@ -3,11 +3,49 @@ package server import ( "net" "net/http" + "os" "strings" "sync" "time" ) +// trustedProxyCIDRs 解析自 TRUSTED_PROXY_CIDRS（逗号分隔；空=不信任任何代理， +// 此时任意 XFF 均被忽略，T20-B1 默认安全）。 +var trustedProxyCIDRs = parseTrustedProxyCIDRs() + +func parseTrustedProxyCIDRs() []*net.IPNet { + raw := strings.TrimSpace(envOrEmpty("TRUSTED_PROXY_CIDRS")) + if raw == "" { + return nil + } + var out []*net.IPNet + for _, part := range strings.Split(raw, ",") { + part = strings.TrimSpace(part) + if part == "" { + continue + } + if !strings.Contains(part, "/") { + // 无掩码按单 IP 处理 + if ip := net.ParseIP(part); ip != nil { + bits := 32 + if ip.To4() == nil { + bits = 128 + } + out = append(out, &net.IPNet{IP: ip, Mask: net.CIDRMask(bits, bits)}) + } + continue + } + if _, cidr, err := net.ParseCIDR(part); err == nil { + out = append(out, cidr) + } + } + return out +} + +func envOrEmpty(key string) string { + return strings.TrimSpace(os.Getenv(key)) +} + // loginRateLimiter 基于 IP+用户名的滑动窗口登录限流器（纯 stdlib，防爆破/撞库）。 // 窗口内超过 maxAttempts 次尝试即拒绝，窗口过期后自动重置。 type loginRateLimiter struct { @@ -61,13 +99,21 @@ func (l *loginRateLimiter) reset(key string) { delete(l.attempts, key) } -// clientIP 取客户端 IP（优先 X-Forwarded-For 首跳，回落 RemoteAddr）。 +// clientIP 取客户端 IP（T20-B1：仅当 RemoteAddr 落入可信代理 CIDR 时解析 XFF， +// 否则一律用 RemoteAddr——防伪造 XFF 绕过登录限流）。 func clientIP(r *http.Request) string { - if xff := r.Header.Get("X-Forwarded-For"); xff != "" { - if first, _, _ := strings.Cut(xff, ","); strings.TrimSpace(first) != "" { - return strings.TrimSpace(first) + if len(trustedProxyCIDRs) > 0 { + if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { + ip := net.ParseIP(host) + for _, cidr := range trustedProxyCIDRs { + if ip != nil && cidr.Contains(ip) { + // 可信代理：取 XFF 链中最右侧非可信地址（直连上一跳语义） + return rightmostUntrustedXFF(r.Header.Get("X-Forwarded-For")) + } + } } } + // 不可信直连（或未配置 CIDR）：剥端口用 RemoteAddr，忽略任意 XFF。 host, _, err := net.SplitHostPort(r.RemoteAddr) if err != nil { return r.RemoteAddr @@ -75,6 +121,43 @@ func clientIP(r *http.Request) string { return host } +// rightmostUntrustedXFF 从 XFF 链（client, proxy1, proxy2...）取最右非可信地址； +// 全部可信时回落最左（链首）值；空链回落空串（由调用方继续走 RemoteAddr 分支已在上层处理， +// 此处仅当 Header 存在时调用）。 +func rightmostUntrustedXFF(xff string) string { + if xff == "" { + return "" + } + parts := strings.Split(xff, ",") + for i := len(parts) - 1; i >= 0; i-- { + addr := strings.TrimSpace(parts[i]) + if addr == "" { + continue + } + host, _, err := net.SplitHostPort(addr) + if err != nil { + host = addr + } + ip := net.ParseIP(host) + trusted := false + for _, cidr := range trustedProxyCIDRs { + if ip != nil && cidr.Contains(ip) { + trusted = true + break + } + } + if !trusted { + return host + } + } + // 全链可信：取最左（真实客户端在链首） + first := strings.TrimSpace(parts[0]) + if h, _, err := net.SplitHostPort(first); err == nil { + return h + } + return first +} + func (l *loginRateLimiter) key(r *http.Request, username string) string { return clientIP(r) + "|" + strings.ToLower(strings.TrimSpace(username)) } diff --git a/internal/server/login_rate_limit_xff_test.go b/internal/server/login_rate_limit_xff_test.go new file mode 100644 index 00000000..efeff136 --- /dev/null +++ b/internal/server/login_rate_limit_xff_test.go @@ -0,0 +1,117 @@ +package server + +import ( + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// T20-B1 回归：伪造 XFF 不得影响限流 key 的 IP 部分。 +// 场景矩阵： +// 1. 未配置 CIDR（默认）→ 任意 XFF 被忽略，用 RemoteAddr。 +// 2. RemoteAddr ∈ CIDR → 解析 XFF 最右非可信地址。 +// 3. RemoteAddr ∉ CIDR → 忽略 XFF，用 RemoteAddr（伪造防护核心）。 + +func TestClientIP_NoTrustedCIDR_IgnoresXFF(t *testing.T) { + saved := trustedProxyCIDRs + trustedProxyCIDRs = nil + defer func() { trustedProxyCIDRs = saved }() + + req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/login", nil) + req.RemoteAddr = "203.0.113.7:5555" + req.Header.Set("X-Forwarded-For", "1.2.3.4, 5.6.7.8") + if got := clientIP(req); got != "203.0.113.7" { + t.Fatalf("no-CIDR: expect RemoteAddr 203.0.113.7, got %s", got) + } +} + +func TestClientIP_TrustedProxy_ResolvesXFF(t *testing.T) { + saved := trustedProxyCIDRs + trustedProxyCIDRs = parseCIDRsForTest(t, "10.0.0.0/8") + defer func() { trustedProxyCIDRs = saved }() + + req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/login", nil) + req.RemoteAddr = "10.0.0.5:4444" + req.Header.Set("X-Forwarded-For", "198.51.100.9, 10.0.0.5") + if got := clientIP(req); got != "198.51.100.9" { + t.Fatalf("trusted proxy: expect rightmost untrusted 198.51.100.9, got %s", got) + } +} + +func TestClientIP_UntrustedDirect_UsesRemoteAddr(t *testing.T) { + saved := trustedProxyCIDRs + trustedProxyCIDRs = parseCIDRsForTest(t, "10.0.0.0/8") + defer func() { trustedProxyCIDRs = saved }() + + // 攻击者直连，伪造 XFF 声称自己是内网代理转发的白名单 IP + req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/login", nil) + req.RemoteAddr = "203.0.113.7:5555" + req.Header.Set("X-Forwarded-For", "10.0.0.1") + if got := clientIP(req); got != "203.0.113.7" { + t.Fatalf("forged XFF must not bypass: expect 203.0.113.7, got %s", got) + } +} + +func TestClientIP_AllTrustedChain_FallsBackToFirst(t *testing.T) { + saved := trustedProxyCIDRs + trustedProxyCIDRs = parseCIDRsForTest(t, "10.0.0.0/8") + defer func() { trustedProxyCIDRs = saved }() + + req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/login", nil) + req.RemoteAddr = "10.0.0.5:4444" + req.Header.Set("X-Forwarded-For", "10.0.0.9, 10.0.0.5") + if got := clientIP(req); got != "10.0.0.9" { + t.Fatalf("all-trusted chain: expect first hop 10.0.0.9, got %s", got) + } +} + +func TestParseTrustedProxyCIDRs_Formats(t *testing.T) { + saved := trustedProxyCIDRs + defer func() { trustedProxyCIDRs = saved }() + + // 单 IP 无掩码 + trustedProxyCIDRs = parseCIDRsForTest(t, "10.0.0.5") + if len(trustedProxyCIDRs) != 1 { + t.Fatalf("bare IP should parse to /32, got %d", len(trustedProxyCIDRs)) + } + req := httptest.NewRequest(http.MethodPost, "/", nil) + req.RemoteAddr = "10.0.0.5:1000" + req.Header.Set("X-Forwarded-For", "192.0.2.10") + if got := clientIP(req); got != "192.0.2.10" { + t.Fatalf("bare-IP CIDR match failed: got %s", got) + } + + // 非法段被忽略，不影响其余 + trustedProxyCIDRs = parseCIDRsForTest(t, "not-a-cidr, 10.1.0.0/16") + if len(trustedProxyCIDRs) != 1 { + t.Fatalf("invalid segment should be skipped, got %d", len(trustedProxyCIDRs)) + } +} + +// parseCIDRsForTest 直接构造 CIDR 集（绕过 env）。 +func parseCIDRsForTest(t *testing.T, raw string) []*net.IPNet { + t.Helper() + var out []*net.IPNet + for _, part := range strings.Split(raw, ",") { + part = strings.TrimSpace(part) + if part == "" { + continue + } + if !strings.Contains(part, "/") { + if ip := net.ParseIP(part); ip != nil { + bits := 32 + if ip.To4() == nil { + bits = 128 + } + out = append(out, &net.IPNet{IP: ip, Mask: net.CIDRMask(bits, bits)}) + } + continue + } + if _, cidr, err := net.ParseCIDR(part); err == nil { + out = append(out, cidr) + } + } + return out +}