修复代理与日志链路的可靠性问题
This commit is contained in:
@@ -11,6 +11,8 @@ CLICKHOUSE_URL=
|
||||
|
||||
# 单条请求/响应 body 最大记录字节数
|
||||
MAX_BODY_BYTES=1048576
|
||||
# 单个代理请求体硬上限,超出时返回 413
|
||||
MAX_REQUEST_BYTES=16777216
|
||||
|
||||
# 异步日志队列容量
|
||||
LOG_QUEUE_SIZE=256
|
||||
|
||||
@@ -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 恢复后重新启用。
|
||||
|
||||
@@ -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}"
|
||||
|
||||
+12
-1
@@ -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
|
||||
}
|
||||
|
||||
|
||||
+25
-2
@@ -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()
|
||||
|
||||
@@ -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") {
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+52
-12
@@ -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++ {
|
||||
|
||||
@@ -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)
|
||||
|
||||
+7
-8
@@ -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)
|
||||
}
|
||||
})
|
||||
|
||||
+11
-2
@@ -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
|
||||
}
|
||||
|
||||
+47
-23
@@ -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
|
||||
}
|
||||
|
||||
+63
-17
@@ -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 {
|
||||
|
||||
+18
-1
@@ -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) {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
+6
-3
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"})
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user