From 9e6e7ca3c7084b37b7c1459222ee38d4b4d068dd Mon Sep 17 00:00:00 2001 From: m1saka Date: Sat, 11 Jul 2026 17:23:32 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BB=A3=E7=90=86=E4=B8=8E?= =?UTF-8?q?=E6=97=A5=E5=BF=97=E9=93=BE=E8=B7=AF=E7=9A=84=E5=8F=AF=E9=9D=A0?= =?UTF-8?q?=E6=80=A7=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 2 + README.md | 5 +- compose.yml | 1 + config/config.go | 13 ++- db/clickhouse.go | 27 ++++++- db/clickhouse_test.go | 20 +++++ .../reports/reliability-security-fixes.md | 8 +- logger/queue.go | 64 ++++++++++++--- main.go | 13 ++- main_test.go | 15 ++-- proxy/capture.go | 13 ++- proxy/proxy.go | 70 ++++++++++------ proxy/sse.go | 80 +++++++++++++++---- proxy/sse_event.go | 19 ++++- proxy/sse_event_test.go | 12 +++ proxy/timeout.go | 18 +++++ proxy/writer.go | 9 ++- tests/config/config_test.go | 22 +++++ tests/logger/queue_test.go | 5 +- tests/proxy/robustness_test.go | 32 ++++++++ tests/proxy/sse_capture_test.go | 52 ++++++++++++ 21 files changed, 417 insertions(+), 83 deletions(-) diff --git a/.env.example b/.env.example index f06044a..a55a270 100644 --- a/.env.example +++ b/.env.example @@ -11,6 +11,8 @@ CLICKHOUSE_URL= # 单条请求/响应 body 最大记录字节数 MAX_BODY_BYTES=1048576 +# 单个代理请求体硬上限，超出时返回 413 +MAX_REQUEST_BYTES=16777216 # 异步日志队列容量 LOG_QUEUE_SIZE=256 diff --git a/README.md b/README.md index f5a4e76..35c204d 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,7 @@ go vet ./... | `UPSTREAM_TLS_INSECURE_SKIP_VERIFY` | 跳过上游 HTTPS 证书校验；仅开发/可信内网自签证书场景使用 | `false` | | `CLICKHOUSE_URL` | ClickHouse 原生协议地址（必填）；支持严格 `host:port`、`clickhouse://` 或 TLS `clickhouses://` | - | | `MAX_BODY_BYTES` | 单个请求体和响应体的记录上限 | `1048576` | +| `MAX_REQUEST_BYTES` | 单个代理请求体硬上限，超出时返回 413 | `16777216` | | `LOG_QUEUE_SIZE` | 异步队列条数上限 | `256` | | `LOG_QUEUE_BYTES` | 异步队列总字节预算 | `67108864` | | `LOG_BATCH_SIZE` | 批量写入条数 | `50` | @@ -164,7 +165,7 @@ ORDER BY (started_at, request_id); | `client_ip` | `String` | 否 | 默认记录 TCP 对端；仅直接对端命中 `TRUSTED_PROXIES` 时，从右向左解析 `X-Forwarded-For`，无 XFF 时使用有效的 `X-Real-IP` | | `request_headers` | `String` | 否 | 完整请求头，序列化为 `{"Header-Name": ["value1", "value2"], ...}` 的 JSON 对象；**注意 Authorization、Cookie、API-Key 等敏感头未脱敏**，按设计原样存储 | | `request_body` | `String` | 否 | 请求体内容。最多保留 `MAX_BODY_BYTES`（默认 1 MiB）字节，超出部分丢弃；读取失败时请求不会转发 | -| `request_truncated` | `Bool` | 否 | 请求体是否被 `MAX_BODY_BYTES` 截断。截断只影响数据库记录；代理会将已完整读入内存的请求体通过 `bytes.Reader` 重放给下游 newapi | +| `request_truncated` | `Bool` | 否 | 请求体是否被 `MAX_BODY_BYTES` 截断。截断只影响数据库记录；请求体超过 `MAX_REQUEST_BYTES` 时不会转发并返回 413 | | `status_code` | `Int32` | 否 | 上游返回的 HTTP 状态码。`502` 通常意味着上游连接失败；WebSocket 成功升级记录为 `101` | | `response_headers` | `String` | 否 | 响应头，结构同 `request_headers`。对于 SSE，会包含 `Content-Type: text/event-stream` 等 | | `response_body` | `String` | 否 | 响应体内容。对于未截断的 SSE，代理会尝试将 OpenAI Completions、Chat Completions、Responses、Anthropic 或 Gemini 事件组装为单个 JSON，并以该 JSON 替换原始 SSE 字节流；无法识别或组装失败时保留原始 SSE。超过 `MAX_BODY_BYTES` 时不组装，仅保留原始字节流的前 `MAX_BODY_BYTES` 字节 | @@ -219,7 +220,7 @@ SELECT JSONExtractRaw(request_headers, 'Authorization') FROM proxy_logs LIMIT 5; ## 设计要点 -- 转发前使用 `io.ReadAll` 将请求体完整读入内存，读取或关闭失败时不向上游发送请求；数据库仅记录前 `MAX_BODY_BYTES` 字节，随后通过 `bytes.Reader` 重放完整请求体。因此 `MAX_BODY_BYTES` 只限制日志字段大小，不限制代理读取请求体时的内存占用。 +- 转发前最多读取 `MAX_REQUEST_BYTES + 1` 字节，超限返回 413，读取或关闭失败时不向上游发送请求；数据库仅记录前 `MAX_BODY_BYTES` 字节，随后通过 `bytes.Reader` 重放完整请求体。 - `httputil.ReverseProxy` + `FlushInterval = -1`，自定义 `ResponseWriter` 同时实现 `Flusher`/`Hijacker`，写入时先转发再缓冲，保证流式实时性。 - 日志通过非阻塞 channel 投递，队列满或 DB 不健康时直接丢弃（每 30 秒打印 metrics）。 - DB 健康状态机：写入失败立即标记 unhealthy，后台 ping 恢复后重新启用。 diff --git a/compose.yml b/compose.yml index df26843..1aa7a54 100644 --- a/compose.yml +++ b/compose.yml @@ -17,6 +17,7 @@ services: UPSTREAM_TLS_INSECURE_SKIP_VERIFY: "${UPSTREAM_TLS_INSECURE_SKIP_VERIFY:-false}" CLICKHOUSE_URL: "${CLICKHOUSE_URL:?set CLICKHOUSE_URL with URL-encoded credentials}" MAX_BODY_BYTES: "${MAX_BODY_BYTES:-1048576}" + MAX_REQUEST_BYTES: "${MAX_REQUEST_BYTES:-16777216}" LOG_QUEUE_SIZE: "${LOG_QUEUE_SIZE:-256}" LOG_QUEUE_BYTES: "${LOG_QUEUE_BYTES:-67108864}" LOG_BATCH_SIZE: "${LOG_BATCH_SIZE:-50}" diff --git a/config/config.go b/config/config.go index cab6219..072e034 100644 --- a/config/config.go +++ b/config/config.go @@ -3,6 +3,7 @@ package config import ( "errors" "fmt" + "math" "net/netip" "net/url" "os" @@ -19,6 +20,7 @@ type Config struct { UpstreamURL *url.URL ClickHouseURL string MaxBodyBytes int64 + MaxRequestBytes int64 LogQueueSize int LogQueueBytes int64 LogBatchSize int @@ -42,6 +44,10 @@ func Load() (*Config, error) { if err != nil { return nil, err } + maxRequestBytes, err := getEnvInt64("MAX_REQUEST_BYTES", 16<<20) + if err != nil { + return nil, err + } logQueueSize, err := getEnvInt("LOG_QUEUE_SIZE", 256) if err != nil { return nil, err @@ -103,6 +109,7 @@ func Load() (*Config, error) { ListenAddr: getEnv("LISTEN_ADDR", ":8080"), ClickHouseURL: os.Getenv("CLICKHOUSE_URL"), MaxBodyBytes: maxBodyBytes, + MaxRequestBytes: maxRequestBytes, LogQueueSize: logQueueSize, LogQueueBytes: logQueueBytes, LogBatchSize: logBatchSize, @@ -128,7 +135,7 @@ func Load() (*Config, error) { if err != nil { return nil, fmt.Errorf("invalid UPSTREAM_URL: %w", err) } - if u.Scheme == "" || u.Host == "" { + if (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" { return nil, fmt.Errorf("invalid UPSTREAM_URL: %q", upstream) } cfg.UpstreamURL = u @@ -144,6 +151,7 @@ func Load() (*Config, error) { value int64 }{ {key: "MAX_BODY_BYTES", value: cfg.MaxBodyBytes}, + {key: "MAX_REQUEST_BYTES", value: cfg.MaxRequestBytes}, {key: "LOG_QUEUE_SIZE", value: int64(cfg.LogQueueSize)}, {key: "LOG_QUEUE_BYTES", value: cfg.LogQueueBytes}, {key: "LOG_BATCH_SIZE", value: int64(cfg.LogBatchSize)}, @@ -162,6 +170,9 @@ func Load() (*Config, error) { return nil, fmt.Errorf("%s must be greater than zero", item.key) } } + if cfg.MaxRequestBytes == math.MaxInt64 { + return nil, errors.New("MAX_REQUEST_BYTES is too large") + } return cfg, nil } diff --git a/db/clickhouse.go b/db/clickhouse.go index a16d5c3..e49e8df 100644 --- a/db/clickhouse.go +++ b/db/clickhouse.go @@ -22,6 +22,7 @@ type Pool struct { open func(*clickhouse.Options) (clickhouse.Conn, error) mu sync.RWMutex conn clickhouse.Conn + generation uint64 healthy bool closed bool } @@ -82,6 +83,7 @@ func (p *Pool) connect(ctx context.Context) error { _ = p.conn.Close() } p.conn = conn + p.generation++ p.healthy = true p.mu.Unlock() log.Printf("[db] connected and migrated") @@ -192,11 +194,11 @@ func (p *Pool) watch(ctx context.Context) { return case <-t.C: if p.Healthy() { - if conn := p.Get(); conn != nil { + if conn, generation := p.GetWithGeneration(); conn != nil { pctx, cancel := context.WithTimeout(ctx, 3*time.Second) if err := conn.Ping(pctx); err != nil { log.Printf("[db] ping failed, marking unhealthy: %v", err) - p.MarkUnhealthy() + p.MarkUnhealthyGeneration(generation) } cancel() } @@ -223,6 +225,21 @@ func (p *Pool) MarkUnhealthy() { p.mu.Unlock() } +// MarkUnhealthyGeneration only changes health when the failed connection is still current. +func (p *Pool) MarkUnhealthyGeneration(generation uint64) { + p.mu.Lock() + if p.generation == generation { + p.healthy = false + } + p.mu.Unlock() +} + +func (p *Pool) Generation() uint64 { + p.mu.RLock() + defer p.mu.RUnlock() + return p.generation +} + // Get 返回当前连接，可能为 nil。 func (p *Pool) Get() clickhouse.Conn { p.mu.RLock() @@ -230,6 +247,12 @@ func (p *Pool) Get() clickhouse.Conn { return p.conn } +func (p *Pool) GetWithGeneration() (clickhouse.Conn, uint64) { + p.mu.RLock() + defer p.mu.RUnlock() + return p.conn, p.generation +} + // Close starts closing the underlying connection and waits within ctx. func (p *Pool) Close(ctx context.Context) error { p.mu.Lock() diff --git a/db/clickhouse_test.go b/db/clickhouse_test.go index 89496ca..92f73f3 100644 --- a/db/clickhouse_test.go +++ b/db/clickhouse_test.go @@ -214,6 +214,26 @@ func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) { } } +func TestOldGenerationCannotMarkReplacementUnhealthy(t *testing.T) { + pool := &Pool{conn: newFakeClickHouseConn(), healthy: true, generation: 1} + oldGeneration := pool.Generation() + + pool.mu.Lock() + pool.conn = newFakeClickHouseConn() + pool.generation++ + pool.healthy = true + pool.mu.Unlock() + + pool.MarkUnhealthyGeneration(oldGeneration) + if !pool.Healthy() { + t.Fatal("old connection failure marked replacement connection unhealthy") + } + pool.MarkUnhealthyGeneration(pool.Generation()) + if pool.Healthy() { + t.Fatal("current connection failure did not mark pool unhealthy") + } +} + func assertNonDestructiveMigration(t *testing.T, statements []string) { t.Helper() if len(statements) == 0 || !strings.Contains(strings.ToUpper(statements[0]), "CREATE TABLE IF NOT EXISTS") { diff --git a/docs/compose/reports/reliability-security-fixes.md b/docs/compose/reports/reliability-security-fixes.md index 39ffde8..fbf012d 100644 --- a/docs/compose/reports/reliability-security-fixes.md +++ b/docs/compose/reports/reliability-security-fixes.md @@ -1,6 +1,6 @@ --- feature: reliability-security-fixes -status: delivered +status: complete specs: - docs/compose/specs/reliability-security-fixes.md - docs/compose/specs/2026-07-09-clickhouse-migration.md @@ -14,7 +14,7 @@ branch: main ## What Was Built -本轮完成了 token_thief 的可靠性与安全加固。异步日志队列现在安全处理并发 `Submit`/`Stop`、使用条目数和字节双预算、在统一 shutdown deadline 内排空，并区分 ClickHouse 的可安全重试、确定失败和提交结果不明三类写入结果。 +本轮完成了 token_thief 的主要可靠性与安全加固。异步日志队列现在安全处理并发 `Submit`/`Stop`、使用条目数和字节双预算、在统一 shutdown deadline 内排空，并区分 ClickHouse 的可安全重试、确定失败和提交结果不明三类写入结果。 反向代理现在对请求体读取、上游响应复制、普通响应超时、SSE idle timeout、WebSocket 101 元数据和升级连接关闭实施完整性保护。客户端错误文本不再泄露内部信息，转发来源头仅在直接对端属于 `TRUSTED_PROXIES` 时参与客户端 IP 判定。 @@ -37,7 +37,7 @@ branch: main ## Verification -最终验收从工作区根目录运行 `gofmt -l .`、`go test -count=1 ./...`、`go vet ./...` 和 `go build ./`。报告保留可复现命令而不固化会随测试增删失真的测试或包计数；针对 config、queue、ClickHouse、HTTP/SSE/WebSocket 和 Compose 的失败路径均有测试覆盖。 +最终验收从工作区根目录运行 `gofmt -w .`、`go vet ./...`、`go build ./...` 和 `go test -count=1 -timeout 3m ./...`，命令均以 0 退出。另对代理测试执行 20 轮重复验证，对 logger 队列测试执行 20 轮重复验证，均通过。WebSocket hijack 后的普通 HTTP flush 已由 writer 层 guard 阻断，未再出现 server recover panic。 ## Journey Log @@ -47,6 +47,8 @@ branch: main - [lesson] 迭代 1：进程内队列只能提供有界 best-effort；跨进程幂等需要持久化 outbox 和下游幂等协议，不能由本地重试可靠模拟。 - [pivot] 迭代 1：SSE 完成判定限定为真正的 `text/event-stream` 且等待代理正常返回，避免终止标记导致提前记录不完整响应。 - [lesson] 迭代 4：最终报告记录禁用测试缓存的验证命令而非易过期的精确计数，并同时执行格式、静态检查和构建。 +- [finding] 迭代 5：格式、vet、构建和测试均退出成功，但 passing test 中仍可隐藏由 `net/http` recover 的 handler panic；WebSocket hijack 后不得再刷新普通 HTTP response writer。 +- [fix] 最终验收：增加 `MAX_REQUEST_BYTES` 硬上限、上游连接写 deadline、Queue deadline 后有界清理，以及 hijack 后 flush guard；代理和队列压力测试各连续 20 轮通过。 ## Source Materials diff --git a/logger/queue.go b/logger/queue.go index 7d22e6d..e801230 100644 --- a/logger/queue.go +++ b/logger/queue.go @@ -8,6 +8,7 @@ import ( "sync" "sync/atomic" "time" + "unsafe" "github.com/ClickHouse/clickhouse-go/v2/lib/driver" @@ -46,11 +47,16 @@ func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string) } type poolBackend struct { - pool *db.Pool + pool *db.Pool + conn driver.Conn + generation uint64 } func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) { - conn := b.pool.Get() + conn := b.conn + if conn == nil { + conn, _ = b.pool.GetWithGeneration() + } if conn == nil { return nil, errors.New("pool nil") } @@ -58,7 +64,14 @@ func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, err } func (b poolBackend) Healthy() bool { return b.pool.Healthy() } -func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthy() } +func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthyGeneration(b.generation) } +func (b poolBackend) Snapshot() (Backend, uint64) { + conn, generation := b.pool.GetWithGeneration() + return poolBackend{pool: b.pool, conn: conn, generation: generation}, generation +} +func (b poolBackend) MarkUnhealthyGeneration(generation uint64) { + b.pool.MarkUnhealthyGeneration(generation) +} type Stats struct { Enqueued uint64 @@ -153,8 +166,8 @@ func EstimatedBytes(e *LogEntry) int64 { if e == nil { return 0 } - return int64(len(e.RequestID) + len(e.Method) + len(e.Path) + len(e.Query) + len(e.ClientIP) + len(e.Error) + - len(e.RequestHeaders) + len(e.RequestBody) + len(e.ResponseHeaders) + len(e.ResponseBody)) + return int64(unsafe.Sizeof(*e)+unsafe.Sizeof(queuedEntry{})) + int64(len(e.RequestID)+len(e.Method)+len(e.Path)+len(e.Query)+len(e.ClientIP)+len(e.Error)+ + len(e.RequestHeaders)+len(e.RequestBody)+len(e.ResponseHeaders)+len(e.ResponseBody)) } func (q *Queue) Stats() Stats { @@ -195,7 +208,6 @@ func (q *Queue) Stop(contexts ...context.Context) error { if len(contexts) > 0 && contexts[0] != nil { ctx = contexts[0] } - q.mu.Lock() shutdownOwner := false if !q.stopped { @@ -226,7 +238,12 @@ func (q *Queue) Stop(contexts ...context.Context) error { case <-ctx.Done(): if shutdownOwner && workCancel != nil { workCancel() - <-done + timer := time.NewTimer(100 * time.Millisecond) + defer timer.Stop() + select { + case <-done: + case <-timer.C: + } } return ctx.Err() } @@ -238,8 +255,7 @@ func (q *Queue) Submit(e *LogEntry) { q.dropped.Add(1) return } - entry := cloneLogEntry(e) - size := EstimatedBytes(entry) + size := EstimatedBytes(e) if !q.mu.TryRLock() { q.dropped.Add(1) return @@ -249,6 +265,7 @@ func (q *Queue) Submit(e *LogEntry) { q.dropped.Add(1) return } + entry := cloneLogEntry(e) select { case q.ch <- queuedEntry{entry: entry, size: size}: q.enq.Add(1) @@ -367,16 +384,19 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) { retry := entries var lastErr error + var failedGeneration uint64 for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ { if err := ctx.Err(); err != nil { lastErr = err break } - result := Flush(ctx, q.backend, retry) + attemptBackend, generation := backendSnapshot(q.backend) + failedGeneration = generation + result := Flush(ctx, attemptBackend, retry) q.failed.Add(uint64(result.Failed)) q.ambiguous.Add(uint64(result.Ambiguous)) if result.Ambiguous > 0 { - q.backend.MarkUnhealthy() + markBackendUnhealthy(q.backend, failedGeneration) } lastErr = result.Err retry = result.Retry @@ -391,11 +411,31 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) { if len(retry) > 0 { q.failed.Add(uint64(len(retry))) - q.backend.MarkUnhealthy() + markBackendUnhealthy(q.backend, failedGeneration) log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr) } } +type generationBackend interface { + Snapshot() (Backend, uint64) + MarkUnhealthyGeneration(uint64) +} + +func backendSnapshot(backend Backend) (Backend, uint64) { + if versioned, ok := backend.(generationBackend); ok { + return versioned.Snapshot() + } + return backend, 0 +} + +func markBackendUnhealthy(backend Backend, generation uint64) { + if versioned, ok := backend.(generationBackend); ok { + versioned.MarkUnhealthyGeneration(generation) + return + } + backend.MarkUnhealthy() +} + func waitBackoff(ctx context.Context, attempt int) bool { wait := BaseBackoff for i := 1; i < attempt; i++ { diff --git a/main.go b/main.go index 9961a60..ab77fb4 100644 --- a/main.go +++ b/main.go @@ -30,7 +30,7 @@ type shutdownStep struct { run func(context.Context) error } -func waitForShutdown(stop <-chan os.Signal, serverErr <-chan error) error { +func waitForShutdown(stop <-chan struct{}, serverErr <-chan error) error { select { case <-stop: log.Printf("[main] shutdown signal received") @@ -83,11 +83,9 @@ func run() error { } log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns)) - stop := make(chan os.Signal, 1) - signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM) - defer signal.Stop(stop) - - rootCtx, cancel := context.WithCancel(context.Background()) + rootCtx, stopSignals := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stopSignals() + rootCtx, cancel := context.WithCancel(rootCtx) defer cancel() pool := db.NewPool(rootCtx, cfg.ClickHouseURL, cfg.DBReconnectInterval) @@ -101,6 +99,7 @@ func run() error { SSEIdleTimeout: cfg.UpstreamStreamIdleTimeout, UpstreamTLSInsecureSkipVerify: cfg.UpstreamTLSInsecureSkipVerify, TrustedProxies: cfg.TrustedProxies, + MaxRequestBytes: cfg.MaxRequestBytes, }) srv := &http.Server{ @@ -118,7 +117,7 @@ func run() error { serverErr <- srv.ListenAndServe() }() - runErr := waitForShutdown(stop, serverErr) + runErr := waitForShutdown(rootCtx.Done(), serverErr) // HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。 shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) diff --git a/main_test.go b/main_test.go index 4fc13bb..3c249c8 100644 --- a/main_test.go +++ b/main_test.go @@ -7,7 +7,6 @@ import ( "go/parser" "go/token" "net/http" - "os" "reflect" "strings" "testing" @@ -49,13 +48,13 @@ func TestRunRegistersSignalsBeforeStartingResources(t *testing.T) { } name := owner.Name + "." + selector.Sel.Name switch name { - case "signal.Notify", "db.NewPool", "queue.Start", "srv.ListenAndServe": + case "signal.NotifyContext", "db.NewPool", "queue.Start", "srv.ListenAndServe": positions[name] = call.Pos() } return true }) - notifyPos, ok := positions["signal.Notify"] + notifyPos, ok := positions["signal.NotifyContext"] if !ok { t.Fatal("run does not register for shutdown signals") } @@ -65,7 +64,7 @@ func TestRunRegistersSignalsBeforeStartingResources(t *testing.T) { t.Fatalf("run does not call %s", start) } if notifyPos >= startPos { - t.Errorf("signal.Notify at line %d must precede %s at line %d", + t.Errorf("signal.NotifyContext at line %d must precede %s at line %d", fset.Position(notifyPos).Line, start, fset.Position(startPos).Line) } } @@ -76,7 +75,7 @@ func TestWaitForShutdownReturnsListenerError(t *testing.T) { serverErr := make(chan error, 1) serverErr <- listenErr - err := waitForShutdown(make(chan os.Signal), serverErr) + err := waitForShutdown(make(chan struct{}), serverErr) if !errors.Is(err, listenErr) || !strings.Contains(err.Error(), "server") { t.Fatalf("waitForShutdown() error = %v, want wrapped listener error", err) } @@ -84,8 +83,8 @@ func TestWaitForShutdownReturnsListenerError(t *testing.T) { func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) { t.Run("signal", func(t *testing.T) { - stop := make(chan os.Signal, 1) - stop <- os.Interrupt + stop := make(chan struct{}) + close(stop) if err := waitForShutdown(stop, make(chan error)); err != nil { t.Fatalf("waitForShutdown() error = %v, want nil", err) } @@ -94,7 +93,7 @@ func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) { t.Run("server closed", func(t *testing.T) { serverErr := make(chan error, 1) serverErr <- http.ErrServerClosed - if err := waitForShutdown(make(chan os.Signal), serverErr); err != nil { + if err := waitForShutdown(make(chan struct{}), serverErr); err != nil { t.Fatalf("waitForShutdown() error = %v, want nil", err) } }) diff --git a/proxy/capture.go b/proxy/capture.go index af1b937..3e7acb0 100644 --- a/proxy/capture.go +++ b/proxy/capture.go @@ -5,21 +5,30 @@ import ( "crypto/rand" "encoding/hex" "encoding/json" + "errors" "io" "net/http" "net/netip" "strings" ) +var errRequestBodyTooLarge = errors.New("request body too large") + // readRequestBody 在转发前完整读取请求体，确保读取失败时不会向上游发送损坏请求。 -func readRequestBody(r *http.Request, max int64) (captured []byte, truncated bool, err error) { +func readRequestBody(r *http.Request, max, requestLimit int64) (captured []byte, truncated bool, err error) { if r.Body == nil || r.Body == http.NoBody { return nil, false, nil } - body, err := io.ReadAll(r.Body) + if r.ContentLength > requestLimit { + return nil, false, errRequestBodyTooLarge + } + body, err := io.ReadAll(io.LimitReader(r.Body, requestLimit+1)) if err != nil { return nil, false, err } + if int64(len(body)) > requestLimit { + return nil, false, errRequestBodyTooLarge + } if err := r.Body.Close(); err != nil { return nil, false, err } diff --git a/proxy/proxy.go b/proxy/proxy.go index 40f1b67..05f75aa 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -3,6 +3,7 @@ package proxy import ( "context" "crypto/tls" + "errors" "log" "net" "net/http" @@ -25,16 +26,17 @@ type LogSubmitter interface { // Handler 构造反代 HTTP handler。 type Handler struct { - rp *httputil.ReverseProxy - filter *config.Filter - queue LogSubmitter - maxBodyBytes int64 - trusted []netip.Prefix - connMu sync.Mutex - conns map[net.Conn]struct{} - connChanged chan struct{} - closing bool - closed bool + rp *httputil.ReverseProxy + filter *config.Filter + queue LogSubmitter + maxBodyBytes int64 + maxRequestBytes int64 + trusted []netip.Prefix + connMu sync.Mutex + conns map[net.Conn]struct{} + connChanged chan struct{} + closing bool + closed bool } // Options 控制反代连接上游时的网络行为。 @@ -44,6 +46,7 @@ type Options struct { SSEIdleTimeout time.Duration UpstreamTLSInsecureSkipVerify bool TrustedProxies []netip.Prefix + MaxRequestBytes int64 } // requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。 @@ -77,12 +80,22 @@ func New(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody i } func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody int64, opts Options) *Handler { + if opts.MaxRequestBytes <= 0 { + opts.MaxRequestBytes = 16 << 20 + } rp := httputil.NewSingleHostReverseProxy(upstream) rp.FlushInterval = -1 // 让流式 chunk 立即转发 if opts.UpstreamTimeout > 0 || opts.UpstreamTLSInsecureSkipVerify { transport := http.DefaultTransport.(*http.Transport).Clone() if opts.UpstreamTimeout > 0 { - transport.DialContext = (&net.Dialer{Timeout: opts.UpstreamTimeout, KeepAlive: 30 * time.Second}).DialContext + dialer := &net.Dialer{Timeout: opts.UpstreamTimeout, KeepAlive: 30 * time.Second} + transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) { + conn, err := dialer.DialContext(ctx, network, address) + if err != nil { + return nil, err + } + return &writeTimeoutConn{Conn: conn, timeout: opts.UpstreamTimeout}, nil + } transport.ResponseHeaderTimeout = opts.UpstreamTimeout transport.TLSHandshakeTimeout = opts.UpstreamTimeout } @@ -130,13 +143,14 @@ func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter } return &Handler{ - rp: rp, - filter: filter, - queue: queue, - maxBodyBytes: maxBody, - trusted: append([]netip.Prefix(nil), opts.TrustedProxies...), - conns: make(map[net.Conn]struct{}), - connChanged: make(chan struct{}), + rp: rp, + filter: filter, + queue: queue, + maxBodyBytes: maxBody, + maxRequestBytes: opts.MaxRequestBytes, + trusted: append([]netip.Prefix(nil), opts.TrustedProxies...), + conns: make(map[net.Conn]struct{}), + connChanged: make(chan struct{}), } } @@ -149,11 +163,15 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } shouldLog := h.filter.ShouldLog(r.URL.Path) - reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes) + reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes, h.maxRequestBytes) if err != nil && !shouldLog { log.Printf("[proxy] read request body failed: %v", err) _ = r.Body.Close() - http.Error(w, "bad request", http.StatusBadRequest) + if errors.Is(err, errRequestBodyTooLarge) { + http.Error(w, "request body too large", http.StatusRequestEntityTooLarge) + } else { + http.Error(w, "bad request", http.StatusBadRequest) + } return } if !shouldLog { @@ -176,7 +194,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { _ = r.Body.Close() s := "read request body: " + err.Error() st.lastErr.Store(&s) - h.serveRequestBodyError(cw, r, st, started, reqID) + h.serveRequestBodyError(cw, r, st, started, reqID, errors.Is(err, errRequestBodyTooLarge)) return } @@ -230,9 +248,15 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } } -func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string) { +func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string, tooLarge bool) { cw.Header().Set("X-Request-Id", requestID) - http.Error(cw, "bad request", http.StatusBadRequest) + status := http.StatusBadRequest + message := "bad request" + if tooLarge { + status = http.StatusRequestEntityTooLarge + message = "request body too large" + } + http.Error(cw, message, status) if !cw.Complete() { return } diff --git a/proxy/sse.go b/proxy/sse.go index 5f416ff..f9bd92d 100644 --- a/proxy/sse.go +++ b/proxy/sse.go @@ -552,11 +552,16 @@ func (g *geminiChunk) UnmarshalJSON(data []byte) error { } type geminiCandidate struct { - raw map[string]json.RawMessage - index int - role string - text strings.Builder - partRaw map[string]json.RawMessage + raw map[string]json.RawMessage + index int + role string + parts []*geminiPart +} + +type geminiPart struct { + raw map[string]json.RawMessage + text strings.Builder + kind string } func assembleGeminiSSE(payloads []string) ([]byte, bool) { @@ -588,15 +593,31 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) { assembled.role = candidate.Content.Role } if len(candidate.Content.Parts) > 0 { - for _, part := range candidate.Content.Parts { - assembled.text.WriteString(part.Text) + var contentRaw struct { + Parts []map[string]json.RawMessage `json:"parts"` } - if len(assembled.partRaw) == 0 { - var contentRaw struct { - Parts []map[string]json.RawMessage `json:"parts"` + if rawContent, ok := candidate.Raw["content"]; ok { + _ = json.Unmarshal(rawContent, &contentRaw) + } + mergeByIndex := len(assembled.parts) == len(candidate.Content.Parts) + if mergeByIndex { + for i := range candidate.Content.Parts { + if i >= len(contentRaw.Parts) || assembled.parts[i].kind != geminiPartKind(contentRaw.Parts[i]) { + mergeByIndex = false + break + } } - if rawContent, ok := candidate.Raw["content"]; ok && json.Unmarshal(rawContent, &contentRaw) == nil && len(contentRaw.Parts) > 0 { - assembled.partRaw = cloneRawMap(contentRaw.Parts[0]) + } + for i, part := range candidate.Content.Parts { + target := i + if !mergeByIndex { + target = len(assembled.parts) + assembled.parts = append(assembled.parts, &geminiPart{}) + } + assembled.parts[target].text.WriteString(part.Text) + if i < len(contentRaw.Parts) { + assembled.parts[target].kind = geminiPartKind(contentRaw.Parts[i]) + assembled.parts[target].raw = mergeRawMap(assembled.parts[target].raw, contentRaw.Parts[i]) } } } @@ -614,12 +635,15 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) { candidate := candidates[index] candidateRaw := cloneRawMap(candidate.raw) candidateRaw["index"] = mustJSON(candidate.index) - contentRaw := map[string]any{"role": candidate.role, "parts": []any{map[string]any{"text": candidate.text.String()}}} - if len(candidate.partRaw) > 0 { - partRaw := cloneRawMap(candidate.partRaw) - partRaw["text"] = mustJSON(candidate.text.String()) - contentRaw["parts"] = []any{rawMapToMap(partRaw)} + parts := make([]any, 0, len(candidate.parts)) + for _, part := range candidate.parts { + partRaw := cloneRawMap(part.raw) + if part.text.Len() > 0 || len(partRaw) == 0 { + partRaw["text"] = mustJSON(part.text.String()) + } + parts = append(parts, rawMapToMap(partRaw)) } + contentRaw := map[string]any{"role": candidate.role, "parts": parts} candidateRaw["content"] = mustJSON(contentRaw) assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw)) } @@ -631,6 +655,28 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) { return data, err == nil } +func geminiPartKind(part map[string]json.RawMessage) string { + for _, key := range []string{"functionCall", "functionResponse", "inlineData", "fileData", "executableCode", "codeExecutionResult", "text"} { + if _, ok := part[key]; ok { + return key + } + } + return "unknown" +} + +func mergeRawMap(dst, src map[string]json.RawMessage) map[string]json.RawMessage { + if dst == nil { + dst = make(map[string]json.RawMessage, len(src)) + } + for key, value := range src { + if key == "text" { + continue + } + dst[key] = append(json.RawMessage(nil), value...) + } + return dst +} + func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage { out := make(map[string]json.RawMessage, len(in)) for k, v := range in { diff --git a/proxy/sse_event.go b/proxy/sse_event.go index f572905..c3da2cc 100644 --- a/proxy/sse_event.go +++ b/proxy/sse_event.go @@ -8,12 +8,29 @@ import ( type sseEventTracker struct { buf []byte + limit int + overflow bool recognized bool terminal bool afterTerminal bool } +func newSSEEventTracker(limit int64) sseEventTracker { + if limit > int64(^uint(0)>>1) { + limit = int64(^uint(0) >> 1) + } + return sseEventTracker{limit: int(limit)} +} + func (t *sseEventTracker) Write(p []byte) { + if t.overflow { + return + } + if t.limit > 0 && len(p) > t.limit-len(t.buf) { + t.buf = nil + t.overflow = true + return + } t.buf = append(t.buf, p...) for { end, separator := completeSSEEvent(t.buf) @@ -34,7 +51,7 @@ func (t *sseEventTracker) Write(p []byte) { } func (t *sseEventTracker) Complete() bool { - return len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal) + return !t.overflow && len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal) } func completeSSEEvent(buf []byte) (int, int) { diff --git a/proxy/sse_event_test.go b/proxy/sse_event_test.go index ea96855..e5ed922 100644 --- a/proxy/sse_event_test.go +++ b/proxy/sse_event_test.go @@ -19,3 +19,15 @@ func TestSSEEventTrackerAcceptsCROnlyEventBoundary(t *testing.T) { t.Fatal("terminal SSE event with CR-only boundary should be complete") } } + +func TestSSEEventTrackerStopsBufferingOverLimit(t *testing.T) { + tracker := newSSEEventTracker(8) + tracker.Write([]byte("data: 123456789")) + + if tracker.Complete() { + t.Fatal("overflowed SSE tracker must not report a complete stream") + } + if len(tracker.buf) != 0 || !tracker.overflow { + t.Fatalf("overflow state=%v buffered=%d, want overflow with released buffer", tracker.overflow, len(tracker.buf)) + } +} diff --git a/proxy/timeout.go b/proxy/timeout.go index a4d9769..2cc9653 100644 --- a/proxy/timeout.go +++ b/proxy/timeout.go @@ -2,12 +2,30 @@ package proxy import ( "io" + "net" "net/http" "sync" "sync/atomic" "time" ) +type writeTimeoutConn struct { + net.Conn + timeout time.Duration +} + +func (c *writeTimeoutConn) Write(p []byte) (int, error) { + if err := c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil { + return 0, err + } + n, err := c.Conn.Write(p) + clearErr := c.Conn.SetWriteDeadline(time.Time{}) + if err == nil { + err = clearErr + } + return n, err +} + type trackingBody struct { io.ReadCloser failed *atomic.Bool diff --git a/proxy/writer.go b/proxy/writer.go index 6ee60fd..003d3df 100644 --- a/proxy/writer.go +++ b/proxy/writer.go @@ -26,7 +26,7 @@ type captureWriter struct { } func newCaptureWriter(w http.ResponseWriter, max int64) *captureWriter { - return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK} + return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK, sse: newSSEEventTracker(max)} } func (c *captureWriter) WriteHeader(code int) { @@ -63,7 +63,7 @@ func (c *captureWriter) Write(p []byte) (int, error) { c.truncated = true } c.written += int64(n) - if isStreamResponse(c.Header()) { + if c.max > 0 && isStreamResponse(c.Header()) { c.sse.Write(p[:n]) } c.flush() @@ -78,7 +78,7 @@ func (c *captureWriter) Flush() { // ConfirmDelivery establishes an observable delivery boundary for responses // whose headers were not followed by a body write. func (c *captureWriter) ConfirmDelivery() { - if c.written != 0 { + if c.hijacked || c.written != 0 { return } if _, ok := c.ResponseWriter.(http.Flusher); !ok { @@ -89,6 +89,9 @@ func (c *captureWriter) ConfirmDelivery() { } func (c *captureWriter) flush() { + if c.hijacked { + return + } if _, ok := c.ResponseWriter.(http.Flusher); !ok { return } diff --git a/tests/config/config_test.go b/tests/config/config_test.go index 21379a2..7e9e73e 100644 --- a/tests/config/config_test.go +++ b/tests/config/config_test.go @@ -1,7 +1,9 @@ package config_test import ( + "math" "net/netip" + "strconv" "strings" "testing" "time" @@ -50,6 +52,9 @@ func TestLoadQueueDefaults(t *testing.T) { if cfg.MaxBodyBytes != 1<<20 { t.Fatalf("MaxBodyBytes=%d want %d", cfg.MaxBodyBytes, 1<<20) } + if cfg.MaxRequestBytes != 16<<20 { + t.Fatalf("MaxRequestBytes=%d want %d", cfg.MaxRequestBytes, 16<<20) + } if cfg.LogQueueSize != 256 { t.Fatalf("LogQueueSize=%d want 256", cfg.LogQueueSize) } @@ -242,6 +247,23 @@ func TestLoadRequiresClickHouseURL(t *testing.T) { } } +func TestLoadRejectsUnsupportedUpstreamScheme(t *testing.T) { + t.Setenv("UPSTREAM_URL", "ftp://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + if _, err := config.Load(); err == nil { + t.Fatal("Load succeeded with unsupported UPSTREAM_URL scheme") + } +} + +func TestLoadRejectsMaxRequestBytesOverflowBoundary(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + t.Setenv("MAX_REQUEST_BYTES", strconv.FormatInt(math.MaxInt64, 10)) + if _, err := config.Load(); err == nil || !strings.Contains(err.Error(), "MAX_REQUEST_BYTES") { + t.Fatalf("Load error=%v, want MAX_REQUEST_BYTES rejection", err) + } +} + func TestLoadValidatesClickHouseURL(t *testing.T) { tests := []struct { name string diff --git a/tests/logger/queue_test.go b/tests/logger/queue_test.go index 2195340..99e0fa6 100644 --- a/tests/logger/queue_test.go +++ b/tests/logger/queue_test.go @@ -244,8 +244,8 @@ func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) { RequestHeaders: []byte("7777777"), RequestBody: []byte("88888888"), ResponseHeaders: []byte("999999999"), ResponseBody: []byte("0000000000"), } - if got, want := logger.EstimatedBytes(entry), int64(55); got != want { - t.Fatalf("EstimatedBytes=%d want %d", got, want) + if got := logger.EstimatedBytes(entry); got <= 55 { + t.Fatalf("EstimatedBytes=%d must include fixed entry overhead", got) } } @@ -485,6 +485,7 @@ func TestConcurrentStopCannotCancelFirstStopDrain(t *testing.T) { firstCtx, cancelFirst := context.WithTimeout(context.Background(), time.Second) defer cancelFirst() go func() { firstResult <- q.Stop(firstCtx) }() + time.Sleep(10 * time.Millisecond) for q.Stats().Dropped == 0 { q.Submit(&logger.LogEntry{RequestID: "stop-probe"}) } diff --git a/tests/proxy/robustness_test.go b/tests/proxy/robustness_test.go index 2d2e579..1158181 100644 --- a/tests/proxy/robustness_test.go +++ b/tests/proxy/robustness_test.go @@ -108,6 +108,38 @@ func TestRequestBodyReadFailureReturnsFixed400WithoutUpstream(t *testing.T) { } } +func TestRequestBodyOverLimitReturns413WithoutUpstream(t *testing.T) { + var upstreamCalls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{MaxRequestBytes: 4}) + + for _, tc := range []struct { + name string + contentLength int64 + }{ + {name: "known length", contentLength: 5}, + {name: "chunked", contentLength: -1}, + } { + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", strings.NewReader("12345")) + req.ContentLength = tc.contentLength + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusRequestEntityTooLarge || rec.Body.String() != "request body too large\n" { + t.Fatalf("response=(%d, %q), want fixed 413", rec.Code, rec.Body.String()) + } + }) + } + if upstreamCalls.Load() != 0 { + t.Fatalf("upstream called %d times", upstreamCalls.Load()) + } +} + func TestRequestBodyReadFailureWithZeroContentLengthDoesNotReachUpstream(t *testing.T) { var upstreamCalls atomic.Int32 upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/tests/proxy/sse_capture_test.go b/tests/proxy/sse_capture_test.go index 51d7227..4cf87b8 100644 --- a/tests/proxy/sse_capture_test.go +++ b/tests/proxy/sse_capture_test.go @@ -643,6 +643,58 @@ func TestGeminiStreamPreservesFunctionCallParts(t *testing.T) { } } +func TestGeminiStreamPreservesMultiplePartsByIndex(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"candidates":[{"content":{"parts":[{"text":"hello "},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"world"}],"role":"model"},"index":0}]}` + "\n\n", + `data: {"candidates":[{"content":{"parts":[{"text":"again"},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"!"}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent") + + var captured struct { + Candidates []struct { + Content struct { + Parts []map[string]any `json:"parts"` + } `json:"content"` + } `json:"candidates"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody) + } + parts := captured.Candidates[0].Content.Parts + if len(parts) != 3 || parts[0]["text"] != "hello again" || parts[2]["text"] != "world!" { + t.Fatalf("multiple Gemini parts not preserved: %+v", parts) + } + if _, ok := parts[1]["functionCall"]; !ok { + t.Fatalf("middle functionCall part missing: %+v", parts) + } +} + +func TestGeminiStreamAppendsDifferentPartKindsAcrossChunks(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"candidates":[{"content":{"parts":[{"text":"answer"}],"role":"model"},"index":0}]}` + "\n\n", + `data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent") + + var captured struct { + Candidates []struct { + Content struct { + Parts []map[string]any `json:"parts"` + } `json:"content"` + } `json:"candidates"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody) + } + parts := captured.Candidates[0].Content.Parts + if len(parts) != 2 || parts[0]["text"] != "answer" { + t.Fatalf("different Gemini part kinds were merged: %+v", parts) + } + if _, ok := parts[1]["functionCall"]; !ok { + t.Fatalf("functionCall part missing: %+v", parts) + } +} + func TestUnknownSSEKeepsRawBody(t *testing.T) { e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions")