Compare commits

...
10 Commits
22 changed files with 824 additions and 100 deletions
+2
View File
@@ -11,6 +11,8 @@ CLICKHOUSE_URL=
# 单条请求/响应 body 最大记录字节数 # 单条请求/响应 body 最大记录字节数
MAX_BODY_BYTES=1048576 MAX_BODY_BYTES=1048576
# 单个代理请求体硬上限,超出时返回 413
MAX_REQUEST_BYTES=16777216
# 异步日志队列容量 # 异步日志队列容量
LOG_QUEUE_SIZE=256 LOG_QUEUE_SIZE=256
+3 -2
View File
@@ -79,6 +79,7 @@ go vet ./...
| `UPSTREAM_TLS_INSECURE_SKIP_VERIFY` | 跳过上游 HTTPS 证书校验;仅开发/可信内网自签证书场景使用 | `false` | | `UPSTREAM_TLS_INSECURE_SKIP_VERIFY` | 跳过上游 HTTPS 证书校验;仅开发/可信内网自签证书场景使用 | `false` |
| `CLICKHOUSE_URL` | ClickHouse 原生协议地址(必填);支持严格 `host:port``clickhouse://` 或 TLS `clickhouses://` | - | | `CLICKHOUSE_URL` | ClickHouse 原生协议地址(必填);支持严格 `host:port``clickhouse://` 或 TLS `clickhouses://` | - |
| `MAX_BODY_BYTES` | 单个请求体和响应体的记录上限 | `1048576` | | `MAX_BODY_BYTES` | 单个请求体和响应体的记录上限 | `1048576` |
| `MAX_REQUEST_BYTES` | 单个代理请求体硬上限,超出时返回 413 | `16777216` |
| `LOG_QUEUE_SIZE` | 异步队列条数上限 | `256` | | `LOG_QUEUE_SIZE` | 异步队列条数上限 | `256` |
| `LOG_QUEUE_BYTES` | 异步队列总字节预算 | `67108864` | | `LOG_QUEUE_BYTES` | 异步队列总字节预算 | `67108864` |
| `LOG_BATCH_SIZE` | 批量写入条数 | `50` | | `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` | | `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_headers` | `String` | 否 | 完整请求头,序列化为 `{"Header-Name": ["value1", "value2"], ...}` 的 JSON 对象;**注意 Authorization、Cookie、API-Key 等敏感头未脱敏**,按设计原样存储 |
| `request_body` | `String` | 否 | 请求体内容。最多保留 `MAX_BODY_BYTES`(默认 1 MiB)字节,超出部分丢弃;读取失败时请求不会转发 | | `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` | | `status_code` | `Int32` | 否 | 上游返回的 HTTP 状态码。`502` 通常意味着上游连接失败;WebSocket 成功升级记录为 `101` |
| `response_headers` | `String` | 否 | 响应头,结构同 `request_headers`。对于 SSE,会包含 `Content-Type: text/event-stream` 等 | | `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` 字节 | | `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`,写入时先转发再缓冲,保证流式实时性。 - `httputil.ReverseProxy` + `FlushInterval = -1`,自定义 `ResponseWriter` 同时实现 `Flusher`/`Hijacker`,写入时先转发再缓冲,保证流式实时性。
- 日志通过非阻塞 channel 投递,队列满或 DB 不健康时直接丢弃(每 30 秒打印 metrics)。 - 日志通过非阻塞 channel 投递,队列满或 DB 不健康时直接丢弃(每 30 秒打印 metrics)。
- DB 健康状态机:写入失败立即标记 unhealthy,后台 ping 恢复后重新启用。 - DB 健康状态机:写入失败立即标记 unhealthy,后台 ping 恢复后重新启用。
+1
View File
@@ -17,6 +17,7 @@ services:
UPSTREAM_TLS_INSECURE_SKIP_VERIFY: "${UPSTREAM_TLS_INSECURE_SKIP_VERIFY:-false}" UPSTREAM_TLS_INSECURE_SKIP_VERIFY: "${UPSTREAM_TLS_INSECURE_SKIP_VERIFY:-false}"
CLICKHOUSE_URL: "${CLICKHOUSE_URL:?set CLICKHOUSE_URL with URL-encoded credentials}" CLICKHOUSE_URL: "${CLICKHOUSE_URL:?set CLICKHOUSE_URL with URL-encoded credentials}"
MAX_BODY_BYTES: "${MAX_BODY_BYTES:-1048576}" MAX_BODY_BYTES: "${MAX_BODY_BYTES:-1048576}"
MAX_REQUEST_BYTES: "${MAX_REQUEST_BYTES:-16777216}"
LOG_QUEUE_SIZE: "${LOG_QUEUE_SIZE:-256}" LOG_QUEUE_SIZE: "${LOG_QUEUE_SIZE:-256}"
LOG_QUEUE_BYTES: "${LOG_QUEUE_BYTES:-67108864}" LOG_QUEUE_BYTES: "${LOG_QUEUE_BYTES:-67108864}"
LOG_BATCH_SIZE: "${LOG_BATCH_SIZE:-50}" LOG_BATCH_SIZE: "${LOG_BATCH_SIZE:-50}"
+12 -1
View File
@@ -3,6 +3,7 @@ package config
import ( import (
"errors" "errors"
"fmt" "fmt"
"math"
"net/netip" "net/netip"
"net/url" "net/url"
"os" "os"
@@ -19,6 +20,7 @@ type Config struct {
UpstreamURL *url.URL UpstreamURL *url.URL
ClickHouseURL string ClickHouseURL string
MaxBodyBytes int64 MaxBodyBytes int64
MaxRequestBytes int64
LogQueueSize int LogQueueSize int
LogQueueBytes int64 LogQueueBytes int64
LogBatchSize int LogBatchSize int
@@ -42,6 +44,10 @@ func Load() (*Config, error) {
if err != nil { if err != nil {
return nil, err return nil, err
} }
maxRequestBytes, err := getEnvInt64("MAX_REQUEST_BYTES", 16<<20)
if err != nil {
return nil, err
}
logQueueSize, err := getEnvInt("LOG_QUEUE_SIZE", 256) logQueueSize, err := getEnvInt("LOG_QUEUE_SIZE", 256)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -103,6 +109,7 @@ func Load() (*Config, error) {
ListenAddr: getEnv("LISTEN_ADDR", ":8080"), ListenAddr: getEnv("LISTEN_ADDR", ":8080"),
ClickHouseURL: os.Getenv("CLICKHOUSE_URL"), ClickHouseURL: os.Getenv("CLICKHOUSE_URL"),
MaxBodyBytes: maxBodyBytes, MaxBodyBytes: maxBodyBytes,
MaxRequestBytes: maxRequestBytes,
LogQueueSize: logQueueSize, LogQueueSize: logQueueSize,
LogQueueBytes: logQueueBytes, LogQueueBytes: logQueueBytes,
LogBatchSize: logBatchSize, LogBatchSize: logBatchSize,
@@ -128,7 +135,7 @@ func Load() (*Config, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("invalid UPSTREAM_URL: %w", err) 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) return nil, fmt.Errorf("invalid UPSTREAM_URL: %q", upstream)
} }
cfg.UpstreamURL = u cfg.UpstreamURL = u
@@ -144,6 +151,7 @@ func Load() (*Config, error) {
value int64 value int64
}{ }{
{key: "MAX_BODY_BYTES", value: cfg.MaxBodyBytes}, {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_SIZE", value: int64(cfg.LogQueueSize)},
{key: "LOG_QUEUE_BYTES", value: cfg.LogQueueBytes}, {key: "LOG_QUEUE_BYTES", value: cfg.LogQueueBytes},
{key: "LOG_BATCH_SIZE", value: int64(cfg.LogBatchSize)}, {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) 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 return cfg, nil
} }
+42 -2
View File
@@ -2,6 +2,7 @@ package db
import ( import (
"context" "context"
"errors"
"fmt" "fmt"
"log" "log"
"net" "net"
@@ -21,9 +22,13 @@ type Pool struct {
open func(*clickhouse.Options) (clickhouse.Conn, error) open func(*clickhouse.Options) (clickhouse.Conn, error)
mu sync.RWMutex mu sync.RWMutex
conn clickhouse.Conn conn clickhouse.Conn
generation uint64
healthy bool healthy bool
closed bool
} }
var errPoolClosed = errors.New("db pool is closed")
// NewPool 创建 Pool。即使首次连接失败也返回非 nil 实例,后台会持续重试。 // NewPool 创建 Pool。即使首次连接失败也返回非 nil 实例,后台会持续重试。
func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *Pool { func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *Pool {
p := &Pool{dsn: dsn, reconnectInterval: reconnectInterval} p := &Pool{dsn: dsn, reconnectInterval: reconnectInterval}
@@ -35,6 +40,13 @@ func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *
} }
func (p *Pool) connect(ctx context.Context) error { func (p *Pool) connect(ctx context.Context) error {
p.mu.RLock()
closed := p.closed
p.mu.RUnlock()
if closed {
return errPoolClosed
}
cctx, cancel := context.WithTimeout(ctx, 5*time.Second) cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel() defer cancel()
opts, err := ClickHouseOptions(p.dsn) opts, err := ClickHouseOptions(p.dsn)
@@ -62,10 +74,16 @@ func (p *Pool) connect(ctx context.Context) error {
return err return err
} }
p.mu.Lock() p.mu.Lock()
if p.closed {
p.mu.Unlock()
_ = conn.Close()
return errPoolClosed
}
if p.conn != nil { if p.conn != nil {
_ = p.conn.Close() _ = p.conn.Close()
} }
p.conn = conn p.conn = conn
p.generation++
p.healthy = true p.healthy = true
p.mu.Unlock() p.mu.Unlock()
log.Printf("[db] connected and migrated") log.Printf("[db] connected and migrated")
@@ -176,11 +194,11 @@ func (p *Pool) watch(ctx context.Context) {
return return
case <-t.C: case <-t.C:
if p.Healthy() { if p.Healthy() {
if conn := p.Get(); conn != nil { if conn, generation := p.GetWithGeneration(); conn != nil {
pctx, cancel := context.WithTimeout(ctx, 3*time.Second) pctx, cancel := context.WithTimeout(ctx, 3*time.Second)
if err := conn.Ping(pctx); err != nil { if err := conn.Ping(pctx); err != nil {
log.Printf("[db] ping failed, marking unhealthy: %v", err) log.Printf("[db] ping failed, marking unhealthy: %v", err)
p.MarkUnhealthy() p.MarkUnhealthyGeneration(generation)
} }
cancel() cancel()
} }
@@ -207,6 +225,21 @@ func (p *Pool) MarkUnhealthy() {
p.mu.Unlock() 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。 // Get 返回当前连接,可能为 nil。
func (p *Pool) Get() clickhouse.Conn { func (p *Pool) Get() clickhouse.Conn {
p.mu.RLock() p.mu.RLock()
@@ -214,9 +247,16 @@ func (p *Pool) Get() clickhouse.Conn {
return p.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. // Close starts closing the underlying connection and waits within ctx.
func (p *Pool) Close(ctx context.Context) error { func (p *Pool) Close(ctx context.Context) error {
p.mu.Lock() p.mu.Lock()
p.closed = true
conn := p.conn conn := p.conn
p.conn = nil p.conn = nil
p.healthy = false p.healthy = false
+65 -1
View File
@@ -183,6 +183,57 @@ func TestInitialConnectWithIncompatibleSchemaStaysUnhealthy(t *testing.T) {
assertNonDestructiveMigration(t, candidate.statements) assertNonDestructiveMigration(t, candidate.statements)
} }
func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) {
old := newFakeClickHouseConn()
pingStarted := make(chan struct{})
pingRelease := make(chan struct{})
candidate := newFakeClickHouseConn().withBlockedPing(pingStarted, pingRelease)
pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true}
pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
connectDone := make(chan error, 1)
go func() { connectDone <- pool.connect(context.Background()) }()
<-pingStarted
if err := pool.Close(context.Background()); err != nil {
t.Fatalf("Close: %v", err)
}
close(pingRelease)
if err := <-connectDone; err == nil {
t.Fatal("connect succeeded after pool was closed")
}
if pool.Get() != nil || pool.Healthy() {
t.Fatalf("closed pool connection = %T, healthy = %v; want nil, false", pool.Get(), pool.Healthy())
}
if old.closeCalls != 1 {
t.Fatalf("old connection close calls = %d, want 1", old.closeCalls)
}
if candidate.closeCalls != 1 {
t.Fatalf("candidate close calls = %d, want 1", candidate.closeCalls)
}
}
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) { func assertNonDestructiveMigration(t *testing.T, statements []string) {
t.Helper() t.Helper()
if len(statements) == 0 || !strings.Contains(strings.ToUpper(statements[0]), "CREATE TABLE IF NOT EXISTS") { if len(statements) == 0 || !strings.Contains(strings.ToUpper(statements[0]), "CREATE TABLE IF NOT EXISTS") {
@@ -203,6 +254,8 @@ type fakeClickHouseConn struct {
queryErr error queryErr error
statements []string statements []string
closeCalls int closeCalls int
pingStarted chan struct{}
pingRelease chan struct{}
} }
func newFakeClickHouseConn() *fakeClickHouseConn { func newFakeClickHouseConn() *fakeClickHouseConn {
@@ -229,6 +282,11 @@ func (c *fakeClickHouseConn) withColumnType(index int, typ string) *fakeClickHou
c.rows.values[index][1] = typ c.rows.values[index][1] = typ
return c return c
} }
func (c *fakeClickHouseConn) withBlockedPing(started, release chan struct{}) *fakeClickHouseConn {
c.pingStarted = started
c.pingRelease = release
return c
}
func (c *fakeClickHouseConn) Contributors() []string { return nil } func (c *fakeClickHouseConn) Contributors() []string { return nil }
func (c *fakeClickHouseConn) ServerVersion() (*driver.ServerVersion, error) { return nil, nil } func (c *fakeClickHouseConn) ServerVersion() (*driver.ServerVersion, error) { return nil, nil }
@@ -249,7 +307,13 @@ func (c *fakeClickHouseConn) Exec(_ context.Context, query string, _ ...any) err
return c.execErr return c.execErr
} }
func (c *fakeClickHouseConn) AsyncInsert(context.Context, string, bool, ...any) error { return nil } func (c *fakeClickHouseConn) AsyncInsert(context.Context, string, bool, ...any) error { return nil }
func (c *fakeClickHouseConn) Ping(context.Context) error { return nil } func (c *fakeClickHouseConn) Ping(context.Context) error {
if c.pingStarted != nil {
close(c.pingStarted)
<-c.pingRelease
}
return nil
}
func (c *fakeClickHouseConn) Stats() driver.Stats { return driver.Stats{} } func (c *fakeClickHouseConn) Stats() driver.Stats { return driver.Stats{} }
func (c *fakeClickHouseConn) Close() error { c.closeCalls++; return nil } func (c *fakeClickHouseConn) Close() error { c.closeCalls++; return nil }
@@ -1,6 +1,6 @@
--- ---
feature: reliability-security-fixes feature: reliability-security-fixes
status: delivered status: complete
specs: specs:
- docs/compose/specs/reliability-security-fixes.md - docs/compose/specs/reliability-security-fixes.md
- docs/compose/specs/2026-07-09-clickhouse-migration.md - docs/compose/specs/2026-07-09-clickhouse-migration.md
@@ -14,7 +14,7 @@ branch: main
## What Was Built ## 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 判定。 反向代理现在对请求体读取、上游响应复制、普通响应超时、SSE idle timeout、WebSocket 101 元数据和升级连接关闭实施完整性保护。客户端错误文本不再泄露内部信息,转发来源头仅在直接对端属于 `TRUSTED_PROXIES` 时参与客户端 IP 判定。
@@ -37,7 +37,7 @@ branch: main
## Verification ## Verification
迭代 1 验证全部通过:`gofmt -l .` 无输出,`go test -json ./...` 共 118 个测试通过、0 失败、5 个测试包通过,`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 ## Journey Log
@@ -46,6 +46,9 @@ branch: main
- [lesson] 迭代 1:网络写入的 `Send` 错误无法证明提交与否;最小安全策略是不自动重试、显式统计 ambiguous,并将连接标记为不健康。 - [lesson] 迭代 1:网络写入的 `Send` 错误无法证明提交与否;最小安全策略是不自动重试、显式统计 ambiguous,并将连接标记为不健康。
- [lesson] 迭代 1:进程内队列只能提供有界 best-effort;跨进程幂等需要持久化 outbox 和下游幂等协议,不能由本地重试可靠模拟。 - [lesson] 迭代 1:进程内队列只能提供有界 best-effort;跨进程幂等需要持久化 outbox 和下游幂等协议,不能由本地重试可靠模拟。
- [pivot] 迭代 1:SSE 完成判定限定为真正的 `text/event-stream` 且等待代理正常返回,避免终止标记导致提前记录不完整响应。 - [pivot] 迭代 1:SSE 完成判定限定为真正的 `text/event-stream` 且等待代理正常返回,避免终止标记导致提前记录不完整响应。
- [lesson] 迭代 4:最终报告记录禁用测试缓存的验证命令而非易过期的精确计数,并同时执行格式、静态检查和构建。
- [finding] 迭代 5:格式、vet、构建和测试均退出成功,但 passing test 中仍可隐藏由 `net/http` recover 的 handler panicWebSocket hijack 后不得再刷新普通 HTTP response writer。
- [fix] 最终验收:增加 `MAX_REQUEST_BYTES` 硬上限、上游连接写 deadline、Queue deadline 后有界清理,以及 hijack 后 flush guard;代理和队列压力测试各连续 20 轮通过。
## Source Materials ## Source Materials
+63 -12
View File
@@ -8,6 +8,7 @@ import (
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "time"
"unsafe"
"github.com/ClickHouse/clickhouse-go/v2/lib/driver" "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
@@ -47,10 +48,15 @@ func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string)
type poolBackend struct { 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) { 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 { if conn == nil {
return nil, errors.New("pool 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) 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 { type Stats struct {
Enqueued uint64 Enqueued uint64
@@ -153,8 +166,8 @@ func EstimatedBytes(e *LogEntry) int64 {
if e == nil { if e == nil {
return 0 return 0
} }
return int64(len(e.RequestID) + len(e.Method) + len(e.Path) + len(e.Query) + len(e.ClientIP) + len(e.Error) + 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)) len(e.RequestHeaders)+len(e.RequestBody)+len(e.ResponseHeaders)+len(e.ResponseBody))
} }
func (q *Queue) Stats() Stats { func (q *Queue) Stats() Stats {
@@ -195,9 +208,10 @@ func (q *Queue) Stop(contexts ...context.Context) error {
if len(contexts) > 0 && contexts[0] != nil { if len(contexts) > 0 && contexts[0] != nil {
ctx = contexts[0] ctx = contexts[0]
} }
q.mu.Lock() q.mu.Lock()
shutdownOwner := false
if !q.stopped { if !q.stopped {
shutdownOwner = true
q.stopped = true q.stopped = true
q.shutdownCtx = ctx q.shutdownCtx = ctx
close(q.ch) close(q.ch)
@@ -222,8 +236,14 @@ func (q *Queue) Stop(contexts ...context.Context) error {
case <-done: case <-done:
return nil return nil
case <-ctx.Done(): case <-ctx.Done():
if workCancel != nil { if shutdownOwner && workCancel != nil {
workCancel() workCancel()
timer := time.NewTimer(100 * time.Millisecond)
defer timer.Stop()
select {
case <-done:
case <-timer.C:
}
} }
return ctx.Err() return ctx.Err()
} }
@@ -235,8 +255,7 @@ func (q *Queue) Submit(e *LogEntry) {
q.dropped.Add(1) q.dropped.Add(1)
return return
} }
entry := cloneLogEntry(e) size := EstimatedBytes(e)
size := EstimatedBytes(entry)
if !q.mu.TryRLock() { if !q.mu.TryRLock() {
q.dropped.Add(1) q.dropped.Add(1)
return return
@@ -246,6 +265,7 @@ func (q *Queue) Submit(e *LogEntry) {
q.dropped.Add(1) q.dropped.Add(1)
return return
} }
entry := cloneLogEntry(e)
select { select {
case q.ch <- queuedEntry{entry: entry, size: size}: case q.ch <- queuedEntry{entry: entry, size: size}:
q.enq.Add(1) q.enq.Add(1)
@@ -364,16 +384,19 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
retry := entries retry := entries
var lastErr error var lastErr error
var failedGeneration uint64
for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ { for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ {
if err := ctx.Err(); err != nil { if err := ctx.Err(); err != nil {
lastErr = err lastErr = err
break 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.failed.Add(uint64(result.Failed))
q.ambiguous.Add(uint64(result.Ambiguous)) q.ambiguous.Add(uint64(result.Ambiguous))
if result.Ambiguous > 0 { if result.Ambiguous > 0 {
q.backend.MarkUnhealthy() markBackendUnhealthy(q.backend, failedGeneration)
} }
lastErr = result.Err lastErr = result.Err
retry = result.Retry retry = result.Retry
@@ -388,11 +411,31 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
if len(retry) > 0 { if len(retry) > 0 {
q.failed.Add(uint64(len(retry))) 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) 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 { func waitBackoff(ctx context.Context, attempt int) bool {
wait := BaseBackoff wait := BaseBackoff
for i := 1; i < attempt; i++ { for i := 1; i < attempt; i++ {
@@ -428,12 +471,20 @@ func Flush(ctx context.Context, conn BatchPreparer, entries []*LogEntry) FlushRe
} }
combined := FlushResult{} combined := FlushResult{}
for _, entry := range entries { for i, entry := range entries {
single := flushOnce(ctx, conn, []*LogEntry{entry}).public([]*LogEntry{entry}) single := flushOnce(ctx, conn, []*LogEntry{entry}).public([]*LogEntry{entry})
combined.Retry = append(combined.Retry, single.Retry...) combined.Retry = append(combined.Retry, single.Retry...)
combined.Failed += single.Failed combined.Failed += single.Failed
combined.Ambiguous += single.Ambiguous combined.Ambiguous += single.Ambiguous
combined.Err = errors.Join(combined.Err, single.Err) combined.Err = errors.Join(combined.Err, single.Err)
if single.Ambiguous > 0 {
if backend, ok := conn.(interface{ MarkUnhealthy() }); ok {
backend.MarkUnhealthy()
}
combined.Failed += len(combined.Retry) + len(entries) - i - 1
combined.Retry = nil
break
}
} }
return combined return combined
} }
+6 -6
View File
@@ -30,7 +30,7 @@ type shutdownStep struct {
run func(context.Context) error 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 { select {
case <-stop: case <-stop:
log.Printf("[main] shutdown signal received") log.Printf("[main] shutdown signal received")
@@ -83,7 +83,9 @@ func run() error {
} }
log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns)) log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns))
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() defer cancel()
pool := db.NewPool(rootCtx, cfg.ClickHouseURL, cfg.DBReconnectInterval) pool := db.NewPool(rootCtx, cfg.ClickHouseURL, cfg.DBReconnectInterval)
@@ -97,6 +99,7 @@ func run() error {
SSEIdleTimeout: cfg.UpstreamStreamIdleTimeout, SSEIdleTimeout: cfg.UpstreamStreamIdleTimeout,
UpstreamTLSInsecureSkipVerify: cfg.UpstreamTLSInsecureSkipVerify, UpstreamTLSInsecureSkipVerify: cfg.UpstreamTLSInsecureSkipVerify,
TrustedProxies: cfg.TrustedProxies, TrustedProxies: cfg.TrustedProxies,
MaxRequestBytes: cfg.MaxRequestBytes,
}) })
srv := &http.Server{ srv := &http.Server{
@@ -114,10 +117,7 @@ func run() error {
serverErr <- srv.ListenAndServe() serverErr <- srv.ListenAndServe()
}() }()
stop := make(chan os.Signal, 1) runErr := waitForShutdown(rootCtx.Done(), serverErr)
signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
runErr := waitForShutdown(stop, serverErr)
signal.Stop(stop)
// HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。 // HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+64 -5
View File
@@ -3,20 +3,79 @@ package main
import ( import (
"context" "context"
"errors" "errors"
"go/ast"
"go/parser"
"go/token"
"net/http" "net/http"
"os"
"reflect" "reflect"
"strings" "strings"
"testing" "testing"
"time" "time"
) )
func TestRunRegistersSignalsBeforeStartingResources(t *testing.T) {
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "main.go", nil, 0)
if err != nil {
t.Fatalf("parse main.go: %v", err)
}
var runBody *ast.BlockStmt
for _, declaration := range file.Decls {
function, ok := declaration.(*ast.FuncDecl)
if ok && function.Name.Name == "run" {
runBody = function.Body
break
}
}
if runBody == nil {
t.Fatal("main.go does not define run")
}
positions := make(map[string]token.Pos)
ast.Inspect(runBody, func(node ast.Node) bool {
call, ok := node.(*ast.CallExpr)
if !ok {
return true
}
selector, ok := call.Fun.(*ast.SelectorExpr)
if !ok {
return true
}
owner, ok := selector.X.(*ast.Ident)
if !ok {
return true
}
name := owner.Name + "." + selector.Sel.Name
switch name {
case "signal.NotifyContext", "db.NewPool", "queue.Start", "srv.ListenAndServe":
positions[name] = call.Pos()
}
return true
})
notifyPos, ok := positions["signal.NotifyContext"]
if !ok {
t.Fatal("run does not register for shutdown signals")
}
for _, start := range []string{"db.NewPool", "queue.Start", "srv.ListenAndServe"} {
startPos, ok := positions[start]
if !ok {
t.Fatalf("run does not call %s", start)
}
if notifyPos >= startPos {
t.Errorf("signal.NotifyContext at line %d must precede %s at line %d",
fset.Position(notifyPos).Line, start, fset.Position(startPos).Line)
}
}
}
func TestWaitForShutdownReturnsListenerError(t *testing.T) { func TestWaitForShutdownReturnsListenerError(t *testing.T) {
listenErr := errors.New("listen failed") listenErr := errors.New("listen failed")
serverErr := make(chan error, 1) serverErr := make(chan error, 1)
serverErr <- listenErr 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") { if !errors.Is(err, listenErr) || !strings.Contains(err.Error(), "server") {
t.Fatalf("waitForShutdown() error = %v, want wrapped listener error", err) t.Fatalf("waitForShutdown() error = %v, want wrapped listener error", err)
} }
@@ -24,8 +83,8 @@ func TestWaitForShutdownReturnsListenerError(t *testing.T) {
func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) { func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) {
t.Run("signal", func(t *testing.T) { t.Run("signal", func(t *testing.T) {
stop := make(chan os.Signal, 1) stop := make(chan struct{})
stop <- os.Interrupt close(stop)
if err := waitForShutdown(stop, make(chan error)); err != nil { if err := waitForShutdown(stop, make(chan error)); err != nil {
t.Fatalf("waitForShutdown() error = %v, want nil", err) t.Fatalf("waitForShutdown() error = %v, want nil", err)
} }
@@ -34,7 +93,7 @@ func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) {
t.Run("server closed", func(t *testing.T) { t.Run("server closed", func(t *testing.T) {
serverErr := make(chan error, 1) serverErr := make(chan error, 1)
serverErr <- http.ErrServerClosed 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) t.Fatalf("waitForShutdown() error = %v, want nil", err)
} }
}) })
+12 -3
View File
@@ -5,21 +5,30 @@ import (
"crypto/rand" "crypto/rand"
"encoding/hex" "encoding/hex"
"encoding/json" "encoding/json"
"errors"
"io" "io"
"net/http" "net/http"
"net/netip" "net/netip"
"strings" "strings"
) )
var errRequestBodyTooLarge = errors.New("request body too large")
// readRequestBody 在转发前完整读取请求体,确保读取失败时不会向上游发送损坏请求。 // 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.ContentLength == 0 { if r.Body == nil || r.Body == http.NoBody {
return nil, false, nil 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 { if err != nil {
return nil, false, err return nil, false, err
} }
if int64(len(body)) > requestLimit {
return nil, false, errRequestBodyTooLarge
}
if err := r.Body.Close(); err != nil { if err := r.Body.Close(); err != nil {
return nil, false, err return nil, false, err
} }
+39 -5
View File
@@ -3,6 +3,7 @@ package proxy
import ( import (
"context" "context"
"crypto/tls" "crypto/tls"
"errors"
"log" "log"
"net" "net"
"net/http" "net/http"
@@ -29,10 +30,13 @@ type Handler struct {
filter *config.Filter filter *config.Filter
queue LogSubmitter queue LogSubmitter
maxBodyBytes int64 maxBodyBytes int64
maxRequestBytes int64
trusted []netip.Prefix trusted []netip.Prefix
connMu sync.Mutex connMu sync.Mutex
conns map[net.Conn]struct{} conns map[net.Conn]struct{}
connChanged chan struct{} connChanged chan struct{}
closing bool
closed bool
} }
// Options 控制反代连接上游时的网络行为。 // Options 控制反代连接上游时的网络行为。
@@ -42,6 +46,7 @@ type Options struct {
SSEIdleTimeout time.Duration SSEIdleTimeout time.Duration
UpstreamTLSInsecureSkipVerify bool UpstreamTLSInsecureSkipVerify bool
TrustedProxies []netip.Prefix TrustedProxies []netip.Prefix
MaxRequestBytes int64
} }
// requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。 // requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。
@@ -75,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 { 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 := httputil.NewSingleHostReverseProxy(upstream)
rp.FlushInterval = -1 // 让流式 chunk 立即转发 rp.FlushInterval = -1 // 让流式 chunk 立即转发
if opts.UpstreamTimeout > 0 || opts.UpstreamTLSInsecureSkipVerify { if opts.UpstreamTimeout > 0 || opts.UpstreamTLSInsecureSkipVerify {
transport := http.DefaultTransport.(*http.Transport).Clone() transport := http.DefaultTransport.(*http.Transport).Clone()
if opts.UpstreamTimeout > 0 { 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.ResponseHeaderTimeout = opts.UpstreamTimeout
transport.TLSHandshakeTimeout = opts.UpstreamTimeout transport.TLSHandshakeTimeout = opts.UpstreamTimeout
} }
@@ -132,6 +147,7 @@ func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter
filter: filter, filter: filter,
queue: queue, queue: queue,
maxBodyBytes: maxBody, maxBodyBytes: maxBody,
maxRequestBytes: opts.MaxRequestBytes,
trusted: append([]netip.Prefix(nil), opts.TrustedProxies...), trusted: append([]netip.Prefix(nil), opts.TrustedProxies...),
conns: make(map[net.Conn]struct{}), conns: make(map[net.Conn]struct{}),
connChanged: make(chan struct{}), connChanged: make(chan struct{}),
@@ -147,11 +163,15 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
} }
shouldLog := h.filter.ShouldLog(r.URL.Path) 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 { if err != nil && !shouldLog {
log.Printf("[proxy] read request body failed: %v", err) log.Printf("[proxy] read request body failed: %v", err)
_ = r.Body.Close() _ = r.Body.Close()
if errors.Is(err, errRequestBodyTooLarge) {
http.Error(w, "request body too large", http.StatusRequestEntityTooLarge)
} else {
http.Error(w, "bad request", http.StatusBadRequest) http.Error(w, "bad request", http.StatusBadRequest)
}
return return
} }
if !shouldLog { if !shouldLog {
@@ -174,7 +194,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
_ = r.Body.Close() _ = r.Body.Close()
s := "read request body: " + err.Error() s := "read request body: " + err.Error()
st.lastErr.Store(&s) st.lastErr.Store(&s)
h.serveRequestBodyError(cw, r, st, started, reqID) h.serveRequestBodyError(cw, r, st, started, reqID, errors.Is(err, errRequestBodyTooLarge))
return return
} }
@@ -221,15 +241,22 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
h.rp.ServeHTTP(cw, r) h.rp.ServeHTTP(cw, r)
finished := time.Now() finished := time.Now()
cw.ConfirmDelivery()
responseComplete := !isStreamResponse(cw.Header()) || cw.SSEComplete() responseComplete := !isStreamResponse(cw.Header()) || cw.SSEComplete()
if cw.Complete() && responseComplete && !st.readFailed.Load() && r.Context().Err() == nil { if cw.Complete() && responseComplete && !st.readFailed.Load() && r.Context().Err() == nil {
submitOnce.Do(func() { submit(finished) }) submitOnce.Do(func() { submit(finished) })
} }
} }
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) 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() { if !cw.Complete() {
return return
} }
@@ -252,7 +279,9 @@ func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *
func (h *Handler) Shutdown(ctx context.Context) error { func (h *Handler) Shutdown(ctx context.Context) error {
for { for {
h.connMu.Lock() h.connMu.Lock()
h.closing = true
if len(h.conns) == 0 { if len(h.conns) == 0 {
h.closed = true
h.connMu.Unlock() h.connMu.Unlock()
return nil return nil
} }
@@ -283,6 +312,11 @@ func (h *Handler) trackConn(conn net.Conn) net.Conn {
h.connMu.Unlock() h.connMu.Unlock()
} }
h.connMu.Lock() h.connMu.Lock()
if h.closing || h.closed {
h.connMu.Unlock()
_ = conn.Close()
return conn
}
h.conns[tracked] = struct{}{} h.conns[tracked] = struct{}{}
close(h.connChanged) close(h.connChanged)
h.connChanged = make(chan struct{}) h.connChanged = make(chan struct{})
+58 -12
View File
@@ -555,8 +555,13 @@ type geminiCandidate struct {
raw map[string]json.RawMessage raw map[string]json.RawMessage
index int index int
role string role string
parts []*geminiPart
}
type geminiPart struct {
raw map[string]json.RawMessage
text strings.Builder text strings.Builder
partRaw map[string]json.RawMessage kind string
} }
func assembleGeminiSSE(payloads []string) ([]byte, bool) { func assembleGeminiSSE(payloads []string) ([]byte, bool) {
@@ -588,15 +593,31 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
assembled.role = candidate.Content.Role assembled.role = candidate.Content.Role
} }
if len(candidate.Content.Parts) > 0 { if len(candidate.Content.Parts) > 0 {
for _, part := range candidate.Content.Parts {
assembled.text.WriteString(part.Text)
}
if len(assembled.partRaw) == 0 {
var contentRaw struct { var contentRaw struct {
Parts []map[string]json.RawMessage `json:"parts"` Parts []map[string]json.RawMessage `json:"parts"`
} }
if rawContent, ok := candidate.Raw["content"]; ok && json.Unmarshal(rawContent, &contentRaw) == nil && len(contentRaw.Parts) > 0 { if rawContent, ok := candidate.Raw["content"]; ok {
assembled.partRaw = cloneRawMap(contentRaw.Parts[0]) _ = 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
}
}
}
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] candidate := candidates[index]
candidateRaw := cloneRawMap(candidate.raw) candidateRaw := cloneRawMap(candidate.raw)
candidateRaw["index"] = mustJSON(candidate.index) candidateRaw["index"] = mustJSON(candidate.index)
contentRaw := map[string]any{"role": candidate.role, "parts": []any{map[string]any{"text": candidate.text.String()}}} parts := make([]any, 0, len(candidate.parts))
if len(candidate.partRaw) > 0 { for _, part := range candidate.parts {
partRaw := cloneRawMap(candidate.partRaw) partRaw := cloneRawMap(part.raw)
partRaw["text"] = mustJSON(candidate.text.String()) if part.text.Len() > 0 || len(partRaw) == 0 {
contentRaw["parts"] = []any{rawMapToMap(partRaw)} partRaw["text"] = mustJSON(part.text.String())
} }
parts = append(parts, rawMapToMap(partRaw))
}
contentRaw := map[string]any{"role": candidate.role, "parts": parts}
candidateRaw["content"] = mustJSON(contentRaw) candidateRaw["content"] = mustJSON(contentRaw)
assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw)) assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw))
} }
@@ -631,6 +655,28 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
return data, err == nil 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 { func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage {
out := make(map[string]json.RawMessage, len(in)) out := make(map[string]json.RawMessage, len(in))
for k, v := range in { for k, v := range in {
+18 -1
View File
@@ -8,12 +8,29 @@ import (
type sseEventTracker struct { type sseEventTracker struct {
buf []byte buf []byte
limit int
overflow bool
recognized bool recognized bool
terminal bool terminal bool
afterTerminal 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) { 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...) t.buf = append(t.buf, p...)
for { for {
end, separator := completeSSEEvent(t.buf) end, separator := completeSSEEvent(t.buf)
@@ -34,7 +51,7 @@ func (t *sseEventTracker) Write(p []byte) {
} }
func (t *sseEventTracker) Complete() bool { 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) { func completeSSEEvent(buf []byte) (int, int) {
+12
View File
@@ -19,3 +19,15 @@ func TestSSEEventTrackerAcceptsCROnlyEventBoundary(t *testing.T) {
t.Fatal("terminal SSE event with CR-only boundary should be complete") 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))
}
}
+18
View File
@@ -2,12 +2,30 @@ package proxy
import ( import (
"io" "io"
"net"
"net/http" "net/http"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "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 { type trackingBody struct {
io.ReadCloser io.ReadCloser
failed *atomic.Bool failed *atomic.Bool
+18 -2
View File
@@ -26,7 +26,7 @@ type captureWriter struct {
} }
func newCaptureWriter(w http.ResponseWriter, max int64) *captureWriter { 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) { func (c *captureWriter) WriteHeader(code int) {
@@ -63,7 +63,7 @@ func (c *captureWriter) Write(p []byte) (int, error) {
c.truncated = true c.truncated = true
} }
c.written += int64(n) c.written += int64(n)
if isStreamResponse(c.Header()) { if c.max > 0 && isStreamResponse(c.Header()) {
c.sse.Write(p[:n]) c.sse.Write(p[:n])
} }
c.flush() c.flush()
@@ -75,7 +75,23 @@ func (c *captureWriter) Flush() {
c.flush() c.flush()
} }
// ConfirmDelivery establishes an observable delivery boundary for responses
// whose headers were not followed by a body write.
func (c *captureWriter) ConfirmDelivery() {
if c.hijacked || c.written != 0 {
return
}
if _, ok := c.ResponseWriter.(http.Flusher); !ok {
c.writeFailed = true
return
}
c.flush()
}
func (c *captureWriter) flush() { func (c *captureWriter) flush() {
if c.hijacked {
return
}
if _, ok := c.ResponseWriter.(http.Flusher); !ok { if _, ok := c.ResponseWriter.(http.Flusher); !ok {
return return
} }
+22
View File
@@ -1,7 +1,9 @@
package config_test package config_test
import ( import (
"math"
"net/netip" "net/netip"
"strconv"
"strings" "strings"
"testing" "testing"
"time" "time"
@@ -50,6 +52,9 @@ func TestLoadQueueDefaults(t *testing.T) {
if cfg.MaxBodyBytes != 1<<20 { if cfg.MaxBodyBytes != 1<<20 {
t.Fatalf("MaxBodyBytes=%d want %d", 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 { if cfg.LogQueueSize != 256 {
t.Fatalf("LogQueueSize=%d want 256", cfg.LogQueueSize) 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) { func TestLoadValidatesClickHouseURL(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
+69 -11
View File
@@ -42,8 +42,11 @@ type fakePool struct {
queries []string queries []string
batches []*fakeBatch batches []*fakeBatch
scripts []batchScript scripts []batchScript
unhealthy bool
} }
func (p *fakePool) MarkUnhealthy() { p.unhealthy = true }
type queueBackend struct { type queueBackend struct {
prepare func(context.Context, string) (logger.Batch, error) prepare func(context.Context, string) (logger.Batch, error)
healthy atomic.Bool healthy atomic.Bool
@@ -215,14 +218,34 @@ func TestFlushAppendRecoveryCountsAmbiguousSingleSend(t *testing.T) {
} }
} }
func TestFlushAppendRecoveryStopsAfterFirstAmbiguousSend(t *testing.T) {
p := &fakePool{scripts: []batchScript{
{appendAt: 1, appendErr: errors.New("bad batch")},
{sendErr: errors.New("ack lost")},
{},
}}
entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}, {RequestID: "c"}}
result := logger.Flush(context.Background(), p, entries)
if result.Failed != 3 || result.Ambiguous != 1 || len(result.Retry) != 0 {
t.Fatalf("unexpected result: %+v", result)
}
if p.calls != 2 {
t.Fatalf("PrepareBatch calls=%d want 2; entries after ambiguous Send must not be attempted", p.calls)
}
if !p.unhealthy {
t.Fatal("ambiguous singleton Send did not immediately mark backend unhealthy")
}
}
func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) { func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) {
entry := &logger.LogEntry{ entry := &logger.LogEntry{
RequestID: "1", Method: "22", Path: "333", Query: "4444", ClientIP: "55555", Error: "666666", RequestID: "1", Method: "22", Path: "333", Query: "4444", ClientIP: "55555", Error: "666666",
RequestHeaders: []byte("7777777"), RequestBody: []byte("88888888"), RequestHeaders: []byte("7777777"), RequestBody: []byte("88888888"),
ResponseHeaders: []byte("999999999"), ResponseBody: []byte("0000000000"), ResponseHeaders: []byte("999999999"), ResponseBody: []byte("0000000000"),
} }
if got, want := logger.EstimatedBytes(entry), int64(55); got != want { if got := logger.EstimatedBytes(entry); got <= 55 {
t.Fatalf("EstimatedBytes=%d want %d", got, want) t.Fatalf("EstimatedBytes=%d must include fixed entry overhead", got)
} }
} }
@@ -393,7 +416,7 @@ func TestQueueByteBudgetTracksOwnedEntrySnapshot(t *testing.T) {
} }
} }
func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) { func TestQueueCanceledStopReleasesAllBudgetBeforeReturning(t *testing.T) {
entry := &logger.LogEntry{RequestID: "queued"} entry := &logger.LogEntry{RequestID: "queued"}
q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry)) q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry))
q.Start(context.Background()) q.Start(context.Background())
@@ -405,10 +428,6 @@ func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) {
cancel() cancel()
_ = q.Stop(ctx) _ = q.Stop(ctx)
deadline := time.Now().Add(time.Second)
for q.Stats().Bytes != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 32 { if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 32 {
t.Fatalf("stats=%+v", stats) t.Fatalf("stats=%+v", stats)
} }
@@ -443,6 +462,49 @@ func TestQueueStopAndSubmitAreConcurrentAndRepeatSafe(t *testing.T) {
} }
} }
func TestConcurrentStopCannotCancelFirstStopDrain(t *testing.T) {
const initialEntries = 32
firstPrepareStarted := make(chan struct{})
releaseFirstPrepare := make(chan struct{})
var prepareCalls atomic.Int32
backend := newQueueBackend(func(context.Context, string) (logger.Batch, error) {
if prepareCalls.Add(1) == 1 {
close(firstPrepareStarted)
<-releaseFirstPrepare
}
return &fakeBatch{}, nil
})
q := logger.NewQueueWithBackend(backend, 1024, 1, 1, time.Hour)
q.Start(context.Background())
for i := 0; i < initialEntries; i++ {
q.Submit(&logger.LogEntry{RequestID: "queued"})
}
<-firstPrepareStarted
firstResult := make(chan error, 1)
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"})
}
enqueued := q.Stats().Enqueued
secondCtx, cancelSecond := context.WithCancel(context.Background())
cancelSecond()
if err := q.Stop(secondCtx); !errors.Is(err, context.Canceled) {
t.Fatalf("second Stop error=%v want canceled", err)
}
close(releaseFirstPrepare)
if err := <-firstResult; err != nil {
t.Fatalf("first Stop: %v", err)
}
if calls := uint64(backend.calls.Load()); calls != enqueued {
t.Fatalf("PrepareBatch calls=%d want %d; second Stop interrupted drain", calls, enqueued)
}
}
func TestStopIsIndependentFromStartContext(t *testing.T) { func TestStopIsIndependentFromStartContext(t *testing.T) {
root, cancelRoot := context.WithCancel(context.Background()) root, cancelRoot := context.WithCancel(context.Background())
q := logger.NewQueue(nil, 4, 4, 1, time.Hour, 1024) q := logger.NewQueue(nil, 4, 4, 1, time.Hour, 1024)
@@ -498,10 +560,6 @@ func TestQueueStopDeadlineCancelsBlockedSend(t *testing.T) {
if err := q.Stop(ctx); !errors.Is(err, context.DeadlineExceeded) { if err := q.Stop(ctx); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Stop error=%v want deadline exceeded", err) t.Fatalf("Stop error=%v want deadline exceeded", err)
} }
deadline := time.Now().Add(250 * time.Millisecond)
for q.Stats().Bytes != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if stats := q.Stats(); stats.Bytes != 0 || stats.Ambiguous != 1 || stats.Failed != 1 { if stats := q.Stats(); stats.Bytes != 0 || stats.Ambiguous != 1 || stats.Failed != 1 {
t.Fatalf("stats after canceled Send=%+v", stats) t.Fatalf("stats after canceled Send=%+v", stats)
} }
+128
View File
@@ -67,6 +67,14 @@ func (*failingFlushResponseWriter) FlushError() error {
return errors.New("flush failed") return errors.New("flush failed")
} }
type unflushableResponseWriter struct{ header http.Header }
func (w *unflushableResponseWriter) Header() http.Header { return w.header }
func (*unflushableResponseWriter) WriteHeader(int) {}
func (*unflushableResponseWriter) Write(p []byte) (int, error) {
return len(p), nil
}
func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler { func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler {
t.Helper() t.Helper()
filter, err := config.NewFilter(config.FilterDisabled, nil) filter, err := config.NewFilter(config.FilterDisabled, nil)
@@ -100,6 +108,62 @@ 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) {
upstreamCalls.Add(1)
w.WriteHeader(http.StatusOK)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{})
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", nil)
req.Body = failingBody{err: errors.New("secret read failure")}
req.ContentLength = 0
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest || rec.Body.String() != "bad request\n" {
t.Fatalf("response=(%d, %q), want fixed 400 bad request", rec.Code, rec.Body.String())
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestRequestBodyReadFailureAfterCaptureLimitDoesNotReachUpstream(t *testing.T) { func TestRequestBodyReadFailureAfterCaptureLimitDoesNotReachUpstream(t *testing.T) {
var upstreamCalls atomic.Int32 var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -196,6 +260,70 @@ func TestResponseFlushFailurePreventsCommit(t *testing.T) {
} }
} }
func TestHeaderOnlyResponseFlushFailurePreventsCommit(t *testing.T) {
tests := []struct {
name string
method string
status int
}{
{name: "head", method: http.MethodHead, status: http.StatusOK},
{name: "no content", method: http.MethodGet, status: http.StatusNoContent},
{name: "not modified", method: http.MethodGet, status: http.StatusNotModified},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(tc.status)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
w := &failingFlushResponseWriter{header: make(http.Header)}
h.ServeHTTP(w, httptest.NewRequest(tc.method, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatal("response with undelivered headers was committed")
}
})
}
}
func TestHeaderOnlyResponseSuccessfulFlushCommits(t *testing.T) {
for _, status := range []int{http.StatusNoContent, http.StatusNotModified} {
t.Run(http.StatusText(status), func(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(status)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 1 {
t.Fatalf("entries=%d, want successfully delivered response committed", sub.Len())
}
})
}
}
func TestHeaderOnlyResponseWithoutDeliveryBoundaryDoesNotCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusNoContent)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
w := &unflushableResponseWriter{header: make(http.Header)}
h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatal("header-only response without a delivery boundary was committed")
}
}
func TestTrustedProxyClientIP(t *testing.T) { func TestTrustedProxyClientIP(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, "ok") _, _ = io.WriteString(w, "ok")
+52
View File
@@ -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) { func TestUnknownSSEKeepsRawBody(t *testing.T) {
e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions") e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions")
+80
View File
@@ -11,6 +11,7 @@ import (
"net/http/httptest" "net/http/httptest"
"net/url" "net/url"
"strings" "strings"
"sync"
"testing" "testing"
"time" "time"
@@ -45,6 +46,30 @@ type failingWriteConn struct{ net.Conn }
func (failingWriteConn) Write([]byte) (int, error) { return 0, errors.New("handshake write failed") } func (failingWriteConn) Write([]byte) (int, error) { return 0, errors.New("handshake write failed") }
type closeNotifyConn struct {
net.Conn
closed chan struct{}
once sync.Once
}
func (c *closeNotifyConn) Close() error {
err := c.Conn.Close()
c.once.Do(func() { close(c.closed) })
return err
}
type blockingHijackResponseWriter struct {
*hijackResponseWriter
hijacked chan struct{}
release chan struct{}
}
func (w *blockingHijackResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
close(w.hijacked)
<-w.release
return w.hijackResponseWriter.Hijack()
}
// fakeUpstream 模拟一个最简 WebSocket 升级: // fakeUpstream 模拟一个最简 WebSocket 升级:
// 收到 GET + Upgrade: websocket 后回 101,然后做字节回声直到对端关闭。 // 收到 GET + Upgrade: websocket 后回 101,然后做字节回声直到对端关闭。
func fakeUpstream(t *testing.T) *httptest.Server { func fakeUpstream(t *testing.T) *httptest.Server {
@@ -161,6 +186,61 @@ func TestWebSocketHandshakeCapturedAndShutdownClosesConnection(t *testing.T) {
} }
} }
func TestShutdownClosesConnectionHijackedBeforeRegistration(t *testing.T) {
upstream := fakeUpstream(t)
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, noopSubmitter{}, 1024)
downstream, peer := net.Pipe()
defer peer.Close()
closed := make(chan struct{})
conn := &closeNotifyConn{Conn: downstream, closed: closed}
w := &blockingHijackResponseWriter{
hijackResponseWriter: &hijackResponseWriter{header: make(http.Header), conn: conn},
hijacked: make(chan struct{}),
release: make(chan struct{}),
}
req := httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/realtime", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
serveDone := make(chan struct{})
var releaseOnce sync.Once
releaseHijack := func() { releaseOnce.Do(func() { close(w.release) }) }
go func() {
h.ServeHTTP(w, req)
close(serveDone)
}()
defer func() {
releaseHijack()
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
_ = h.Shutdown(ctx)
<-serveDone
}()
select {
case <-w.hijacked:
case <-time.After(time.Second):
t.Fatal("downstream connection was not hijacked")
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := h.Shutdown(ctx); err != nil {
t.Fatal(err)
}
releaseHijack()
select {
case <-closed:
case <-time.After(100 * time.Millisecond):
t.Fatal("connection hijacked during Shutdown remained open")
}
}
func TestWebSocketHandshakeWriteFailureIsNotCaptured(t *testing.T) { func TestWebSocketHandshakeWriteFailureIsNotCaptured(t *testing.T) {
upstream := fakeUpstream(t) upstream := fakeUpstream(t)
defer upstream.Close() defer upstream.Close()