From d1bbb5370c0841af6cfe447c92d2ca6875ac3e8a Mon Sep 17 00:00:00 2001 From: MiMoCode Date: Fri, 10 Jul 2026 18:26:48 +0800 Subject: [PATCH] fix: restore reviewable migration evidence --- .dockerignore | 8 + .env.example | 57 ++ .gitignore | 54 ++ Dockerfile | 18 + README.md | 228 +++++ compose.yml | 65 ++ config/config.go | 245 ++++++ config/filter.go | 134 +++ db/clickhouse.go | 221 +++++ db/clickhouse_test.go | 127 +++ db/migrate.go | 134 +++ .../plans/2026-07-09-clickhouse-migration.md | 66 ++ .../plans/reliability-security-fixes.md | 5 + ...-10-clickhouse-migration-security-fixes.md | 32 + .../reports/reliability-security-fixes.md | 55 ++ .../specs/2026-07-09-clickhouse-migration.md | 54 ++ .../specs/reliability-security-fixes.md | 5 + filter.yaml | 25 + go.mod | 29 + go.sum | 61 ++ logger/model.go | 24 + logger/queue.go | 447 ++++++++++ main.go | 109 +++ proxy/capture.go | 104 +++ proxy/proxy.go | 339 +++++++ proxy/sse.go | 661 ++++++++++++++ proxy/sse_event.go | 65 ++ proxy/timeout.go | 101 +++ proxy/writer.go | 128 +++ tests/README.md | 61 ++ tests/config/config_test.go | 283 ++++++ tests/config/filter_test.go | 164 ++++ tests/deployment/compose_test.go | 66 ++ tests/logger/queue_test.go | 323 +++++++ tests/proxy/error_test.go | 139 +++ tests/proxy/robustness_test.go | 336 +++++++ tests/proxy/sse_capture_test.go | 824 ++++++++++++++++++ tests/proxy/websocket_test.go | 299 +++++++ tests/scripts/dbcheck/main.go | 109 +++ tests/scripts/dump-stream-body.ps1 | 54 ++ tests/scripts/load-env.ps1 | 25 + tests/scripts/smoke.ps1 | 135 +++ 42 files changed, 6419 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 compose.yml create mode 100644 config/config.go create mode 100644 config/filter.go create mode 100644 db/clickhouse.go create mode 100644 db/clickhouse_test.go create mode 100644 db/migrate.go create mode 100644 docs/compose/plans/2026-07-09-clickhouse-migration.md create mode 100644 docs/compose/plans/reliability-security-fixes.md create mode 100644 docs/compose/reports/2026-07-10-clickhouse-migration-security-fixes.md create mode 100644 docs/compose/reports/reliability-security-fixes.md create mode 100644 docs/compose/specs/2026-07-09-clickhouse-migration.md create mode 100644 docs/compose/specs/reliability-security-fixes.md create mode 100644 filter.yaml create mode 100644 go.mod create mode 100644 go.sum create mode 100644 logger/model.go create mode 100644 logger/queue.go create mode 100644 main.go create mode 100644 proxy/capture.go create mode 100644 proxy/proxy.go create mode 100644 proxy/sse.go create mode 100644 proxy/sse_event.go create mode 100644 proxy/timeout.go create mode 100644 proxy/writer.go create mode 100644 tests/README.md create mode 100644 tests/config/config_test.go create mode 100644 tests/config/filter_test.go create mode 100644 tests/deployment/compose_test.go create mode 100644 tests/logger/queue_test.go create mode 100644 tests/proxy/error_test.go create mode 100644 tests/proxy/robustness_test.go create mode 100644 tests/proxy/sse_capture_test.go create mode 100644 tests/proxy/websocket_test.go create mode 100644 tests/scripts/dbcheck/main.go create mode 100644 tests/scripts/dump-stream-body.ps1 create mode 100644 tests/scripts/load-env.ps1 create mode 100644 tests/scripts/smoke.ps1 diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..5d9ceeb --- /dev/null +++ b/.dockerignore @@ -0,0 +1,8 @@ +tests/ +.env +.env.* +*.log +*.exe +models.json +tmp_*.go +dbcheck.exe diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..f06044a --- /dev/null +++ b/.env.example @@ -0,0 +1,57 @@ +# 监听地址 +LISTEN_ADDR=:8080 + +# newapi 后端；本地运行和 Compose 均必须显式设置 +UPSTREAM_URL= +# 是否跳过上游 HTTPS 证书校验；仅开发/可信内网自签证书场景使用 +UPSTREAM_TLS_INSECURE_SKIP_VERIFY=false + +# ClickHouse 原生协议地址；必须显式设置，支持 host:port、clickhouse:// 或 clickhouses:// +CLICKHOUSE_URL= + +# 单条请求/响应 body 最大记录字节数 +MAX_BODY_BYTES=1048576 + +# 异步日志队列容量 +LOG_QUEUE_SIZE=256 + +# 队列中所有待写日志的总字节预算 +LOG_QUEUE_BYTES=67108864 + +# 批量写入条数 +LOG_BATCH_SIZE=50 + +# 批量写入间隔（Go duration: 500ms, 2s, 1m...） +LOG_BATCH_INTERVAL=2s + +# 后台日志 worker 数 +LOG_WORKERS=2 + +# 黑白名单文件 +FILTER_FILE=./filter.yaml + +# DB 故障后重连尝试间隔 +DB_RECONNECT_INTERVAL=10s + +# HTTP 超时（Go duration: 30s, 10m...） +READ_TIMEOUT=30s +WRITE_TIMEOUT=10m +IDLE_TIMEOUT=5m +UPSTREAM_TIMEOUT=30s +UPSTREAM_RESPONSE_TIMEOUT=30s +UPSTREAM_STREAM_IDLE_TIMEOUT=2m + +# 可信反向代理 IP/CIDR，逗号分隔；为空时忽略所有客户端转发头 +TRUSTED_PROXIES= + +# ===== Docker Compose 可选配置 ===== +# 对外暴露的代理端口（容器内固定 :8080） +LISTEN_PORT=8080 + +# Compose 内置 ClickHouse 配置 +CLICKHOUSE_USER=tokenthief +CLICKHOUSE_PASSWORD= +CLICKHOUSE_DB=tokenthief + +# 容器时区 +TZ=Asia/Shanghai diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..cd1e4b3 --- /dev/null +++ b/.gitignore @@ -0,0 +1,54 @@ +# Environment / secrets +.env +.env.* +!.env.example + +# Local binaries +TokenThief +TokenThief.exe +tokenthief +tokenthief.exe +*.exe +*.dll +*.so +*.dylib + +# Go build/test outputs +bin/ +dist/ +build/ +*.test +*.out +coverage.out +coverage.html +*.coverprofile + +# Runtime logs and temporary outputs +*.log +proxy.log +proxy.err.log +models.json +tmp_*.go +dbcheck.exe + +# Local database / volumes +clickhouse_data/ +data/ +*.db +*.sqlite +*.sqlite3 + +# IDE / editor +.idea/ +.vscode/ +*.swp +*.swo +*~ + +# OS files +.DS_Store +Thumbs.db + +# Docker compose overrides +compose.override.yml +docker-compose.override.yml diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..ba1e626 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,18 @@ +# syntax=docker/dockerfile:1 +FROM golang:1.25-alpine AS build +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY . . +RUN CGO_ENABLED=0 go build -ldflags="-s -w" -o /out/TokenThief ./ + +FROM alpine:3.20 +RUN apk add --no-cache ca-certificates tzdata && \ + addgroup -S app && adduser -S -G app app +COPY --from=build /out/TokenThief /usr/local/bin/TokenThief +COPY filter.yaml /etc/tokenthief/filter.yaml +ENV FILTER_FILE=/etc/tokenthief/filter.yaml \ + TZ=Asia/Shanghai +USER app +EXPOSE 8080 +ENTRYPOINT ["/usr/local/bin/TokenThief"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..b2e5ba9 --- /dev/null +++ b/README.md @@ -0,0 +1,228 @@ +# TokenThief + +newapi 反向代理，采集请求/响应（包括 SSE/chunked 流式响应）并异步批量写入 ClickHouse。 + +## 特性 + +- 透明反代任意 HTTP 后端（默认目标为 newapi）。 +- 流式响应（SSE、chunked、OpenAI 兼容 chat completions）边转发边缓冲，结束后整体入库。 +- 异步队列 + 批量写入，主路径零阻塞。 +- **DB 故障不影响代理服务**：连接失败时丢弃日志，后台持续重连。 +- 通过 yaml 配置 glob 风格的黑白名单。 +- 原样存储 headers/body（String 字段保存 JSON 文本与 body 内容）。 + +## 快速开始 + +### 本地运行 + +```bash +cp .env.example .env +# 编辑 .env 设置 UPSTREAM_URL / CLICKHOUSE_URL +set -a; source .env; set +a +go run . +``` + +### Docker + +```bash +docker build -t tokenthief . +docker run --rm -p 8080:8080 \ + -e UPSTREAM_URL=http://newapi:3000 \ + -e CLICKHOUSE_URL='clickhouses://tokenthief:@clickhouse:9440/tokenthief' \ + -v $(pwd)/filter.yaml:/etc/tokenthief/filter.yaml \ + tokenthief +``` + +镜像内二进制路径为 `/usr/local/bin/TokenThief`。 + +### Docker Compose + +`compose.yml` 已包含 TokenThief 与 ClickHouse： + +```bash +cp .env.example .env +# 编辑 .env，设置 UPSTREAM_URL、强密码 CLICKHOUSE_PASSWORD 和对应的 CLICKHOUSE_URL +docker compose -f compose.yml up -d --build +``` + +默认端口： + +| 服务 | 地址 | +|---|---| +| TokenThief | `http://localhost:8080` | +| ClickHouse | 仅 Compose 内部网络，不默认发布宿主机端口 | + +Compose 在渲染配置时要求显式提供 `UPSTREAM_URL`、应用 DSN 和非空 ClickHouse 密码。密码中的 URL 特殊字符必须编码： + +```text +clickhouse://tokenthief:@thief_clickhouse:9000/tokenthief +``` + +### 运行测试 + +测试代码主要集中在 `tests/` 目录，数据库包还包含同包测试。 + +```bash +gofmt -w . +go test ./... +go vet ./... +``` + +构建二进制时不会引入测试内容：`_test.go` 不参与 `go build`，`tests/scripts/` 也不被主程序 import；Docker 构建时 `.dockerignore` 会把整个 `tests/` 目录排除在 build context 之外。 + +## 环境变量 + +| 变量 | 说明 | 默认 | +|---|---|---| +| `LISTEN_ADDR` | 监听地址 | `:8080` | +| `UPSTREAM_URL` | 后端地址（必填） | - | +| `UPSTREAM_TLS_INSECURE_SKIP_VERIFY` | 跳过上游 HTTPS 证书校验；仅开发/可信内网自签证书场景使用 | `false` | +| `CLICKHOUSE_URL` | ClickHouse 原生协议地址（必填）；支持严格 `host:port`、`clickhouse://` 或 TLS `clickhouses://` | - | +| `MAX_BODY_BYTES` | 单个请求体和响应体的记录上限 | `1048576` | +| `LOG_QUEUE_SIZE` | 异步队列条数上限 | `256` | +| `LOG_QUEUE_BYTES` | 异步队列总字节预算 | `67108864` | +| `LOG_BATCH_SIZE` | 批量写入条数 | `50` | +| `LOG_BATCH_INTERVAL` | 批量刷新间隔 | `2s` | +| `LOG_WORKERS` | worker 数 | `2` | +| `FILTER_FILE` | 黑白名单文件 | `./filter.yaml` | +| `DB_RECONNECT_INTERVAL` | DB 重连间隔 | `10s` | +| `READ_TIMEOUT` | 请求读取总超时 | `30s` | +| `WRITE_TIMEOUT` | 响应写入总超时 | `10m` | +| `IDLE_TIMEOUT` | HTTP keep-alive 空闲超时 | `5m` | +| `UPSTREAM_TIMEOUT` | 上游连接、TLS 握手、响应头等待超时 | `30s` | +| `UPSTREAM_RESPONSE_TIMEOUT` | 普通上游响应体总超时 | `30s` | +| `UPSTREAM_STREAM_IDLE_TIMEOUT` | SSE 上游响应体空闲超时，每次成功读取后重置 | `2m` | +| `TRUSTED_PROXIES` | 可信反向代理 IP/CIDR，逗号分隔；为空时忽略转发头 | 空 | + +## 黑白名单（filter.yaml） + +`CLICKHOUSE_URL` 中的账号、密码和数据库名会传给 ClickHouse 原生协议连接。密码如果包含 `@`、`:`、`/`、`#` 等 URL 特殊字符，需要先做 URL encode。仅填写 `host:port` 时会使用 ClickHouse 默认用户、空密码和默认数据库。 + +默认 `filter.yaml` 已按 [newapi 官方文档](https://docs.newapi.pro/zh/docs/api) 列出全部 AI 模型接口（chat、completions、embeddings、moderations、rerank、realtime、audio、images、videos、Claude、Gemini 等）。 + +```yaml +mode: whitelist # whitelist | blacklist | disabled +patterns: + - /v1/chat/completions + - /v1/audio/* # 单段通配 + - /v1/videos/** # 跨段通配 + - /v1beta/models/*:generateContent # 段内通配 +``` + +Pattern 语法： + +| 通配符 | 含义 | +|---|---| +| `*` | 匹配单个路径段内除 `/` 之外的任意字符 | +| `**` | 跨段匹配任意字符（包括 `/`） | +| `?` | 匹配单个非 `/` 字符 | + +## 端点 + +- `/healthz` — 始终返回 200，不参与反代与日志。 +- 其余路径 — 全部反代到 `UPSTREAM_URL`。 + +## 数据库表 + +启动时自动创建单张 ClickHouse MergeTree 表 `proxy_logs`（详见 `db/migrate.go`），按 `started_at` 月份分区。DDL 如下： + +```sql +CREATE TABLE IF NOT EXISTS proxy_logs ( + request_id String, + method String, + path String, + query String, + client_ip String, + request_headers String, + request_body String, + request_truncated Bool DEFAULT false, + status_code Int32, + response_headers String, + response_body String, + response_truncated Bool DEFAULT false, + is_stream Bool DEFAULT false, + latency_ms Int64, + started_at DateTime64(3), + finished_at DateTime64(3), + error String +) +ENGINE = MergeTree +PARTITION BY toYYYYMM(started_at) +ORDER BY (started_at, request_id); +``` + +完整字段说明如下。 + +### 字段一览 + +| 字段 | 类型 | 可空 | 说明 | +|---|---|---|---| +| `request_id` | `String` | 否 | 每次请求由代理生成的 16 字节随机 hex（32 个字符），同时写入响应头 `X-Request-Id`，便于客户端日志关联 | +| `method` | `String` | 否 | HTTP 方法（`GET` / `POST` / …），与客户端实际发送一致 | +| `path` | `String` | 否 | 请求路径，不含 query string，例如 `/v1/chat/completions` | +| `query` | `String` | 否 | 原始 query string（不含 `?`），例如 `model=gpt-4&stream=true`；无 query 时为空串 | +| `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` 截断。即使截断，下游 newapi 仍会通过 `MultiReader` 接收到完整 body，不影响功能 | +| `status_code` | `Int32` | 否 | 上游返回的 HTTP 状态码。`502` 通常意味着上游连接失败；WebSocket 成功升级记录为 `101` | +| `response_headers` | `String` | 否 | 响应头，结构同 `request_headers`。对于 SSE，会包含 `Content-Type: text/event-stream` 等 | +| `response_body` | `String` | 否 | 响应体内容。对于流式响应（SSE / chunked），这里保存的是**所有 chunk 拼接后的完整字节流**（包含 `data:` 前缀、`\n\n` 分隔符以及最后的 `[DONE]`），方便事后离线解析。最多保留 `MAX_BODY_BYTES` 字节 | +| `response_truncated` | `Bool` | 否 | 响应体是否被截断。截断只影响数据库存储，客户端始终收到完整数据 | +| `is_stream` | `Bool` | 否 | 是否 SSE 响应，仅接受媒体类型 `text/event-stream` | +| `latency_ms` | `Int64` | 否 | 端到端耗时（毫秒），从代理接收到请求到响应完成。对流式响应 = 从首请求到最后一个 chunk 发出 | +| `started_at` | `DateTime64(3)` | 否 | 代理接收到请求的时刻；ClickHouse 按 `toYYYYMM(started_at)` 月度分区 | +| `finished_at` | `DateTime64(3)` | 否 | 响应完全写回客户端（包括所有 chunk）的时刻 | +| `error` | `String` | 否 | 仅在反代过程中出现错误时填充。常见值：上游不可达、超时、读取请求体失败等 | + +### 常用查询示例 + +```sql +-- 查看最近 20 次失败请求 +SELECT started_at, path, status_code, error +FROM proxy_logs +WHERE status_code >= 400 OR error != '' +ORDER BY started_at DESC +LIMIT 20; + +-- 查看某次请求的完整内容 +SELECT + request_id, + method, path, + request_body AS req_text, + response_body AS resp_text, + latency_ms, is_stream +FROM proxy_logs +WHERE request_id = '0123456789abcdef0123456789abcdef'; + +-- 按模型统计调用量（从请求体里提取 JSON 字段） +SELECT + JSONExtractString(request_body, 'model') AS model, + count(*) AS calls, + toInt32(avg(latency_ms)) AS avg_ms +FROM proxy_logs +WHERE path = '/v1/chat/completions' + AND started_at > now() - INTERVAL 1 DAY +GROUP BY 1 +ORDER BY calls DESC; + +-- 查 Authorization（注意：敏感信息） +SELECT JSONExtractRaw(request_headers, 'Authorization') FROM proxy_logs LIMIT 5; +``` + +### 注意事项 + +- **敏感信息**：请求头中的 `Authorization`、`Cookie`、`X-Api-Key` 等**未脱敏**。如需脱敏请在 `proxy/capture.go` 的 `headersJSON` 中改造，或对数据库做列级权限控制。 +- **body 编码**：ClickHouse 以 `String` 保存 body 内容；文本接口可直接查询，二进制或压缩内容需按业务格式离线解析。 +- **WebSocket**：只记录 101 握手元数据，`response_body` 为空；WS 帧内容不采集，进程关闭时会关闭受管升级连接。 +- **截断**：`request_truncated` / `response_truncated` 为 `true` 时，对应 `*_body` 仅包含前 `MAX_BODY_BYTES` 字节。需保留完整内容请调高 `MAX_BODY_BYTES`，但要警惕数据库膨胀。 + +## 设计要点 + +- 请求体读取使用 `LimitReader`，超长仅记录前 `MAX_BODY_BYTES` 字节，下游仍能拿到完整 body。 +- `httputil.ReverseProxy` + `FlushInterval = -1`，自定义 `ResponseWriter` 同时实现 `Flusher`/`Hijacker`，写入时先转发再缓冲，保证流式实时性。 +- 日志通过非阻塞 channel 投递，队列满或 DB 不健康时直接丢弃（每 30 秒打印 metrics）。 +- DB 健康状态机：写入失败立即标记 unhealthy，后台 ping 恢复后重新启用。 +- ClickHouse `PrepareBatch` 失败会保留整批重试；逐项 `Append` 失败只保留明确失败项，已成功追加项继续发送。`Send` 返回错误时提交结果可能不明，系统不会自动重发该批，避免静默重复，并通过 `ambiguous_send` 计数暴露可能丢失；进程内 best-effort 队列不承诺分布式 exactly-once。 +- 启动会校验现有 `proxy_logs` 的列、引擎、分区键和排序键；不兼容 schema 会保持数据库 unhealthy，不自动重建数据表。 +- 所有 batch 路径都会执行清理；`Abort`/关闭失败会与原始错误合并记录，但不能使模糊提交变得可判定。 diff --git a/compose.yml b/compose.yml new file mode 100644 index 0000000..df26843 --- /dev/null +++ b/compose.yml @@ -0,0 +1,65 @@ +services: + tokenthief: + build: + context: . + dockerfile: Dockerfile + image: tokenthief:0.1.0 + container_name: tokenthief + restart: unless-stopped + depends_on: + thief_clickhouse: + condition: service_healthy + ports: + - "${LISTEN_PORT:-8080}:8080" + environment: + LISTEN_ADDR: ":8080" + UPSTREAM_URL: "${UPSTREAM_URL:?set UPSTREAM_URL}" + 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}" + LOG_QUEUE_SIZE: "${LOG_QUEUE_SIZE:-256}" + LOG_QUEUE_BYTES: "${LOG_QUEUE_BYTES:-67108864}" + LOG_BATCH_SIZE: "${LOG_BATCH_SIZE:-50}" + LOG_BATCH_INTERVAL: "${LOG_BATCH_INTERVAL:-2s}" + LOG_WORKERS: "${LOG_WORKERS:-2}" + FILTER_FILE: "/etc/tokenthief/filter.yaml" + DB_RECONNECT_INTERVAL: "${DB_RECONNECT_INTERVAL:-10s}" + READ_TIMEOUT: "${READ_TIMEOUT:-30s}" + WRITE_TIMEOUT: "${WRITE_TIMEOUT:-10m}" + IDLE_TIMEOUT: "${IDLE_TIMEOUT:-5m}" + UPSTREAM_TIMEOUT: "${UPSTREAM_TIMEOUT:-30s}" + UPSTREAM_RESPONSE_TIMEOUT: "${UPSTREAM_RESPONSE_TIMEOUT:-30s}" + UPSTREAM_STREAM_IDLE_TIMEOUT: "${UPSTREAM_STREAM_IDLE_TIMEOUT:-2m}" + TRUSTED_PROXIES: "${TRUSTED_PROXIES:-}" + TZ: "${TZ:-Asia/Shanghai}" + volumes: + - ./filter.yaml:/etc/tokenthief/filter.yaml:ro + networks: + - tokenthief + + thief_clickhouse: + image: clickhouse/clickhouse-server:25.3.3.42-alpine + container_name: tokenthief-clickhouse + restart: unless-stopped + environment: + CLICKHOUSE_USER: "${CLICKHOUSE_USER:-tokenthief}" + CLICKHOUSE_PASSWORD: "${CLICKHOUSE_PASSWORD:?set a strong CLICKHOUSE_PASSWORD}" + CLICKHOUSE_DB: "${CLICKHOUSE_DB:-tokenthief}" + TZ: "${TZ:-Asia/Shanghai}" + volumes: + - clickhouse_data:/var/lib/clickhouse + healthcheck: + test: ["CMD-SHELL", "clickhouse-client --user \"$${CLICKHOUSE_USER}\" --password \"$${CLICKHOUSE_PASSWORD}\" --database \"$${CLICKHOUSE_DB}\" --query 'SELECT 1'"] + interval: 10s + timeout: 5s + retries: 5 + start_period: 10s + networks: + - tokenthief + +volumes: + clickhouse_data: + +networks: + tokenthief: + driver: bridge diff --git a/config/config.go b/config/config.go new file mode 100644 index 0000000..cab6219 --- /dev/null +++ b/config/config.go @@ -0,0 +1,245 @@ +package config + +import ( + "errors" + "fmt" + "net/netip" + "net/url" + "os" + "strconv" + "strings" + "time" + + "git.misaka.ren/M1saka/token_thief/db" +) + +// Config 保存从环境变量加载的全部运行配置。 +type Config struct { + ListenAddr string + UpstreamURL *url.URL + ClickHouseURL string + MaxBodyBytes int64 + LogQueueSize int + LogQueueBytes int64 + LogBatchSize int + LogBatchInterval time.Duration + LogWorkers int + FilterFile string + DBReconnectInterval time.Duration + ReadTimeout time.Duration + WriteTimeout time.Duration + IdleTimeout time.Duration + UpstreamTimeout time.Duration + UpstreamResponseTimeout time.Duration + UpstreamStreamIdleTimeout time.Duration + UpstreamTLSInsecureSkipVerify bool + TrustedProxies []netip.Prefix +} + +// Load 从进程环境读取配置；缺少必填字段时返回错误。 +func Load() (*Config, error) { + maxBodyBytes, err := getEnvInt64("MAX_BODY_BYTES", 1<<20) + if err != nil { + return nil, err + } + logQueueSize, err := getEnvInt("LOG_QUEUE_SIZE", 256) + if err != nil { + return nil, err + } + logQueueBytes, err := getEnvInt64("LOG_QUEUE_BYTES", 64<<20) + if err != nil { + return nil, err + } + logBatchSize, err := getEnvInt("LOG_BATCH_SIZE", 50) + if err != nil { + return nil, err + } + logBatchInterval, err := getEnvDuration("LOG_BATCH_INTERVAL", 2*time.Second) + if err != nil { + return nil, err + } + logWorkers, err := getEnvInt("LOG_WORKERS", 2) + if err != nil { + return nil, err + } + dbReconnectInterval, err := getEnvDuration("DB_RECONNECT_INTERVAL", 10*time.Second) + if err != nil { + return nil, err + } + readTimeout, err := getEnvDuration("READ_TIMEOUT", 30*time.Second) + if err != nil { + return nil, err + } + writeTimeout, err := getEnvDuration("WRITE_TIMEOUT", 10*time.Minute) + if err != nil { + return nil, err + } + idleTimeout, err := getEnvDuration("IDLE_TIMEOUT", 5*time.Minute) + if err != nil { + return nil, err + } + upstreamTimeout, err := getEnvDuration("UPSTREAM_TIMEOUT", 30*time.Second) + if err != nil { + return nil, err + } + upstreamResponseTimeout, err := getEnvDuration("UPSTREAM_RESPONSE_TIMEOUT", 30*time.Second) + if err != nil { + return nil, err + } + upstreamStreamIdleTimeout, err := getEnvDuration("UPSTREAM_STREAM_IDLE_TIMEOUT", 2*time.Minute) + if err != nil { + return nil, err + } + upstreamTLSInsecureSkipVerify, err := getEnvBool("UPSTREAM_TLS_INSECURE_SKIP_VERIFY", false) + if err != nil { + return nil, err + } + trustedProxies, err := getEnvPrefixes("TRUSTED_PROXIES") + if err != nil { + return nil, err + } + + cfg := &Config{ + ListenAddr: getEnv("LISTEN_ADDR", ":8080"), + ClickHouseURL: os.Getenv("CLICKHOUSE_URL"), + MaxBodyBytes: maxBodyBytes, + LogQueueSize: logQueueSize, + LogQueueBytes: logQueueBytes, + LogBatchSize: logBatchSize, + LogBatchInterval: logBatchInterval, + LogWorkers: logWorkers, + FilterFile: getEnv("FILTER_FILE", "./filter.yaml"), + DBReconnectInterval: dbReconnectInterval, + ReadTimeout: readTimeout, + WriteTimeout: writeTimeout, + IdleTimeout: idleTimeout, + UpstreamTimeout: upstreamTimeout, + UpstreamResponseTimeout: upstreamResponseTimeout, + UpstreamStreamIdleTimeout: upstreamStreamIdleTimeout, + UpstreamTLSInsecureSkipVerify: upstreamTLSInsecureSkipVerify, + TrustedProxies: trustedProxies, + } + + upstream := os.Getenv("UPSTREAM_URL") + if upstream == "" { + return nil, errors.New("UPSTREAM_URL is required") + } + u, err := url.Parse(upstream) + if err != nil { + return nil, fmt.Errorf("invalid UPSTREAM_URL: %w", err) + } + if u.Scheme == "" || u.Host == "" { + return nil, fmt.Errorf("invalid UPSTREAM_URL: %q", upstream) + } + cfg.UpstreamURL = u + + if cfg.ClickHouseURL == "" { + return nil, errors.New("CLICKHOUSE_URL is required") + } + if _, err := db.ClickHouseOptions(cfg.ClickHouseURL); err != nil { + return nil, err + } + positiveValues := []struct { + key string + value int64 + }{ + {key: "MAX_BODY_BYTES", value: cfg.MaxBodyBytes}, + {key: "LOG_QUEUE_SIZE", value: int64(cfg.LogQueueSize)}, + {key: "LOG_QUEUE_BYTES", value: cfg.LogQueueBytes}, + {key: "LOG_BATCH_SIZE", value: int64(cfg.LogBatchSize)}, + {key: "LOG_BATCH_INTERVAL", value: int64(cfg.LogBatchInterval)}, + {key: "LOG_WORKERS", value: int64(cfg.LogWorkers)}, + {key: "DB_RECONNECT_INTERVAL", value: int64(cfg.DBReconnectInterval)}, + {key: "READ_TIMEOUT", value: int64(cfg.ReadTimeout)}, + {key: "WRITE_TIMEOUT", value: int64(cfg.WriteTimeout)}, + {key: "IDLE_TIMEOUT", value: int64(cfg.IdleTimeout)}, + {key: "UPSTREAM_TIMEOUT", value: int64(cfg.UpstreamTimeout)}, + {key: "UPSTREAM_RESPONSE_TIMEOUT", value: int64(cfg.UpstreamResponseTimeout)}, + {key: "UPSTREAM_STREAM_IDLE_TIMEOUT", value: int64(cfg.UpstreamStreamIdleTimeout)}, + } + for _, item := range positiveValues { + if item.value <= 0 { + return nil, fmt.Errorf("%s must be greater than zero", item.key) + } + } + return cfg, nil +} + +func getEnv(key, def string) string { + if v := os.Getenv(key); v != "" { + return v + } + return def +} + +func getEnvInt(key string, def int) (int, error) { + if v := os.Getenv(key); v != "" { + n, err := strconv.Atoi(v) + if err != nil { + return 0, fmt.Errorf("invalid %s: %w", key, err) + } + return n, nil + } + return def, nil +} + +func getEnvInt64(key string, def int64) (int64, error) { + if v := os.Getenv(key); v != "" { + n, err := strconv.ParseInt(v, 10, 64) + if err != nil { + return 0, fmt.Errorf("invalid %s: %w", key, err) + } + return n, nil + } + return def, nil +} + +func getEnvDuration(key string, def time.Duration) (time.Duration, error) { + if v := os.Getenv(key); v != "" { + d, err := time.ParseDuration(v) + if err != nil { + return 0, fmt.Errorf("invalid %s: %w", key, err) + } + return d, nil + } + return def, nil +} + +func getEnvBool(key string, def bool) (bool, error) { + if v := os.Getenv(key); v != "" { + b, err := strconv.ParseBool(v) + if err != nil { + return false, fmt.Errorf("invalid %s: %w", key, err) + } + return b, nil + } + return def, nil +} + +func getEnvPrefixes(key string) ([]netip.Prefix, error) { + v := os.Getenv(key) + if v == "" { + return nil, nil + } + + parts := strings.Split(v, ",") + prefixes := make([]netip.Prefix, 0, len(parts)) + for _, part := range parts { + value := strings.TrimSpace(part) + if value == "" { + return nil, fmt.Errorf("invalid %s: empty IP or CIDR", key) + } + + if addr, err := netip.ParseAddr(value); err == nil { + prefixes = append(prefixes, netip.PrefixFrom(addr, addr.BitLen())) + continue + } + + prefix, err := netip.ParsePrefix(value) + if err != nil || !prefix.IsValid() || prefix != prefix.Masked() { + return nil, fmt.Errorf("invalid %s entry %q", key, value) + } + prefixes = append(prefixes, prefix) + } + return prefixes, nil +} diff --git a/config/filter.go b/config/filter.go new file mode 100644 index 0000000..c2be411 --- /dev/null +++ b/config/filter.go @@ -0,0 +1,134 @@ +package config + +import ( + "fmt" + "os" + "regexp" + "strings" + + "gopkg.in/yaml.v3" +) + +// FilterMode 决定如何应用 patterns。 +type FilterMode string + +const ( + FilterWhitelist FilterMode = "whitelist" + FilterBlacklist FilterMode = "blacklist" + FilterDisabled FilterMode = "disabled" +) + +// Filter 表示路径匹配规则。 +type Filter struct { + Mode FilterMode `yaml:"mode"` + Patterns []string `yaml:"patterns"` + + compiled []*regexp.Regexp +} + +// LoadFilter 从 yaml 文件加载过滤配置；文件不存在则默认全部记录。 +// +// Pattern 语法（glob 风格）： +// - `*` 匹配单个路径段内除 `/` 之外的任意字符（包括零个）。 +// - `**` 匹配任意字符，含 `/`，可跨段。 +// - `?` 匹配单个非 `/` 字符。 +// - 其它字符按字面匹配。 +// +// 示例： +// - `/v1/audio/*` 匹配 /v1/audio/speech、/v1/audio/transcriptions +// - `/v1/videos/**` 匹配 /v1/videos/任意子路径 +// - `/v1beta/models/*:generateContent` 匹配 Gemini 风格端点 +func LoadFilter(filePath string) (*Filter, error) { + data, err := os.ReadFile(filePath) + if err != nil { + if os.IsNotExist(err) { + return &Filter{Mode: FilterDisabled}, nil + } + return nil, err + } + var f Filter + if err := yaml.Unmarshal(data, &f); err != nil { + return nil, fmt.Errorf("parse filter yaml: %w", err) + } + switch f.Mode { + case FilterWhitelist, FilterBlacklist, FilterDisabled: + case "": + f.Mode = FilterDisabled + default: + return nil, fmt.Errorf("unknown filter mode: %q", f.Mode) + } + if err := f.compile(); err != nil { + return nil, err + } + return &f, nil +} + +// NewFilter 程序化构造一个 Filter（主要供测试使用）。 +func NewFilter(mode FilterMode, patterns []string) (*Filter, error) { + f := &Filter{Mode: mode, Patterns: append([]string(nil), patterns...)} + if err := f.compile(); err != nil { + return nil, err + } + return f, nil +} + +func (f *Filter) compile() error { + f.compiled = make([]*regexp.Regexp, 0, len(f.Patterns)) + for _, p := range f.Patterns { + re, err := CompileGlob(p) + if err != nil { + return fmt.Errorf("invalid pattern %q: %w", p, err) + } + f.compiled = append(f.compiled, re) + } + return nil +} + +// ShouldLog 决定一个请求 path 是否需要被记录。 +func (f *Filter) ShouldLog(reqPath string) bool { + if f == nil || f.Mode == FilterDisabled { + return true + } + matched := false + for _, re := range f.compiled { + if re.MatchString(reqPath) { + matched = true + break + } + } + switch f.Mode { + case FilterWhitelist: + return matched + case FilterBlacklist: + return !matched + default: + return true + } +} + +// CompileGlob 将 glob 风格 pattern 转换为 anchored 正则表达式。 +func CompileGlob(pattern string) (*regexp.Regexp, error) { + var sb strings.Builder + sb.WriteString("^") + for i := 0; i < len(pattern); i++ { + c := pattern[i] + switch c { + case '*': + if i+1 < len(pattern) && pattern[i+1] == '*' { + sb.WriteString(".*") + i++ + } else { + sb.WriteString("[^/]*") + } + case '?': + sb.WriteString("[^/]") + case '.', '+', '(', ')', '|', '^', '$', '{', '}', '[', ']', '\\': + sb.WriteByte('\\') + sb.WriteByte(c) + default: + sb.WriteByte(c) + } + } + sb.WriteString("$") + return regexp.Compile(sb.String()) +} diff --git a/db/clickhouse.go b/db/clickhouse.go new file mode 100644 index 0000000..932b26a --- /dev/null +++ b/db/clickhouse.go @@ -0,0 +1,221 @@ +package db + +import ( + "context" + "fmt" + "log" + "net" + "net/url" + "strconv" + "strings" + "sync" + "time" + + "github.com/ClickHouse/clickhouse-go/v2" +) + +// Pool 封装 ClickHouse 连接并维护健康状态，DB 故障时不阻塞调用方。 +type Pool struct { + dsn string + reconnectInterval time.Duration + mu sync.RWMutex + conn clickhouse.Conn + healthy bool +} + +// NewPool 创建 Pool。即使首次连接失败也返回非 nil 实例，后台会持续重试。 +func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *Pool { + p := &Pool{dsn: dsn, reconnectInterval: reconnectInterval} + if err := p.connect(ctx); err != nil { + log.Printf("[db] initial connect failed: %v (service continues without DB)", err) + } + go p.watch(ctx) + return p +} + +func (p *Pool) connect(ctx context.Context) error { + cctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + opts, err := ClickHouseOptions(p.dsn) + if err != nil { + return err + } + conn, err := clickhouse.Open(opts) + if err != nil { + return err + } + if err := conn.Ping(cctx); err != nil { + _ = conn.Close() + return err + } + // 先 migrate，再 swap：避免新连接 migrate 失败时取代掉旧的可用连接。 + mctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + if err := migrate(mctx, conn); err != nil { + _ = conn.Close() + log.Printf("[db] migrate failed: %v", err) + return err + } + p.mu.Lock() + if p.conn != nil { + _ = p.conn.Close() + } + p.conn = conn + p.healthy = true + p.mu.Unlock() + log.Printf("[db] connected and migrated") + return nil +} + +// ClickHouseOptions converts the supported CLICKHOUSE_URL subset into driver options. +func ClickHouseOptions(raw string) (*clickhouse.Options, error) { + if !strings.Contains(raw, "://") { + if err := validateHostPort(raw); err != nil { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) + } + return &clickhouse.Options{Addr: []string{raw}}, nil + } + + u, err := url.Parse(raw) + if err != nil { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) + } + if u.Scheme != "clickhouse" && u.Scheme != "clickhouses" { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: unsupported scheme %q", u.Scheme) + } + if u.Fragment != "" { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: fragment is not allowed") + } + if err := validateHostPort(u.Host); err != nil { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) + } + + query, err := url.ParseQuery(u.RawQuery) + if err != nil { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: query: %w", err) + } + for key, values := range query { + switch key { + case "secure", "skip_verify", "compress": + default: + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: unsupported parameter %q", key) + } + if len(values) != 1 { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: parameter %q must occur once", key) + } + } + + wantSecure := u.Scheme == "clickhouses" + if value, ok := query["secure"]; ok { + secure, err := strconv.ParseBool(value[0]) + if err != nil { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: secure: %w", err) + } + if secure != wantSecure { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: secure conflicts with %s", u.Scheme) + } + } + if _, ok := query["skip_verify"]; ok && !wantSecure { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: skip_verify requires clickhouses") + } + if value, ok := query["skip_verify"]; ok && value[0] != "" { + if _, err := strconv.ParseBool(value[0]); err != nil { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: skip_verify: %w", err) + } + } + if value, ok := query["compress"]; ok { + switch value[0] { + case "true", "false", "none", "zstd", "lz4", "lz4hc", "gzip", "deflate", "br": + default: + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: unsupported compression %q", value[0]) + } + } + + // Let the driver decode userinfo/path and interpret compression and TLS values. + if wantSecure { + query.Set("secure", "true") + u.RawQuery = query.Encode() + } + opts, err := clickhouse.ParseDSN(u.String()) + if err != nil { + return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) + } + return opts, nil +} + +func validateHostPort(address string) error { + if address == "" { + return fmt.Errorf("missing host and port") + } + host, port, err := net.SplitHostPort(address) + if err != nil { + return fmt.Errorf("expected host:port: %w", err) + } + if host == "" { + return fmt.Errorf("missing host") + } + n, err := strconv.ParseUint(port, 10, 16) + if err != nil || n == 0 { + return fmt.Errorf("invalid port %q", port) + } + return nil +} + +// watch 定期探活；不健康时尝试重连。 +func (p *Pool) watch(ctx context.Context) { + t := time.NewTicker(p.reconnectInterval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + if p.Healthy() { + if conn := p.Get(); 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() + } + cancel() + } + continue + } + if err := p.connect(ctx); err != nil { + log.Printf("[db] reconnect failed: %v", err) + } + } + } +} + +// Healthy 报告连接是否可用。 +func (p *Pool) Healthy() bool { + p.mu.RLock() + defer p.mu.RUnlock() + return p.healthy +} + +// MarkUnhealthy 由调用方在写入失败后调用。 +func (p *Pool) MarkUnhealthy() { + p.mu.Lock() + p.healthy = false + p.mu.Unlock() +} + +// Get 返回当前连接，可能为 nil。 +func (p *Pool) Get() clickhouse.Conn { + p.mu.RLock() + defer p.mu.RUnlock() + return p.conn +} + +// Close 释放底层连接。 +func (p *Pool) Close() { + p.mu.Lock() + defer p.mu.Unlock() + if p.conn != nil { + _ = p.conn.Close() + p.conn = nil + } + p.healthy = false +} diff --git a/db/clickhouse_test.go b/db/clickhouse_test.go new file mode 100644 index 0000000..43d6705 --- /dev/null +++ b/db/clickhouse_test.go @@ -0,0 +1,127 @@ +package db + +import ( + "reflect" + "testing" +) + +func TestClickHouseOptionsHostPort(t *testing.T) { + opts, err := ClickHouseOptions("localhost:9000") + if err != nil { + t.Fatalf("ClickHouseOptions: %v", err) + } + if len(opts.Addr) != 1 || opts.Addr[0] != "localhost:9000" { + t.Fatalf("Addr=%v want [localhost:9000]", opts.Addr) + } + if opts.Auth.Username != "" || opts.Auth.Password != "" || opts.Auth.Database != "" { + t.Fatalf("Auth=%+v want empty", opts.Auth) + } +} + +func TestClickHouseOptionsURLWithAuthAndDatabase(t *testing.T) { + opts, err := ClickHouseOptions("clickhouse://user%40name:p%2Fass@localhost:9000/token%20thief?compress=zstd") + if err != nil { + t.Fatalf("ClickHouseOptions: %v", err) + } + if len(opts.Addr) != 1 || opts.Addr[0] != "localhost:9000" { + t.Fatalf("Addr=%v want [localhost:9000]", opts.Addr) + } + if opts.Auth.Username != "user@name" { + t.Fatalf("Username=%q want user@name", opts.Auth.Username) + } + if opts.Auth.Password != "p/ass" { + t.Fatalf("Password=%q want p/ass", opts.Auth.Password) + } + if opts.Auth.Database != "token thief" { + t.Fatalf("Database=%q want token thief", opts.Auth.Database) + } + if opts.Compression == nil { + t.Fatal("Compression=nil want enabled") + } +} + +func TestClickHouseOptionsURLWithDatabaseOnly(t *testing.T) { + opts, err := ClickHouseOptions("clickhouse://localhost:9000/tokenthief") + if err != nil { + t.Fatalf("ClickHouseOptions: %v", err) + } + if len(opts.Addr) != 1 || opts.Addr[0] != "localhost:9000" { + t.Fatalf("Addr=%v want [localhost:9000]", opts.Addr) + } + if opts.Auth.Database != "tokenthief" { + t.Fatalf("Database=%q want tokenthief", opts.Auth.Database) + } +} + +func TestClickHouseOptionsTLS(t *testing.T) { + opts, err := ClickHouseOptions("clickhouses://localhost:9440/tokenthief?skip_verify=true") + if err != nil { + t.Fatalf("ClickHouseOptions: %v", err) + } + if opts.TLS == nil || !opts.TLS.InsecureSkipVerify { + t.Fatalf("TLS=%+v want InsecureSkipVerify", opts.TLS) + } +} + +func TestClickHouseOptionsRejectsInvalidDSN(t *testing.T) { + tests := []string{ + "localhost", + "localhost:http", + "clickhouse://localhost/tokenthief", + "clickhouse:///tokenthief", + "clickhouse://localhost:9000/db#fragment", + "clickhouse://localhost:9000/db?unknown=true", + "clickhouse://localhost:9000/db?compress=snappy", + "clickhouse://localhost:9000/db?compress=lz4&compress=zstd", + "clickhouse://localhost:9000/db?compress=%zz", + "clickhouse://localhost:9000/db?secure=true", + "clickhouses://localhost:9440/db?secure=false", + "clickhouses://localhost:9440/db?skip_verify=maybe", + "clickhouse://localhost:9000/db?skip_verify=true", + "https://localhost:9440/db", + } + for _, dsn := range tests { + t.Run(dsn, func(t *testing.T) { + if _, err := ClickHouseOptions(dsn); err == nil { + t.Fatalf("ClickHouseOptions(%q) succeeded", dsn) + } + }) + } +} + +func TestValidateProxyLogsSchema(t *testing.T) { + columns := append([]schemaColumn(nil), proxyLogsColumns...) + table := schemaTable{ + engine: "MergeTree", + partitionKey: "toYYYYMM(started_at)", + sortingKey: "started_at, request_id", + } + if err := validateProxyLogsSchema(columns, table); err != nil { + t.Fatalf("validateProxyLogsSchema: %v", err) + } + + badColumns := append([]schemaColumn(nil), columns...) + badColumns[8].typ = "UInt16" + tests := []struct { + name string + columns []schemaColumn + table schemaTable + }{ + {"missing column", columns[:16], table}, + {"wrong type", badColumns, table}, + {"wrong engine", columns, schemaTable{"ReplacingMergeTree", table.partitionKey, table.sortingKey}}, + {"wrong partition", columns, schemaTable{table.engine, "toDate(started_at)", table.sortingKey}}, + {"wrong sorting", columns, schemaTable{table.engine, table.partitionKey, "request_id"}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if err := validateProxyLogsSchema(test.columns, test.table); err == nil { + t.Fatal("validateProxyLogsSchema succeeded") + } + }) + } + + if !reflect.DeepEqual(columns, proxyLogsColumns) { + t.Fatal("schema validation mutated columns") + } +} diff --git a/db/migrate.go b/db/migrate.go new file mode 100644 index 0000000..abde6a1 --- /dev/null +++ b/db/migrate.go @@ -0,0 +1,134 @@ +package db + +import ( + "context" + "fmt" + "strings" + + "github.com/ClickHouse/clickhouse-go/v2" +) + +const schemaSQL = ` +CREATE TABLE IF NOT EXISTS proxy_logs ( + request_id String, + method String, + path String, + query String, + client_ip String, + request_headers String, + request_body String, + request_truncated Bool DEFAULT false, + status_code Int32, + response_headers String, + response_body String, + response_truncated Bool DEFAULT false, + is_stream Bool DEFAULT false, + latency_ms Int64, + started_at DateTime64(3), + finished_at DateTime64(3), + error String +) +ENGINE = MergeTree +PARTITION BY toYYYYMM(started_at) +ORDER BY (started_at, request_id) +` + +func migrate(ctx context.Context, conn clickhouse.Conn) error { + if err := conn.Exec(ctx, schemaSQL); err != nil { + return fmt.Errorf("create proxy_logs: %w", err) + } + + rows, err := conn.Query(ctx, ` +SELECT name, type +FROM system.columns +WHERE database = currentDatabase() AND table = 'proxy_logs' +ORDER BY position`) + if err != nil { + return fmt.Errorf("query proxy_logs columns: %w", err) + } + var columns []schemaColumn + for rows.Next() { + var column schemaColumn + if err := rows.Scan(&column.name, &column.typ); err != nil { + rows.Close() + return fmt.Errorf("scan proxy_logs columns: %w", err) + } + columns = append(columns, column) + } + if err := rows.Err(); err != nil { + rows.Close() + return fmt.Errorf("read proxy_logs columns: %w", err) + } + rows.Close() + + var table schemaTable + err = conn.QueryRow(ctx, ` +SELECT engine, partition_key, sorting_key +FROM system.tables +WHERE database = currentDatabase() AND name = 'proxy_logs'`).Scan( + &table.engine, &table.partitionKey, &table.sortingKey, + ) + if err != nil { + return fmt.Errorf("query proxy_logs table: %w", err) + } + if err := validateProxyLogsSchema(columns, table); err != nil { + return fmt.Errorf("incompatible proxy_logs schema: %w", err) + } + return nil +} + +type schemaColumn struct { + name string + typ string +} + +type schemaTable struct { + engine string + partitionKey string + sortingKey string +} + +var proxyLogsColumns = []schemaColumn{ + {"request_id", "String"}, + {"method", "String"}, + {"path", "String"}, + {"query", "String"}, + {"client_ip", "String"}, + {"request_headers", "String"}, + {"request_body", "String"}, + {"request_truncated", "Bool"}, + {"status_code", "Int32"}, + {"response_headers", "String"}, + {"response_body", "String"}, + {"response_truncated", "Bool"}, + {"is_stream", "Bool"}, + {"latency_ms", "Int64"}, + {"started_at", "DateTime64(3)"}, + {"finished_at", "DateTime64(3)"}, + {"error", "String"}, +} + +func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error { + if len(columns) != len(proxyLogsColumns) { + return fmt.Errorf("got %d columns, want %d", len(columns), len(proxyLogsColumns)) + } + for i, want := range proxyLogsColumns { + if columns[i] != want { + return fmt.Errorf("column %d is %s %s, want %s %s", i+1, columns[i].name, columns[i].typ, want.name, want.typ) + } + } + if table.engine != "MergeTree" { + return fmt.Errorf("engine is %q, want MergeTree", table.engine) + } + if compactExpression(table.partitionKey) != "toYYYYMM(started_at)" { + return fmt.Errorf("partition key is %q, want toYYYYMM(started_at)", table.partitionKey) + } + if compactExpression(table.sortingKey) != "started_at,request_id" { + return fmt.Errorf("sorting key is %q, want started_at, request_id", table.sortingKey) + } + return nil +} + +func compactExpression(value string) string { + return strings.Join(strings.Fields(value), "") +} diff --git a/docs/compose/plans/2026-07-09-clickhouse-migration.md b/docs/compose/plans/2026-07-09-clickhouse-migration.md new file mode 100644 index 0000000..a3522c9 --- /dev/null +++ b/docs/compose/plans/2026-07-09-clickhouse-migration.md @@ -0,0 +1,66 @@ +# ClickHouse Migration 安全修复计划 + +> [!NOTE] +> This document may not reflect the current implementation. +> See the final report for up-to-date state: +> [Final Report](../reports/reliability-security-fixes.md) + +**变更规模：** 大型跨模块修复。现有迁移主体复用，仅重新实施受本次审查影响的任务及其集成依赖。 + +## 全局约束 + +- 保留非阻塞 best-effort 日志语义，并通过 Git 提交保留可审查的变更证据。 +- 先用失败测试固定根因，再实施最小修复。 +- `Send` 模糊失败不重试；不宣称分布式 exactly-once。 +- WebSocket 只管理连接并记录 101 元数据，不采集帧。 + +### Task 1: 严格配置与安全默认 + +**文件：** `config/config.go`、`tests/config/config_test.go` + +- [ ] 严格解析整数、布尔值和 duration，拒绝非正 body/队列/timeout。 +- [ ] 增加 `LOG_QUEUE_BYTES`、响应总/idle timeout、`TRUSTED_PROXIES`。 +- [ ] 在配置加载阶段校验 ClickHouse DSN，并补齐边界测试。 + +### Task 2: Queue 生命周期、字节预算与确定性写入 + +**文件：** `logger/queue.go`、`tests/logger/queue_test.go` + +- [ ] 用同步状态机消除 Submit/Stop send-close 竞态，Stop 使用单一 context 总预算。 +- [ ] 增加条数/字节双预算并在所有消费、丢弃和 shutdown 路径释放预留。 +- [ ] 区分 Prepare、Append、Send 错误；Append 失败清理 batch，Send 模糊失败不重试。 +- [ ] 增加并发关闭、预算、清理、部分失败和模糊提交测试。 + +### Task 3: ClickHouse DSN 与 schema 验证 + +**文件：** `db/clickhouse.go`、`db/migrate.go`、`db/clickhouse_test.go` + +- [ ] 严格解析 scheme、TLS 和白名单 query 参数。 +- [ ] 创建 schema 后校验列、引擎、分区和排序键；不兼容时保持 unhealthy。 +- [ ] 增加 DSN 与 schema 元数据校验测试。 + +### Task 4: HTTP/SSE 代理完整性与可信来源 + +**文件：** `proxy/proxy.go`、`proxy/writer.go`、`proxy/capture.go`、`proxy/sse.go`、`tests/proxy/*` + +- [ ] 请求体读取失败 fail closed；502 使用固定客户端文本。 +- [ ] 记录响应写错误/短写，响应中断不提交完整日志。 +- [ ] 仅对真正 SSE 按完整事件检测终止，并等待正常返回后提交。 +- [ ] 为普通响应总 timeout 与 SSE idle timeout 包装 upstream Body。 +- [ ] 默认忽略转发头，仅按可信代理链提取客户端 IP。 + +### Task 5: WebSocket 元数据与统一 shutdown + +**文件：** `proxy/writer.go`、`proxy/proxy.go`、`tests/proxy/websocket_test.go`、`main.go` + +- [ ] Hijack 成功后登记连接，记录 101 握手元数据并在结束时注销。 +- [ ] 提供 Handler shutdown，关闭受管升级连接。 +- [ ] server 启动错误和 signal 共用清理路径，所有关闭步骤受单一总预算约束。 + +### Task 6: Compose、文档与最终验证 + +**文件：** `compose.yml`、`.env.example`、`README.md`、`docs/compose/reports/*` + +- [ ] 固定镜像、取消默认 ClickHouse 端口发布、强制显式密码/DSN并同步新配置。 +- [ ] 记录 best-effort、模糊提交、schema 和 WebSocket 取舍。 +- [ ] 运行 `gofmt`、`go test ./...`、`go vet ./...` 并生成报告。 diff --git a/docs/compose/plans/reliability-security-fixes.md b/docs/compose/plans/reliability-security-fixes.md new file mode 100644 index 0000000..5388f5d --- /dev/null +++ b/docs/compose/plans/reliability-security-fixes.md @@ -0,0 +1,5 @@ +# Reliability And Security Fixes Plan + +The canonical dated plan is [2026-07-09-clickhouse-migration.md](2026-07-09-clickhouse-migration.md). + +This stable path is the review entry point for the reliability and security fixes. diff --git a/docs/compose/reports/2026-07-10-clickhouse-migration-security-fixes.md b/docs/compose/reports/2026-07-10-clickhouse-migration-security-fixes.md new file mode 100644 index 0000000..a4b9a46 --- /dev/null +++ b/docs/compose/reports/2026-07-10-clickhouse-migration-security-fixes.md @@ -0,0 +1,32 @@ +# TokenThief 可靠性与安全修复验证报告 + +## 结果 + +本次修订覆盖队列并发关闭和单一总 shutdown deadline、响应完整性、请求体 fail-closed、ClickHouse 部分/模糊提交、64 MiB 默认队列字节预算、严格正值配置、响应体总/idle timeout、SSE 协议边界、WebSocket 101 与 shutdown、DSN/TLS、schema 校验、Compose 安全默认、固定 502、可信代理、batch 清理和 server 启动统一清理。 + +Compose 使用固定 ClickHouse 镜像 `clickhouse/clickhouse-server:25.3.3.42-alpine`，默认不发布 ClickHouse 端口，并在配置展开阶段拒绝空的 `UPSTREAM_URL`、`CLICKHOUSE_URL` 和 `CLICKHOUSE_PASSWORD`。`.env.example` 不提供密码或 DSN 默认值。 + +## 关键取舍 + +- 日志仍是内存中、非阻塞、best-effort。队列满、字节预算不足或数据库不可用时允许丢弃并计数。 +- ClickHouse `PrepareBatch` 的确定失败保留整批；逐项 `Append` 的确定失败只保留失败项，成功项照常发送。`Send` 错误无法从单机客户端确认服务端是否已提交，因此不自动重试该批并计入 `ambiguous_send`，优先避免静默重复。跨进程 exactly-once 需要持久化 outbox 和服务端幂等协议，未在本次引入。 +- WebSocket 仅记录 101 握手元数据并关闭受管连接，不解析帧或承诺跨进程迁移。 +- schema 会自动创建缺失表，并校验现有表的列类型、MergeTree 引擎、分区键和排序键；不兼容时拒绝标记健康，不执行有数据风险的自动重建。 +- Compose 不默认发布 ClickHouse 端口，强制显式密码和独立 DSN，以同时保留原始服务端密码和 URL 编码凭据。 + +## 行为边界 + +- 仅当请求体完整可重放且 `ReverseProxy` 正常结束时提交 HTTP 日志；读取失败的请求不转发，响应中断不提交不完整日志。 +- SSE 仅由 `Content-Type: text/event-stream` 判定；看到 `[DONE]` 不会提前提交，仍等待上游正常结束。普通响应使用总 timeout，SSE 使用可重置 idle timeout。 +- WebSocket 仅记录成功 `101` 的握手元数据，不采集帧；进程关闭会关闭当前进程管理的升级连接，但不提供跨进程连接迁移或协调。 +- `502` 对客户端使用固定错误文本，内部连接错误只进入服务端日志。转发客户端 IP 仅在 TCP 对端属于 `TRUSTED_PROXIES` 时生效。 + +## 验证 + +- `gofmt -w .`：通过。 +- `go test ./...`：通过。 +- `go vet ./...`：通过。 +- 新增 `tests/deployment` 契约测试，覆盖必填部署参数、固定 ClickHouse 镜像、不发布数据库端口和示例凭据留空。 +- `docker compose config`：未执行，当前环境没有可用 Docker daemon/CLI；Compose 安全约束由 Go 契约测试覆盖。 +- `go test -race`：不属于验收命令，未执行。 +- 未连接真实 ClickHouse 做断链模糊提交和旧 schema 集成测试；相关阶段语义通过 fake batch 与纯 schema 校验单测覆盖。 diff --git a/docs/compose/reports/reliability-security-fixes.md b/docs/compose/reports/reliability-security-fixes.md new file mode 100644 index 0000000..d562aff --- /dev/null +++ b/docs/compose/reports/reliability-security-fixes.md @@ -0,0 +1,55 @@ +--- +feature: reliability-security-fixes +status: delivered +specs: + - docs/compose/specs/reliability-security-fixes.md + - docs/compose/specs/2026-07-09-clickhouse-migration.md +plans: + - docs/compose/plans/reliability-security-fixes.md + - docs/compose/plans/2026-07-09-clickhouse-migration.md +branch: main +--- + +# 可靠性与安全修复 - 最终报告 + +## What Was Built + +本轮完成了 token_thief 的可靠性与安全加固。异步日志队列现在安全处理并发 `Submit`/`Stop`、使用条目数和字节双预算、在统一 shutdown deadline 内排空，并区分 ClickHouse 的可安全重试、确定失败和提交结果不明三类写入结果。 + +反向代理现在对请求体读取、上游响应复制、普通响应超时、SSE idle timeout、WebSocket 101 元数据和升级连接关闭实施完整性保护。客户端错误文本不再泄露内部信息，转发来源头仅在直接对端属于 `TRUSTED_PROXIES` 时参与客户端 IP 判定。 + +配置、ClickHouse DSN/TLS、数据库 schema 和 Compose 部署均采用 fail-closed 校验与更安全默认；服务启动错误和信号关闭共用资源清理路径。 + +## Architecture + +`config/config.go` 在启动前严格解析正值预算、duration、布尔值、可信代理和 DSN。`proxy/` 在转发前完整读取请求体，通过 capture writer 和 response body wrapper 判断响应是否完整，并管理 hijacked 连接。`logger/queue.go` 提供非阻塞有界队列，按 `PrepareBatch`、`Append`、`Send` 阶段决定重试或丢弃；`db/` 只在连接、迁移和 schema 校验均成功后标记健康。`main.go` 用一个 30 秒 shutdown context 依次关闭 HTTP、升级连接、队列和数据库。 + +### Design Decisions + +- 选择仅重试 `PrepareBatch` 失败，因为此时可确定没有提交；`Send` 失败计为 ambiguous 且不重放，避免静默重复。 +- 选择进程内有界 best-effort 队列，因为当前项目没有持久化 outbox；该策略明确不承诺跨进程 exactly-once。 +- 选择 WebSocket 仅记录 101 握手元数据并管理连接，不采集帧，以保持代理边界和关闭行为可测试。 +- 选择拒绝不兼容 ClickHouse schema，而不是自动破坏性迁移或重建表。 + +## Usage + +必须设置 `UPSTREAM_URL`、`CLICKHOUSE_URL`，Compose 部署还必须显式设置强 `CLICKHOUSE_PASSWORD`。关键新增配置包括 `LOG_QUEUE_BYTES`、`UPSTREAM_RESPONSE_TIMEOUT`、`UPSTREAM_STREAM_IDLE_TIMEOUT` 和 `TRUSTED_PROXIES`；所有预算和 timeout 必须大于零。`CLICKHOUSE_URL` 支持严格的 `host:port`、`clickhouse://` 和 `clickhouses://`，TLS 与压缩参数受白名单校验。 + +## Verification + +迭代 1 验证全部通过：`gofmt -l .` 无输出，`go test -json ./...` 共 118 个测试通过、0 失败、5 个测试包通过，`go vet ./...` 无诊断，`go build ./` 成功。针对 config、queue、ClickHouse、HTTP/SSE/WebSocket 和 Compose 的失败路径均有测试覆盖。 + +## Journey Log + +> Brief notes on what informed the final design. Not required reading. + +- [lesson] 迭代 1：网络写入的 `Send` 错误无法证明提交与否；最小安全策略是不自动重试、显式统计 ambiguous，并将连接标记为不健康。 +- [lesson] 迭代 1：进程内队列只能提供有界 best-effort；跨进程幂等需要持久化 outbox 和下游幂等协议，不能由本地重试可靠模拟。 +- [pivot] 迭代 1：SSE 完成判定限定为真正的 `text/event-stream` 且等待代理正常返回，避免终止标记导致提前记录不完整响应。 + +## Source Materials + +| File | Role | Notes | +|------|------|-------| +| `docs/compose/specs/2026-07-09-clickhouse-migration.md` | 安全修复规格 | 定义本轮行为边界与取舍 | +| `docs/compose/plans/2026-07-09-clickhouse-migration.md` | 实施计划 | 覆盖跨模块修复和验证 | diff --git a/docs/compose/specs/2026-07-09-clickhouse-migration.md b/docs/compose/specs/2026-07-09-clickhouse-migration.md new file mode 100644 index 0000000..c1a7dda --- /dev/null +++ b/docs/compose/specs/2026-07-09-clickhouse-migration.md @@ -0,0 +1,54 @@ +# ClickHouse Migration 安全修复规格 + +> [!NOTE] +> This document may not reflect the current implementation. +> See the final report for up-to-date state: +> [Final Report](../reports/reliability-security-fixes.md) + +## 修订范围 + +本规格是现有 ClickHouse Migration 的增量修订。保留当前单表 `proxy_logs`、内存异步队列、`httputil.ReverseProxy` 和 best-effort 审计模型，不引入 PostgreSQL 兼容层、持久化 outbox 或 WebSocket 帧采集。 + +## 行为要求 + +### 队列与关闭 + +- `Submit` 始终非阻塞；与 `Stop` 并发、停止后提交和重复停止均不得 panic。 +- 队列由条数和字节双重预算约束。默认 `LOG_QUEUE_SIZE=256`、`LOG_QUEUE_BYTES=67108864`、`MAX_BODY_BYTES=1048576`；任一预算不足即丢弃并计数。 +- `Stop(ctx)` 使用调用方提供的单一总 deadline 排空；deadline 到期取消所有 worker I/O、丢弃剩余条目并返回错误。 +- worker flush 后清空 batch 指针并释放条目的字节预留。 + +### ClickHouse 写入 + +- `PrepareBatch` 错误是确定未提交错误，可有限重试整个 batch。 +- `Append` 错误是确定未提交错误，必须中止/关闭 batch，并允许逐条隔离坏记录；成功记录继续写入，失败记录明确计数。 +- `Send` 错误视为提交结果不明，不自动重试或逐条重发，避免静默重复；该批计为 ambiguous drop 并标记连接不健康。 +- 进程内 best-effort 方案不承诺跨进程 exactly-once。需要更强保证时必须另行引入持久化 outbox 与幂等协议。 + +### 代理完整性 + +- 请求体预读失败时返回固定 400，不调用 upstream，不转发已损坏请求。 +- 502 客户端响应只包含固定 `bad gateway` 和 request ID；内部网络错误仅写服务端日志和 `LogEntry.Error`。 +- 只有 ReverseProxy 正常完成的普通/SSE 响应才提交完整日志。下游写失败、短写、客户端断开或响应复制中断不得提交成完整成功日志。 +- SSE 终止检测只在响应 `Content-Type` 为 `text/event-stream` 时按完整 event 边界解析。终止事件仅用于状态判断，不触发提前提交；提交仍等待代理正常返回。 +- 普通 HTTP 响应体使用可配置总 timeout；SSE 使用可配置 idle timeout，成功读取数据后重置；timeout 必须能关闭阻塞中的上游 Body。 +- WebSocket 成功升级时记录一条 101 握手元数据日志，body 为空。Handler 跟踪已 hijack 连接，并提供受 context 限制的 shutdown 关闭能力；不采集帧。 +- 默认不信任 `X-Forwarded-For`/`X-Real-IP`。仅当直接对端命中 `TRUSTED_PROXIES` CIDR 时，从 XFF 右向左剥离可信代理并选择最近的不可信地址。 + +### 配置与数据库边界 + +- 所有整数、布尔值和 duration 环境变量格式错误时启动失败；要求正值的 body、队列、batch、worker、重连和 timeout 配置必须严格大于零。 +- `CLICKHOUSE_URL` 支持严格 `host:port`、`clickhouse://` 和 `clickhouses://`。仅允许 `secure`、`skip_verify`、`compress` 查询参数；未知参数、冲突 TLS 配置、fragment、空 host/port 均失败。 +- 建表后校验 `system.columns` 与 `system.tables` 的必需列类型、引擎、分区键和排序键；不兼容 schema 拒绝标记健康，不自动重建或改类型。 + +### 启动与部署 + +- server 启动错误与 signal 进入同一清理路径；禁止 goroutine 内 `log.Fatalf` 绕过 defer。 +- shutdown 总预算依次覆盖 HTTP、升级连接、日志队列、DB 和根 context。 +- Compose 固定 ClickHouse 明确版本，不使用 `latest`；默认不发布 ClickHouse 端口；`CLICKHOUSE_PASSWORD` 必须显式设置，应用 DSN 使用独立 `CLICKHOUSE_URL`，避免弱默认和 URL 编码冲突。 + +## 验证 + +- 针对上述失败路径增加 config、logger、proxy、db 单元测试。 +- 运行 `gofmt`、`go test ./...` 和 `go vet ./...`。 +- 使用 Git 提交保留完整变更证据；审查应以提交及其 diff 验证实现和验收项。 diff --git a/docs/compose/specs/reliability-security-fixes.md b/docs/compose/specs/reliability-security-fixes.md new file mode 100644 index 0000000..fed3e21 --- /dev/null +++ b/docs/compose/specs/reliability-security-fixes.md @@ -0,0 +1,5 @@ +# Reliability And Security Fixes Specification + +The canonical dated specification is [2026-07-09-clickhouse-migration.md](2026-07-09-clickhouse-migration.md). + +This stable path is the review entry point for the reliability and security fixes. diff --git a/filter.yaml b/filter.yaml new file mode 100644 index 0000000..74c76fd --- /dev/null +++ b/filter.yaml @@ -0,0 +1,25 @@ +# TokenThief 路径过滤配置（聊天与补全接口） +# +# 仅记录聊天（Chat）和补全（Completions）相关请求。 +# +# mode: +# whitelist - 仅记录命中 patterns 的请求 +# blacklist - 命中 patterns 的请求不记录 +# disabled - 全部记录（忽略 patterns） +# +# Pattern 语法（glob）： +# * 匹配单个段内除 / 之外的任意字符 +# ** 跨段匹配任意字符（包括 /） +# ? 匹配单个非 / 字符 +mode: whitelist +patterns: + # ===== 聊天（Chat） ===== + - /v1/chat/completions # 原生 OpenAI ChatCompletions（也是 Gemini 图片生成的 OpenAI 入口） + - /v1/responses # 原生 OpenAI Responses 格式 + - /v1/messages # 原生 Claude (Anthropic) Messages + - /v1beta/models/*:generateContent # 原生 Gemini generateContent（含媒体识别 / TTS / 图片生成） + - /v1beta/models/*:generateContent/ # 同上，Gemini 图片生成页面所示带尾斜杠形式 + - /v1beta/models/*:streamGenerateContent # 原生 Gemini 流式 generateContent + + # ===== 补全（Completions） ===== + - /v1/completions # 原生 OpenAI 文本补全 diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..c1a6841 --- /dev/null +++ b/go.mod @@ -0,0 +1,29 @@ +module git.misaka.ren/M1saka/token_thief + +go 1.25.0 + +require ( + github.com/ClickHouse/clickhouse-go/v2 v2.47.0 + gopkg.in/yaml.v3 v3.0.1 +) + +require ( + github.com/ClickHouse/ch-go v0.73.0 // indirect + github.com/andybalholm/brotli v1.2.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/go-faster/city v1.0.1 // indirect + github.com/go-faster/errors v0.7.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/klauspost/compress v1.18.6 // indirect + github.com/kr/pretty v0.3.0 // indirect + github.com/paulmach/orb v0.13.0 // indirect + github.com/pierrec/lz4/v4 v4.1.27 // indirect + github.com/rogpeppe/go-internal v1.14.1 // indirect + github.com/segmentio/asm v1.2.1 // indirect + github.com/shopspring/decimal v1.4.0 // indirect + go.opentelemetry.io/otel v1.44.0 // indirect + go.opentelemetry.io/otel/trace v1.44.0 // indirect + go.yaml.in/yaml/v3 v3.0.4 // indirect + golang.org/x/sys v0.46.0 // indirect + gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..ef07565 --- /dev/null +++ b/go.sum @@ -0,0 +1,61 @@ +github.com/ClickHouse/ch-go v0.73.0 h1:jsHiGRbQ3sz+gekvDFJF29LWDo5dzbJm5s1h8TWVP2M= +github.com/ClickHouse/ch-go v0.73.0/go.mod h1:wkFIxrqlXeRJ9cn3r5Fz5Qen9jl5aTMPuGZeuJpANNY= +github.com/ClickHouse/clickhouse-go/v2 v2.47.0 h1:ZDAzrnKSOPTIsm4tdUNfrii2yc8dk4SVRLC77BR7Z5Q= +github.com/ClickHouse/clickhouse-go/v2 v2.47.0/go.mod h1:sPj7C7UYQ2MWHcfX+4eGN6nwnCqwUKfgO6PcwKpd6K8= +github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro= +github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/go-faster/city v1.0.1 h1:4WAxSZ3V2Ws4QRDrscLEDcibJY8uf41H6AhXDrNDcGw= +github.com/go-faster/city v1.0.1/go.mod h1:jKcUJId49qdW3L1qKHH/3wPeUstCVpVSXTM6vO3VcTw= +github.com/go-faster/errors v0.7.1 h1:MkJTnDoEdi9pDabt1dpWf7AA8/BaSYZqibYyhZ20AYg= +github.com/go-faster/errors v0.7.1/go.mod h1:5ySTjWFiphBs07IKuiL69nxdfd5+fzh1u7FPGZP2quo= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/klauspost/compress v1.18.6 h1:2jupLlAwFm95+YDR+NwD2MEfFO9d4z4Prjl1XXDjuao= +github.com/klauspost/compress v1.18.6/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= +github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/paulmach/orb v0.13.0 h1:r7n7mQGGF+cj/CbcivEj9J3HGK+XR+yXnvzRdq9saIw= +github.com/paulmach/orb v0.13.0/go.mod h1:6scRWINywA2Jf05dcjOfLfxrUIMECvTSG2MVbRLxu/k= +github.com/pierrec/lz4/v4 v4.1.27 h1:+PhzhWDrjRj89TH2sw43nE3+4+W8lSxIuQadEHZyjUk= +github.com/pierrec/lz4/v4 v4.1.27/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/segmentio/asm v1.2.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0= +github.com/segmentio/asm v1.2.1/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= +github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k= +github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= +go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= +go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= +go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= +go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/logger/model.go b/logger/model.go new file mode 100644 index 0000000..32d744a --- /dev/null +++ b/logger/model.go @@ -0,0 +1,24 @@ +package logger + +import "time" + +// LogEntry 表示一次代理请求的完整记录。 +type LogEntry struct { + RequestID string + Method string + Path string + Query string + ClientIP string + RequestHeaders []byte // JSON + RequestBody []byte + RequestTruncated bool + StatusCode int + ResponseHeaders []byte // JSON + ResponseBody []byte + ResponseTruncated bool + IsStream bool + LatencyMS int64 + StartedAt time.Time + FinishedAt time.Time + Error string +} diff --git a/logger/queue.go b/logger/queue.go new file mode 100644 index 0000000..0aeb581 --- /dev/null +++ b/logger/queue.go @@ -0,0 +1,447 @@ +package logger + +import ( + "context" + "errors" + "log" + "math/rand" + "sync" + "sync/atomic" + "time" + + "github.com/ClickHouse/clickhouse-go/v2/lib/driver" + + "git.misaka.ren/M1saka/token_thief/db" +) + +const ( + maxAttempts = 3 + BaseBackoff = 200 * time.Millisecond +) + +type Batch interface { + Append(v ...any) error + Send() error + Abort() error +} + +type BatchPreparer interface { + PrepareBatch(ctx context.Context, query string) (Batch, error) +} + +type clickHouseBatchPreparer struct { + conn driver.Conn +} + +// The explicit adapter keeps the local Batch contract aligned with the real driver. +func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string) (Batch, error) { + return p.conn.PrepareBatch(ctx, query) +} + +type Stats struct { + Enqueued uint64 + Dropped uint64 + Failed uint64 + Ambiguous uint64 + Bytes int64 +} + +// FlushResult separates retry-safe Prepare failures from final or ambiguous failures. +type FlushResult struct { + Retry []*LogEntry + Failed int + Ambiguous int + Err error +} + +// Queue is an asynchronous, bounded logger. Submit never waits for database work. +type Queue struct { + ch chan *LogEntry + pool *db.Pool + batchSize int + batchInterval time.Duration + workers int + byteBudget int64 + + dropped atomic.Uint64 + failed atomic.Uint64 + ambiguous atomic.Uint64 + enq atomic.Uint64 + bytes atomic.Int64 + + mu sync.RWMutex + started bool + stopped bool + shutdownCtx context.Context + reporterCancel context.CancelFunc + workCancel context.CancelFunc + wg sync.WaitGroup + done chan struct{} + doneOnce sync.Once +} + +// NewQueue accepts an optional byte budget. A non-positive or omitted budget disables byte limiting. +func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval time.Duration, byteBudget ...int64) *Queue { + var budget int64 + if len(byteBudget) > 0 { + budget = byteBudget[0] + } + if queueSize < 0 { + queueSize = 0 + } + if batchSize < 1 { + batchSize = 1 + } + if workers < 1 { + workers = 1 + } + if batchInterval <= 0 { + batchInterval = time.Second + } + return &Queue{ + ch: make(chan *LogEntry, queueSize), + pool: pool, + batchSize: batchSize, + batchInterval: batchInterval, + workers: workers, + byteBudget: budget, + done: make(chan struct{}), + } +} + +// EstimatedBytes covers all variable-size strings and byte slices retained by an entry. +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)) +} + +func (q *Queue) Stats() Stats { + return Stats{ + Enqueued: q.enq.Load(), + Dropped: q.dropped.Load(), + Failed: q.failed.Load(), + Ambiguous: q.ambiguous.Load(), + Bytes: q.bytes.Load(), + } +} + +// Start starts workers once. The root context controls reporting only; Stop owns worker shutdown. +func (q *Queue) Start(root context.Context) { + q.mu.Lock() + if q.started || q.stopped { + q.mu.Unlock() + return + } + q.started = true + workCtx, workCancel := context.WithCancel(context.Background()) + reportCtx, reporterCancel := context.WithCancel(root) + q.workCancel = workCancel + q.reporterCancel = reporterCancel + for i := 0; i < q.workers; i++ { + q.wg.Add(1) + go q.run(workCtx) + } + q.wg.Add(1) + go q.reportLoop(reportCtx) + q.mu.Unlock() +} + +// Stop closes submissions once and waits under the caller's single total deadline. +// The variadic form permits legacy Stop() calls while new callers should pass a context. +func (q *Queue) Stop(contexts ...context.Context) error { + ctx := context.Background() + if len(contexts) > 0 && contexts[0] != nil { + ctx = contexts[0] + } + + q.mu.Lock() + if !q.stopped { + q.stopped = true + q.shutdownCtx = ctx + close(q.ch) + if q.reporterCancel != nil { + q.reporterCancel() + } + if q.started { + go func() { + q.wg.Wait() + q.doneOnce.Do(func() { close(q.done) }) + }() + } else { + q.discardQueued() + q.doneOnce.Do(func() { close(q.done) }) + } + } + done := q.done + workCancel := q.workCancel + q.mu.Unlock() + + select { + case <-done: + return nil + case <-ctx.Done(): + if workCancel != nil { + workCancel() + } + return ctx.Err() + } +} + +// Submit reserves memory and enqueues without waiting; full, over-budget, and stopped queues drop. +func (q *Queue) Submit(e *LogEntry) { + if e == nil { + q.dropped.Add(1) + return + } + size := EstimatedBytes(e) + if !q.mu.TryRLock() { + q.dropped.Add(1) + return + } + if q.stopped || !q.reserve(size) { + q.mu.RUnlock() + q.dropped.Add(1) + return + } + select { + case q.ch <- e: + q.enq.Add(1) + default: + q.release(size) + q.dropped.Add(1) + } + q.mu.RUnlock() +} + +func (q *Queue) reserve(size int64) bool { + if q.byteBudget <= 0 { + q.bytes.Add(size) + return true + } + for { + used := q.bytes.Load() + if size > q.byteBudget-used { + return false + } + if q.bytes.CompareAndSwap(used, used+size) { + return true + } + } +} + +func (q *Queue) release(size int64) { + q.bytes.Add(-size) +} + +func (q *Queue) discardQueued() { + for e := range q.ch { + q.release(EstimatedBytes(e)) + q.failed.Add(1) + } +} + +func (q *Queue) run(workCtx context.Context) { + defer q.wg.Done() + batch := make([]*LogEntry, 0, q.batchSize) + ticker := time.NewTicker(q.batchInterval) + defer ticker.Stop() + + flush := func() { + if len(batch) == 0 { + return + } + q.flush(q.flushContext(workCtx), batch) + for _, e := range batch { + q.release(EstimatedBytes(e)) + } + clear(batch) + batch = batch[:0] + } + + for { + select { + case e, ok := <-q.ch: + if !ok { + flush() + return + } + batch = append(batch, e) + if len(batch) >= q.batchSize { + flush() + } + case <-ticker.C: + flush() + case <-workCtx.Done(): + flush() + q.discardQueued() + return + } + } +} + +func (q *Queue) flushContext(workCtx context.Context) context.Context { + q.mu.RLock() + defer q.mu.RUnlock() + if q.shutdownCtx != nil { + return q.shutdownCtx + } + return workCtx +} + +func (q *Queue) flush(ctx context.Context, entries []*LogEntry) { + if q.pool == nil || !q.pool.Healthy() { + q.failed.Add(uint64(len(entries))) + return + } + + retry := entries + var lastErr error + for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ { + if err := ctx.Err(); err != nil { + lastErr = err + break + } + conn := q.pool.Get() + if conn == nil { + lastErr = errors.New("pool nil") + break + } + result := Flush(ctx, clickHouseBatchPreparer{conn: conn}, retry) + q.failed.Add(uint64(result.Failed)) + q.ambiguous.Add(uint64(result.Ambiguous)) + if result.Ambiguous > 0 { + q.pool.MarkUnhealthy() + } + lastErr = result.Err + retry = result.Retry + if len(retry) == 0 { + return + } + if attempt < maxAttempts && !waitBackoff(ctx, attempt) { + lastErr = ctx.Err() + break + } + } + + if len(retry) > 0 { + q.failed.Add(uint64(len(retry))) + q.pool.MarkUnhealthy() + log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr) + } +} + +func waitBackoff(ctx context.Context, attempt int) bool { + wait := BaseBackoff + for i := 1; i < attempt; i++ { + wait *= 3 + } + jitter := time.Duration((rand.Float64()*0.4 - 0.2) * float64(wait)) + timer := time.NewTimer(wait + jitter) + defer timer.Stop() + select { + case <-timer.C: + return true + case <-ctx.Done(): + return false + } +} + +const insertStatement = `INSERT INTO proxy_logs ( + request_id, method, path, query, client_ip, + request_headers, request_body, request_truncated, + status_code, response_headers, response_body, response_truncated, + is_stream, latency_ms, started_at, finished_at, error + )` + +// Flush executes one batch. Only Prepare failures are retryable. Append failures are isolated by row; +// Send failures are ambiguous and therefore never replayed. +func Flush(ctx context.Context, conn BatchPreparer, entries []*LogEntry) FlushResult { + if len(entries) == 0 { + return FlushResult{} + } + result := flushOnce(ctx, conn, entries) + if result.stage != flushAppend || len(entries) == 1 { + return result.public(entries) + } + + combined := FlushResult{} + for _, entry := range entries { + single := flushOnce(ctx, conn, []*LogEntry{entry}).public([]*LogEntry{entry}) + combined.Retry = append(combined.Retry, single.Retry...) + combined.Failed += single.Failed + combined.Ambiguous += single.Ambiguous + combined.Err = errors.Join(combined.Err, single.Err) + } + return combined +} + +type flushStage uint8 + +const ( + flushSuccess flushStage = iota + flushPrepare + flushAppend + flushSend +) + +type flushAttempt struct { + stage flushStage + err error +} + +func (r flushAttempt) public(entries []*LogEntry) FlushResult { + switch r.stage { + case flushSuccess: + return FlushResult{} + case flushPrepare: + return FlushResult{Retry: entries, Err: r.err} + case flushSend: + return FlushResult{Failed: len(entries), Ambiguous: len(entries), Err: r.err} + default: + return FlushResult{Failed: len(entries), Err: r.err} + } +} + +func flushOnce(ctx context.Context, conn BatchPreparer, entries []*LogEntry) flushAttempt { + batch, err := conn.PrepareBatch(ctx, insertStatement) + if err != nil { + return flushAttempt{stage: flushPrepare, err: err} + } + abort := func(cause error) error { + return errors.Join(cause, batch.Abort()) + } + for _, e := range entries { + if err := batch.Append( + e.RequestID, e.Method, e.Path, e.Query, e.ClientIP, + string(e.RequestHeaders), string(e.RequestBody), e.RequestTruncated, + int32(e.StatusCode), string(e.ResponseHeaders), string(e.ResponseBody), e.ResponseTruncated, + e.IsStream, e.LatencyMS, e.StartedAt, e.FinishedAt, e.Error, + ); err != nil { + return flushAttempt{stage: flushAppend, err: abort(err)} + } + } + if err := batch.Send(); err != nil { + return flushAttempt{stage: flushSend, err: abort(err)} + } + return flushAttempt{stage: flushSuccess} +} + +func (q *Queue) reportLoop(ctx context.Context) { + defer q.wg.Done() + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + stats := q.Stats() + healthy := q.pool != nil && q.pool.Healthy() + log.Printf("[logger] metrics: enq=%d dropped_queue_full=%d dropped_db_fail=%d ambiguous_send=%d queue_len=%d queue_bytes=%d db_healthy=%v", + stats.Enqueued, stats.Dropped, stats.Failed, stats.Ambiguous, len(q.ch), stats.Bytes, healthy) + } + } +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..ebe8e95 --- /dev/null +++ b/main.go @@ -0,0 +1,109 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "git.misaka.ren/M1saka/token_thief/config" + "git.misaka.ren/M1saka/token_thief/db" + "git.misaka.ren/M1saka/token_thief/logger" + "git.misaka.ren/M1saka/token_thief/proxy" +) + +func main() { + log.SetFlags(log.LstdFlags | log.Lmicroseconds) + if err := run(); err != nil { + log.Printf("[main] fatal: %v", err) + os.Exit(1) + } +} + +func run() error { + cfg, err := config.Load() + if err != nil { + return fmt.Errorf("config: %w", err) + } + if _, err := db.ClickHouseOptions(cfg.ClickHouseURL); err != nil { + return err + } + + filter, err := config.LoadFilter(cfg.FilterFile) + if err != nil { + return fmt.Errorf("load filter: %w", err) + } + log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns)) + + rootCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pool := db.NewPool(rootCtx, cfg.ClickHouseURL, cfg.DBReconnectInterval) + + queue := logger.NewQueue(pool, cfg.LogQueueSize, cfg.LogBatchSize, cfg.LogWorkers, cfg.LogBatchInterval, cfg.LogQueueBytes) + queue.Start(rootCtx) + + h := proxy.NewWithOptions(cfg.UpstreamURL, filter, queue, cfg.MaxBodyBytes, proxy.Options{ + UpstreamTimeout: cfg.UpstreamTimeout, + ResponseTimeout: cfg.UpstreamResponseTimeout, + SSEIdleTimeout: cfg.UpstreamStreamIdleTimeout, + UpstreamTLSInsecureSkipVerify: cfg.UpstreamTLSInsecureSkipVerify, + TrustedProxies: cfg.TrustedProxies, + }) + + srv := &http.Server{ + Addr: cfg.ListenAddr, + Handler: h, + ReadHeaderTimeout: 30 * time.Second, + ReadTimeout: cfg.ReadTimeout, + WriteTimeout: cfg.WriteTimeout, + IdleTimeout: cfg.IdleTimeout, + } + + serverErr := make(chan error, 1) + go func() { + log.Printf("[main] listening on %s, upstream=%s", cfg.ListenAddr, cfg.UpstreamURL) + serverErr <- srv.ListenAndServe() + }() + + stop := make(chan os.Signal, 1) + signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM) + var runErr error + select { + case <-stop: + log.Printf("[main] shutdown signal received") + case err := <-serverErr: + if !errors.Is(err, http.ErrServerClosed) { + runErr = fmt.Errorf("server: %w", err) + } + } + signal.Stop(stop) + + // HTTP、升级连接和日志排空共享同一个关闭总预算。 + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) + defer shutdownCancel() + if err := srv.Shutdown(shutdownCtx); err != nil { + log.Printf("[main] http shutdown: %v", err) + } + if err := h.Shutdown(shutdownCtx); err != nil { + log.Printf("[main] upgraded connection shutdown: %v", err) + } + queueStopped := true + if err := queue.Stop(shutdownCtx); err != nil { + log.Printf("[main] queue shutdown: %v", err) + queueStopped = false + } + cancel() + if queueStopped { + pool.Close() + } else { + log.Printf("[main] skip db close while queue workers are still exiting") + } + log.Printf("[main] bye") + return runErr +} diff --git a/proxy/capture.go b/proxy/capture.go new file mode 100644 index 0000000..5d5cf41 --- /dev/null +++ b/proxy/capture.go @@ -0,0 +1,104 @@ +package proxy + +import ( + "bytes" + "crypto/rand" + "encoding/hex" + "encoding/json" + "io" + "net/http" + "net/netip" + "strings" +) + +// readRequestBody 在转发前完整读取请求体，确保读取失败时不会向上游发送损坏请求。 +func readRequestBody(r *http.Request, max int64) (captured []byte, truncated bool, err error) { + if r.Body == nil || r.ContentLength == 0 { + return nil, false, nil + } + body, err := io.ReadAll(r.Body) + if err != nil { + return nil, false, err + } + if err := r.Body.Close(); err != nil { + return nil, false, err + } + r.Body = io.NopCloser(bytes.NewReader(body)) + if int64(len(body)) > max { + return body[:max], true, nil + } + return body, false, nil +} + +func headersJSON(h http.Header) []byte { + if len(h) == 0 { + return nil + } + b, err := json.Marshal(h) + if err != nil { + return nil + } + return b +} + +func newRequestID() string { + var b [16]byte + if _, err := rand.Read(b[:]); err != nil { + return "unknown" + } + return hex.EncodeToString(b[:]) +} + +// clientIP 从请求中提取客户端 IP。 +func clientIP(r *http.Request, trusted []netip.Prefix) string { + peer, ok := parsePeerAddr(r.RemoteAddr) + if !ok { + return r.RemoteAddr + } + if !isTrusted(peer, trusted) { + return peer.String() + } + xff := strings.Split(r.Header.Get("X-Forwarded-For"), ",") + if len(xff) == 1 && strings.TrimSpace(xff[0]) == "" { + if realIP, err := netip.ParseAddr(strings.TrimSpace(r.Header.Get("X-Real-IP"))); err == nil { + return realIP.Unmap().String() + } + return peer.String() + } + chain := make([]netip.Addr, len(xff)) + for i, raw := range xff { + addr, err := netip.ParseAddr(strings.TrimSpace(raw)) + if err != nil { + return peer.String() + } + chain[i] = addr.Unmap() + } + client := peer + for i := len(chain) - 1; i >= 0 && isTrusted(client, trusted); i-- { + client = chain[i] + } + return client.String() +} + +func parsePeerAddr(remote string) (netip.Addr, bool) { + if addrPort, err := netip.ParseAddrPort(remote); err == nil { + return addrPort.Addr().Unmap(), true + } + addr, err := netip.ParseAddr(remote) + return addr.Unmap(), err == nil +} + +func isTrusted(addr netip.Addr, prefixes []netip.Prefix) bool { + for _, prefix := range prefixes { + if prefix.Contains(addr) { + return true + } + } + return false +} + +// isStreamResponse 通过响应头判断是否为流式响应。 +func isStreamResponse(h http.Header) bool { + ct := strings.ToLower(strings.TrimSpace(strings.SplitN(h.Get("Content-Type"), ";", 2)[0])) + return ct == "text/event-stream" +} diff --git a/proxy/proxy.go b/proxy/proxy.go new file mode 100644 index 0000000..deea9d3 --- /dev/null +++ b/proxy/proxy.go @@ -0,0 +1,339 @@ +package proxy + +import ( + "context" + "crypto/tls" + "log" + "net" + "net/http" + "net/http/httputil" + "net/netip" + "net/url" + "strings" + "sync" + "sync/atomic" + "time" + + "git.misaka.ren/M1saka/token_thief/config" + "git.misaka.ren/M1saka/token_thief/logger" +) + +// LogSubmitter 是 proxy 唯一依赖的日志接收方接口。 +type LogSubmitter interface { + Submit(*logger.LogEntry) +} + +// 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{} +} + +// Options 控制反代连接上游时的网络行为。 +type Options struct { + UpstreamTimeout time.Duration + ResponseTimeout time.Duration + SSEIdleTimeout time.Duration + UpstreamTLSInsecureSkipVerify bool + TrustedProxies []netip.Prefix +} + +// requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。 +type requestState struct { + requestID string + lastErr atomic.Pointer[string] + upgrade atomic.Pointer[upgradeResponse] + readFailed atomic.Bool +} + +type upgradeResponse struct { + status int + header http.Header +} + +type ctxKey struct{} + +func newRequestState(id string) *requestState { return &requestState{requestID: id} } + +func stateFromCtx(ctx context.Context) *requestState { + v, _ := ctx.Value(ctxKey{}).(*requestState) + return v +} + +func New(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody int64, upstreamTimeout ...time.Duration) *Handler { + opts := Options{} + if len(upstreamTimeout) > 0 { + opts.UpstreamTimeout = upstreamTimeout[0] + } + return NewWithOptions(upstream, filter, queue, maxBody, opts) +} + +func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody int64, opts Options) *Handler { + 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 + transport.ResponseHeaderTimeout = opts.UpstreamTimeout + transport.TLSHandshakeTimeout = opts.UpstreamTimeout + } + if opts.UpstreamTLSInsecureSkipVerify { + transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} + } + rp.Transport = transport + } + + origDirector := rp.Director + rp.Director = func(r *http.Request) { + origDirector(r) + r.Host = upstream.Host + } + + // ModifyResponse 在响应头写回客户端之前调用，确保 X-Request-Id 一定生效。 + rp.ModifyResponse = func(resp *http.Response) error { + if st := stateFromCtx(resp.Request.Context()); st != nil { + resp.Header.Set("X-Request-Id", st.requestID) + if resp.StatusCode == http.StatusSwitchingProtocols && + strings.EqualFold(resp.Request.Header.Get("Upgrade"), "websocket") && + strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") { + st.upgrade.Store(&upgradeResponse{status: resp.StatusCode, header: resp.Header.Clone()}) + } + } + if timeout := responseBodyTimeout(resp, opts); timeout > 0 { + resp.Body = newTimeoutBody(resp.Body, timeout, isStreamResponse(resp.Header)) + } + if st := stateFromCtx(resp.Request.Context()); st != nil && resp.StatusCode != http.StatusSwitchingProtocols { + resp.Body = &trackingBody{ReadCloser: resp.Body, failed: &st.readFailed} + } + return nil + } + + rp.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) { + log.Printf("[proxy] upstream error for %s %s: %v", r.Method, r.URL.Path, err) + // 把错误暴露给主 handler，使其能写入日志。 + if st := stateFromCtx(r.Context()); st != nil { + s := err.Error() + st.lastErr.Store(&s) + // ErrorHandler 路径下 ModifyResponse 不会被调用，这里手动写 X-Request-Id。 + w.Header().Set("X-Request-Id", st.requestID) + } + http.Error(w, "bad gateway", http.StatusBadGateway) + } + + 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{}), + } +} + +func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + // 健康检查不参与反代与日志。 + if r.URL.Path == "/healthz" { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + return + } + + shouldLog := h.filter.ShouldLog(r.URL.Path) + reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes) + if err != nil && !shouldLog { + log.Printf("[proxy] read request body failed: %v", err) + _ = r.Body.Close() + http.Error(w, "bad request", http.StatusBadRequest) + return + } + if !shouldLog { + log.Printf("[proxy] skip log by filter method=%s path=%s", r.Method, r.URL.Path) + cw := newCaptureWriter(w, 0) + cw.OnHijack(h.trackConn) + h.rp.ServeHTTP(cw, r) + return + } + + started := time.Now() + reqID := newRequestID() + log.Printf("[proxy] capture start request_id=%s method=%s path=%s", reqID, r.Method, r.URL.Path) + st := newRequestState(reqID) + r = r.WithContext(context.WithValue(r.Context(), ctxKey{}, st)) + cw := newCaptureWriter(w, h.maxBodyBytes) + + if err != nil { + log.Printf("[proxy] read request body failed: %v", err) + _ = r.Body.Close() + s := "read request body: " + err.Error() + st.lastErr.Store(&s) + h.serveRequestBodyError(cw, r, st, started, reqID) + return + } + + reqHeadersJSON := headersJSON(r.Header) + clientAddr := clientIP(r, h.trusted) + method := r.Method + path := r.URL.Path + query := r.URL.RawQuery + + var submitOnce sync.Once + submit := func(finished time.Time) { + entry, ok := h.buildLogEntry(cw, st, logEntryInput{ + requestID: reqID, + method: method, + path: path, + query: query, + clientAddr: clientAddr, + requestHeaders: reqHeadersJSON, + requestBody: reqBody, + requestTruncated: reqTruncated, + started: started, + finished: finished, + }) + if !ok { + return + } + h.queue.Submit(entry) + log.Printf("[proxy] capture finish request_id=%s method=%s path=%s status=%d is_stream=%v latency_ms=%d", + reqID, method, path, entry.StatusCode, entry.IsStream, entry.LatencyMS) + } + cw.OnHijack(func(conn net.Conn) net.Conn { + tracked := h.trackConn(conn) + if upgrade := st.upgrade.Load(); upgrade != nil { + cw.SetHijackedResponse(upgrade.status, upgrade.header) + submitOnce.Do(func() { submit(time.Now()) }) + } + return tracked + }) + h.rp.ServeHTTP(cw, r) + + finished := time.Now() + responseComplete := !isStreamResponse(cw.Header()) || cw.SSEComplete() + if cw.Complete() && responseComplete && !st.readFailed.Load() && r.Context().Err() == nil { + submitOnce.Do(func() { submit(finished) }) + } +} + +func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string) { + cw.Header().Set("X-Request-Id", requestID) + http.Error(cw, "bad request", http.StatusBadRequest) + if !cw.Complete() { + return + } + entry, ok := h.buildLogEntry(cw, st, logEntryInput{ + requestID: requestID, + method: r.Method, + path: r.URL.Path, + query: r.URL.RawQuery, + clientAddr: clientIP(r, h.trusted), + requestHeaders: headersJSON(r.Header), + started: started, + finished: time.Now(), + }) + if ok { + h.queue.Submit(entry) + } +} + +// Shutdown closes all active hijacked connections and waits for their release. +func (h *Handler) Shutdown(ctx context.Context) error { + for { + h.connMu.Lock() + if len(h.conns) == 0 { + h.connMu.Unlock() + return nil + } + conns := make([]net.Conn, 0, len(h.conns)) + for conn := range h.conns { + conns = append(conns, conn) + } + changed := h.connChanged + h.connMu.Unlock() + for _, conn := range conns { + _ = conn.Close() + } + select { + case <-changed: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +func (h *Handler) trackConn(conn net.Conn) net.Conn { + tracked := &trackedConn{Conn: conn} + tracked.onClose = func() { + h.connMu.Lock() + delete(h.conns, tracked) + close(h.connChanged) + h.connChanged = make(chan struct{}) + h.connMu.Unlock() + } + h.connMu.Lock() + h.conns[tracked] = struct{}{} + close(h.connChanged) + h.connChanged = make(chan struct{}) + h.connMu.Unlock() + return tracked +} + +type logEntryInput struct { + requestID string + method string + path string + query string + clientAddr string + requestHeaders []byte + requestBody []byte + requestTruncated bool + started time.Time + finished time.Time +} + +func (h *Handler) buildLogEntry(cw *captureWriter, st *requestState, in logEntryInput) (*logger.LogEntry, bool) { + if cw.Hijacked() && cw.Status() != http.StatusSwitchingProtocols { + return nil, false + } + + var errMsg string + if p := st.lastErr.Load(); p != nil { + errMsg = *p + } + + isStream := isStreamResponse(cw.Header()) + responseBody := append([]byte(nil), cw.Body()...) + responseTruncated := cw.Truncated() + if isStream && !responseTruncated { + if assembled, ok := assembleSSEJSON(responseBody); ok { + responseBody = assembled + } + } + + return &logger.LogEntry{ + RequestID: in.requestID, + Method: in.method, + Path: in.path, + Query: in.query, + ClientIP: in.clientAddr, + RequestHeaders: in.requestHeaders, + RequestBody: in.requestBody, + RequestTruncated: in.requestTruncated, + StatusCode: cw.Status(), + ResponseHeaders: headersJSON(cw.Header()), + ResponseBody: responseBody, + ResponseTruncated: responseTruncated, + IsStream: isStream, + LatencyMS: in.finished.Sub(in.started).Milliseconds(), + StartedAt: in.started, + FinishedAt: in.finished, + Error: errMsg, + }, true +} diff --git a/proxy/sse.go b/proxy/sse.go new file mode 100644 index 0000000..5f416ff --- /dev/null +++ b/proxy/sse.go @@ -0,0 +1,661 @@ +package proxy + +import ( + "bufio" + "bytes" + "encoding/json" + "strings" +) + +func assembleSSEJSON(body []byte) ([]byte, bool) { + payloads := sseDataPayloads(body) + if len(payloads) == 0 { + return nil, false + } + if assembled, ok := assembleOpenAICompletionsSSE(payloads); ok { + return assembled, true + } + if assembled, ok := assembleOpenAIChatSSE(payloads); ok { + return assembled, true + } + if assembled, ok := assembleOpenAIResponsesSSE(payloads); ok { + return assembled, true + } + if assembled, ok := assembleAnthropicSSE(payloads); ok { + return assembled, true + } + if assembled, ok := assembleGeminiSSE(payloads); ok { + return assembled, true + } + return nil, false +} + +func sseDataPayloads(body []byte) []string { + payloads := make([]string, 0) + scanner := bufio.NewScanner(bytes.NewReader(body)) + scanner.Buffer(make([]byte, 0, 64*1024), len(body)+1) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if !strings.HasPrefix(line, "data:") { + continue + } + payload := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if payload == "" || payload == "[DONE]" { + continue + } + payloads = append(payloads, payload) + } + return payloads +} + +type openAICompletionChunk struct { + ID string `json:"id,omitempty"` + Object string `json:"object,omitempty"` + Created int64 `json:"created,omitempty"` + Model string `json:"model,omitempty"` + Usage json.RawMessage `json:"usage"` + Choices []struct { + Index int `json:"index"` + Text *string `json:"text"` + FinishReason *string `json:"finish_reason"` + } `json:"choices"` +} + +type openAICompletionChoice struct { + index int + text strings.Builder + finishReason string +} + +func assembleOpenAICompletionsSSE(payloads []string) ([]byte, bool) { + choices := map[int]*openAICompletionChoice{} + order := make([]int, 0, 1) + var id, object, model string + var created int64 + var usage json.RawMessage + matched := false + + for _, payload := range payloads { + var chunk openAICompletionChunk + if err := json.Unmarshal([]byte(payload), &chunk); err != nil { + return nil, false + } + if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 { + return nil, false + } + for _, choice := range chunk.Choices { + if choice.Text == nil { + return nil, false + } + } + matched = true + if id == "" { + id = chunk.ID + } + if object == "" { + object = chunk.Object + } + if created == 0 { + created = chunk.Created + } + if model == "" { + model = chunk.Model + } + if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" { + usage = chunk.Usage + } + + for _, choice := range chunk.Choices { + assembled := choices[choice.Index] + if assembled == nil { + assembled = &openAICompletionChoice{index: choice.Index} + choices[choice.Index] = assembled + order = append(order, choice.Index) + } + assembled.text.WriteString(*choice.Text) + if choice.FinishReason != nil { + assembled.finishReason = *choice.FinishReason + } + } + } + if !matched { + return nil, false + } + + assembled := struct { + ID string `json:"id,omitempty"` + Object string `json:"object,omitempty"` + Created int64 `json:"created,omitempty"` + Model string `json:"model,omitempty"` + Usage json.RawMessage `json:"usage,omitempty"` + Choices []struct { + Index int `json:"index"` + Text string `json:"text"` + FinishReason string `json:"finish_reason,omitempty"` + } `json:"choices"` + }{ID: id, Object: object, Created: created, Model: model, Usage: usage} + for _, index := range order { + choice := choices[index] + assembled.Choices = append(assembled.Choices, struct { + Index int `json:"index"` + Text string `json:"text"` + FinishReason string `json:"finish_reason,omitempty"` + }{Index: choice.index, Text: choice.text.String(), FinishReason: choice.finishReason}) + } + + data, err := json.Marshal(assembled) + return data, err == nil +} + +type openAIChatChunk struct { + ID string `json:"id,omitempty"` + Object string `json:"object,omitempty"` + Created int64 `json:"created,omitempty"` + Model string `json:"model,omitempty"` + SystemFingerprint string `json:"system_fingerprint,omitempty"` + Usage json.RawMessage `json:"usage"` + Choices []struct { + Index int `json:"index"` + Delta struct { + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []struct { + Index int `json:"index"` + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` + } `json:"function,omitempty"` + } `json:"tool_calls,omitempty"` + } `json:"delta"` + FinishReason *string `json:"finish_reason"` + NativeFinishReason *string `json:"native_finish_reason"` + } `json:"choices"` +} + +type openAIChatChoice struct { + index int + role string + content strings.Builder + reasoning strings.Builder + toolCalls map[int]*openAIToolCall + toolOrder []int + finishReason string + nativeFinish string +} + +type openAIToolCall struct { + id string + callType string + name string + arguments strings.Builder +} + +type openAIChatResponse struct { + ID string `json:"id,omitempty"` + Object string `json:"object,omitempty"` + Created int64 `json:"created,omitempty"` + Model string `json:"model,omitempty"` + SystemFingerprint string `json:"system_fingerprint,omitempty"` + Usage json.RawMessage `json:"usage,omitempty"` + Choices []openAIResponseChoice `json:"choices"` +} + +type openAIResponseChoice struct { + Index int `json:"index"` + Message openAIResponseMessage `json:"message"` + FinishReason string `json:"finish_reason,omitempty"` + NativeFinishReason string `json:"native_finish_reason,omitempty"` +} + +type openAIResponseMessage struct { + Role string `json:"role,omitempty"` + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []openAIResponseTool `json:"tool_calls,omitempty"` +} + +type openAIResponseTool struct { + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function openAIResponseToolFunction `json:"function"` +} + +type openAIResponseToolFunction struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments"` +} + +func assembleOpenAIChatSSE(payloads []string) ([]byte, bool) { + choices := map[int]*openAIChatChoice{} + order := make([]int, 0, 1) + var id, object string + var created int64 + var model string + var systemFingerprint string + var usage json.RawMessage + matched := false + + for _, payload := range payloads { + var chunk openAIChatChunk + if err := json.Unmarshal([]byte(payload), &chunk); err != nil { + return nil, false + } + if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 { + return nil, false + } + matched = true + if id == "" { + id = chunk.ID + } + if object == "" { + object = strings.TrimSuffix(chunk.Object, ".chunk") + } + if created == 0 { + created = chunk.Created + } + if model == "" { + model = chunk.Model + } + if systemFingerprint == "" { + systemFingerprint = chunk.SystemFingerprint + } + if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" { + usage = chunk.Usage + } + + for _, choice := range chunk.Choices { + assembled := choices[choice.Index] + if assembled == nil { + assembled = &openAIChatChoice{index: choice.Index} + choices[choice.Index] = assembled + order = append(order, choice.Index) + } + if choice.Delta.Role != "" { + assembled.role = choice.Delta.Role + } + if choice.Delta.Content != "" { + assembled.content.WriteString(choice.Delta.Content) + } + if choice.Delta.ReasoningContent != "" { + assembled.reasoning.WriteString(choice.Delta.ReasoningContent) + } + for _, toolCall := range choice.Delta.ToolCalls { + if assembled.toolCalls == nil { + assembled.toolCalls = map[int]*openAIToolCall{} + } + assembledTool := assembled.toolCalls[toolCall.Index] + if assembledTool == nil { + assembledTool = &openAIToolCall{} + assembled.toolCalls[toolCall.Index] = assembledTool + assembled.toolOrder = append(assembled.toolOrder, toolCall.Index) + } + if toolCall.ID != "" { + assembledTool.id = toolCall.ID + } + if toolCall.Type != "" { + assembledTool.callType = toolCall.Type + } + if toolCall.Function.Name != "" { + assembledTool.name = toolCall.Function.Name + } + if toolCall.Function.Arguments != "" { + assembledTool.arguments.WriteString(toolCall.Function.Arguments) + } + } + if choice.FinishReason != nil { + assembled.finishReason = *choice.FinishReason + } + if choice.NativeFinishReason != nil { + assembled.nativeFinish = *choice.NativeFinishReason + } + } + } + if !matched { + return nil, false + } + + assembled := openAIChatResponse{ID: id, Object: object, Created: created, Model: model, SystemFingerprint: systemFingerprint, Usage: usage} + for _, index := range order { + choice := choices[index] + out := openAIResponseChoice{Index: choice.index, FinishReason: choice.finishReason, NativeFinishReason: choice.nativeFinish} + out.Message.Role = choice.role + out.Message.Content = choice.content.String() + out.Message.ReasoningContent = choice.reasoning.String() + for _, toolIndex := range choice.toolOrder { + toolCall := choice.toolCalls[toolIndex] + outTool := openAIResponseTool{ID: toolCall.id, Type: toolCall.callType} + outTool.Function.Name = toolCall.name + outTool.Function.Arguments = toolCall.arguments.String() + out.Message.ToolCalls = append(out.Message.ToolCalls, outTool) + } + assembled.Choices = append(assembled.Choices, out) + } + + data, err := json.Marshal(assembled) + return data, err == nil +} + +type openAIResponsesEvent struct { + Type string `json:"type"` + Response json.RawMessage `json:"response"` +} + +func assembleOpenAIResponsesSSE(payloads []string) ([]byte, bool) { + var completed json.RawMessage + matched := false + for _, payload := range payloads { + var event openAIResponsesEvent + if err := json.Unmarshal([]byte(payload), &event); err != nil { + return nil, false + } + if !strings.HasPrefix(event.Type, "response.") { + return nil, false + } + matched = true + if event.Type == "response.completed" && len(event.Response) > 0 { + completed = append(json.RawMessage(nil), event.Response...) + } + } + if !matched || len(completed) == 0 { + return nil, false + } + return completed, true +} + +type anthropicEvent struct { + Type string `json:"type"` + Message *struct { + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role"` + Model string `json:"model,omitempty"` + StopReason string `json:"stop_reason"` + StopSequence string `json:"stop_sequence"` + Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } `json:"usage"` + } `json:"message"` + Index int `json:"index"` + ContentBlock *struct { + Type string `json:"type"` + Text string `json:"text"` + ID string `json:"id"` + Name string `json:"name"` + } `json:"content_block"` + Delta *struct { + StopReason string `json:"stop_reason"` + StopSequence string `json:"stop_sequence"` + Text string `json:"text"` + PartialJSON string `json:"partial_json"` + } `json:"delta"` + Usage *struct { + OutputTokens int `json:"output_tokens"` + } `json:"usage"` +} + +type anthropicContentBlock struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + ID string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Input map[string]any `json:"input,omitempty"` + + partialJSON strings.Builder +} + +func assembleAnthropicSSE(payloads []string) ([]byte, bool) { + var assembled struct { + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role"` + Content []anthropicContentBlock `json:"content"` + Model string `json:"model,omitempty"` + StopReason string `json:"stop_reason,omitempty"` + StopSequence string `json:"stop_sequence,omitempty"` + Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } `json:"usage"` + } + contents := map[int]*anthropicContentBlock{} + order := make([]int, 0, 1) + matched := false + + for _, payload := range payloads { + var event anthropicEvent + if err := json.Unmarshal([]byte(payload), &event); err != nil { + return nil, false + } + if !strings.HasPrefix(event.Type, "message_") && !strings.HasPrefix(event.Type, "content_block_") { + return nil, false + } + matched = true + + if event.Message != nil { + assembled.ID = event.Message.ID + assembled.Type = event.Message.Type + assembled.Role = event.Message.Role + assembled.Model = event.Message.Model + assembled.StopReason = event.Message.StopReason + assembled.StopSequence = event.Message.StopSequence + assembled.Usage.InputTokens = event.Message.Usage.InputTokens + assembled.Usage.OutputTokens = event.Message.Usage.OutputTokens + } + if event.ContentBlock != nil { + block := contents[event.Index] + if block == nil { + block = &anthropicContentBlock{Type: event.ContentBlock.Type, ID: event.ContentBlock.ID, Name: event.ContentBlock.Name} + contents[event.Index] = block + order = append(order, event.Index) + } + if event.ContentBlock.ID != "" { + block.ID = event.ContentBlock.ID + } + if event.ContentBlock.Name != "" { + block.Name = event.ContentBlock.Name + } + block.Text += event.ContentBlock.Text + } + if event.Delta != nil { + if event.Delta.Text != "" { + block := contents[event.Index] + if block == nil { + block = &anthropicContentBlock{Type: "text"} + contents[event.Index] = block + order = append(order, event.Index) + } + block.Text += event.Delta.Text + } + if event.Delta.PartialJSON != "" { + block := contents[event.Index] + if block == nil { + block = &anthropicContentBlock{Type: "tool_use"} + contents[event.Index] = block + order = append(order, event.Index) + } + block.partialJSON.WriteString(event.Delta.PartialJSON) + } + if event.Delta.StopReason != "" { + assembled.StopReason = event.Delta.StopReason + } + if event.Delta.StopSequence != "" { + assembled.StopSequence = event.Delta.StopSequence + } + } + if event.Usage != nil { + assembled.Usage.OutputTokens = event.Usage.OutputTokens + } + } + if !matched || assembled.Type == "" { + return nil, false + } + for _, index := range order { + block := contents[index] + if block.partialJSON.Len() > 0 { + var input map[string]any + if err := json.Unmarshal([]byte(block.partialJSON.String()), &input); err != nil { + return nil, false + } + block.Input = input + } + assembled.Content = append(assembled.Content, *block) + } + + data, err := json.Marshal(assembled) + return data, err == nil +} + +type geminiChunk struct { + Raw map[string]json.RawMessage `json:"-"` + Candidates []struct { + Raw map[string]json.RawMessage `json:"-"` + Content struct { + Parts []struct { + Text string `json:"text"` + } `json:"parts"` + Role string `json:"role"` + } `json:"content"` + FinishReason string `json:"finishReason"` + Index int `json:"index"` + } `json:"candidates"` + UsageMetadata json.RawMessage `json:"usageMetadata"` +} + +func (g *geminiChunk) UnmarshalJSON(data []byte) error { + type alias geminiChunk + var decoded alias + if err := json.Unmarshal(data, &decoded); err != nil { + return err + } + var raw map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + *g = geminiChunk(decoded) + g.Raw = raw + if candidatesRaw, ok := raw["candidates"]; ok { + var rawCandidates []map[string]json.RawMessage + if err := json.Unmarshal(candidatesRaw, &rawCandidates); err != nil { + return err + } + for i := range g.Candidates { + if i < len(rawCandidates) { + g.Candidates[i].Raw = rawCandidates[i] + } + } + } + return nil +} + +type geminiCandidate struct { + raw map[string]json.RawMessage + index int + role string + text strings.Builder + partRaw map[string]json.RawMessage +} + +func assembleGeminiSSE(payloads []string) ([]byte, bool) { + candidates := map[int]*geminiCandidate{} + order := make([]int, 0, 1) + var usage json.RawMessage + matched := false + + for _, payload := range payloads { + var chunk geminiChunk + if err := json.Unmarshal([]byte(payload), &chunk); err != nil { + return nil, false + } + if len(chunk.Candidates) == 0 { + return nil, false + } + matched = true + for _, candidate := range chunk.Candidates { + assembled := candidates[candidate.Index] + if assembled == nil { + assembled = &geminiCandidate{index: candidate.Index} + candidates[candidate.Index] = assembled + order = append(order, candidate.Index) + } + if candidate.Raw != nil { + assembled.raw = cloneRawMap(candidate.Raw) + } + if candidate.Content.Role != "" { + assembled.role = candidate.Content.Role + } + 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 { + Parts []map[string]json.RawMessage `json:"parts"` + } + if rawContent, ok := candidate.Raw["content"]; ok && json.Unmarshal(rawContent, &contentRaw) == nil && len(contentRaw.Parts) > 0 { + assembled.partRaw = cloneRawMap(contentRaw.Parts[0]) + } + } + } + } + if len(chunk.UsageMetadata) > 0 { + usage = chunk.UsageMetadata + } + } + if !matched { + return nil, false + } + + assembled := map[string]any{"candidates": make([]any, 0, len(order))} + for _, index := range order { + 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)} + } + candidateRaw["content"] = mustJSON(contentRaw) + assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw)) + } + if len(usage) > 0 { + assembled["usageMetadata"] = usage + } + + data, err := json.Marshal(assembled) + return data, err == nil +} + +func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage { + out := make(map[string]json.RawMessage, len(in)) + for k, v := range in { + out[k] = append(json.RawMessage(nil), v...) + } + return out +} + +func mustJSON(v any) json.RawMessage { + data, err := json.Marshal(v) + if err != nil { + return nil + } + return data +} + +func rawMapToMap(in map[string]json.RawMessage) map[string]any { + out := make(map[string]any, len(in)) + for k, v := range in { + var decoded any + if err := json.Unmarshal(v, &decoded); err != nil { + out[k] = string(v) + continue + } + out[k] = decoded + } + return out +} diff --git a/proxy/sse_event.go b/proxy/sse_event.go new file mode 100644 index 0000000..59c01b6 --- /dev/null +++ b/proxy/sse_event.go @@ -0,0 +1,65 @@ +package proxy + +import ( + "bytes" + "encoding/json" + "strings" +) + +type sseEventTracker struct { + buf []byte + terminal bool +} + +func (t *sseEventTracker) Write(p []byte) { + t.buf = append(t.buf, p...) + for { + end, separator := completeSSEEvent(t.buf) + if end < 0 { + return + } + event := t.buf[:end] + t.buf = t.buf[end+separator:] + if terminalSSEEvent(event) { + t.terminal = true + } + } +} + +func (t *sseEventTracker) Complete() bool { return len(t.buf) == 0 } + +func completeSSEEvent(buf []byte) (int, int) { + lf := bytes.Index(buf, []byte("\n\n")) + crlf := bytes.Index(buf, []byte("\r\n\r\n")) + if crlf >= 0 && (lf < 0 || crlf < lf) { + return crlf, 4 + } + if lf >= 0 { + return lf, 2 + } + return -1, 0 +} + +func terminalSSEEvent(event []byte) bool { + var data strings.Builder + for _, line := range strings.Split(strings.ReplaceAll(string(event), "\r\n", "\n"), "\n") { + if !strings.HasPrefix(line, "data:") { + continue + } + if data.Len() > 0 { + data.WriteByte('\n') + } + data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + } + payload := data.String() + if payload == "[DONE]" { + return true + } + var envelope struct { + Type string `json:"type"` + } + if json.Unmarshal([]byte(payload), &envelope) != nil { + return false + } + return envelope.Type == "message_stop" || envelope.Type == "response.completed" +} diff --git a/proxy/timeout.go b/proxy/timeout.go new file mode 100644 index 0000000..a4d9769 --- /dev/null +++ b/proxy/timeout.go @@ -0,0 +1,101 @@ +package proxy + +import ( + "io" + "net/http" + "sync" + "sync/atomic" + "time" +) + +type trackingBody struct { + io.ReadCloser + failed *atomic.Bool +} + +func (b *trackingBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if err != nil && err != io.EOF { + b.failed.Store(true) + } + return n, err +} + +func responseBodyTimeout(resp *http.Response, opts Options) time.Duration { + if resp.StatusCode == http.StatusSwitchingProtocols { + return 0 + } + if isStreamResponse(resp.Header) { + return opts.SSEIdleTimeout + } + return opts.ResponseTimeout +} + +type timeoutBody struct { + body io.ReadCloser + idle bool + timeout time.Duration + timer *time.Timer + mu sync.Mutex + done bool + sequence uint64 + closeOnce sync.Once + closeErr error +} + +func newTimeoutBody(body io.ReadCloser, timeout time.Duration, idle bool) *timeoutBody { + t := &timeoutBody{body: body, idle: idle, timeout: timeout} + t.resetLocked() + return t +} + +func (b *timeoutBody) Read(p []byte) (int, error) { + n, err := b.body.Read(p) + b.mu.Lock() + if !b.done { + if err != nil { + b.done = true + b.timer.Stop() + } else if n > 0 && b.idle { + b.resetLocked() + } + } + b.mu.Unlock() + return n, err +} + +func (b *timeoutBody) Close() error { + b.mu.Lock() + if !b.done { + b.done = true + b.sequence++ + b.timer.Stop() + } + b.mu.Unlock() + return b.closeUnderlying() +} + +func (b *timeoutBody) resetLocked() { + if b.timer != nil { + b.timer.Stop() + } + b.sequence++ + sequence := b.sequence + b.timer = time.AfterFunc(b.timeout, func() { b.expire(sequence) }) +} + +func (b *timeoutBody) expire(sequence uint64) { + b.mu.Lock() + if b.done || sequence != b.sequence { + b.mu.Unlock() + return + } + b.done = true + b.mu.Unlock() + _ = b.closeUnderlying() +} + +func (b *timeoutBody) closeUnderlying() error { + b.closeOnce.Do(func() { b.closeErr = b.body.Close() }) + return b.closeErr +} diff --git a/proxy/writer.go b/proxy/writer.go new file mode 100644 index 0000000..8a502da --- /dev/null +++ b/proxy/writer.go @@ -0,0 +1,128 @@ +package proxy + +import ( + "bufio" + "bytes" + "errors" + "net" + "net/http" + "sync" +) + +// captureWriter 包装 http.ResponseWriter，边转发边缓冲响应体。 +// 实现 http.Flusher 与 http.Hijacker 以支持 SSE/chunked/WebSocket。 +type captureWriter struct { + http.ResponseWriter + buf bytes.Buffer + max int64 + written int64 + truncated bool + status int + wroteHeader bool + hijacked bool + writeFailed bool + onHijack func(net.Conn) net.Conn + sse sseEventTracker +} + +func newCaptureWriter(w http.ResponseWriter, max int64) *captureWriter { + return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK} +} + +func (c *captureWriter) WriteHeader(code int) { + if c.wroteHeader { + return + } + c.status = code + c.wroteHeader = true + c.ResponseWriter.WriteHeader(code) +} + +func (c *captureWriter) Write(p []byte) (int, error) { + if !c.wroteHeader { + c.wroteHeader = true + } + n, err := c.ResponseWriter.Write(p) + if err != nil || n != len(p) { + c.writeFailed = true + } + if n > len(p) { + n = len(p) + } + if n > 0 { + // 仅缓冲 max 字节以内的内容。 + remaining := c.max - c.written + if remaining > 0 { + toBuf := n + if int64(toBuf) > remaining { + toBuf = int(remaining) + c.truncated = true + } + c.buf.Write(p[:toBuf]) + } else if c.max > 0 { + c.truncated = true + } + c.written += int64(n) + if isStreamResponse(c.Header()) { + c.sse.Write(p[:n]) + } + if f, ok := c.ResponseWriter.(http.Flusher); ok { + f.Flush() + } + } + return n, err +} + +func (c *captureWriter) Flush() { + if f, ok := c.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + +func (c *captureWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + h, ok := c.ResponseWriter.(http.Hijacker) + if !ok { + return nil, nil, errors.New("hijack not supported") + } + conn, rw, err := h.Hijack() + if err != nil { + c.writeFailed = true + return nil, nil, err + } + c.hijacked = true + if c.onHijack != nil { + conn = c.onHijack(conn) + } + return conn, rw, nil +} + +func (c *captureWriter) Body() []byte { return c.buf.Bytes() } +func (c *captureWriter) Truncated() bool { return c.truncated } +func (c *captureWriter) Status() int { return c.status } +func (c *captureWriter) Hijacked() bool { return c.hijacked } +func (c *captureWriter) Complete() bool { return !c.writeFailed } +func (c *captureWriter) SSEComplete() bool { return c.sse.Complete() } + +func (c *captureWriter) OnHijack(fn func(net.Conn) net.Conn) { c.onHijack = fn } + +func (c *captureWriter) SetHijackedResponse(status int, header http.Header) { + c.status = status + for key := range c.Header() { + c.Header().Del(key) + } + for key, values := range header { + c.Header()[key] = append([]string(nil), values...) + } +} + +type trackedConn struct { + net.Conn + closeOnce sync.Once + onClose func() +} + +func (c *trackedConn) Close() error { + err := c.Conn.Close() + c.closeOnce.Do(c.onClose) + return err +} diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..bc62a9b --- /dev/null +++ b/tests/README.md @@ -0,0 +1,61 @@ +# tests/ + +集中放置项目的测试代码。每个子目录对应被测包，使用 Go 的黑盒测试包模式（`package xxx_test`）通过被测包的导出 API 进行验证。 + +## 目录结构 + +``` +tests/ +├── config/ # 配置与 glob 过滤器测试 +├── logger/ # 异步日志队列与重试逻辑测试 +├── proxy/ # 反向代理测试（WebSocket、SSE、上游错误） +└── scripts/ # 手工端到端测试脚本与 DB 校验工具 +``` + +## 运行 + +```bash +go test ./tests/... +``` + +或者运行某一个子包： + +```bash +go test ./tests/proxy/... +``` + +## 端到端冒烟测试 + +先构建本地二进制： + +```powershell +go build -o TokenThief.exe . +``` + +加载 `.env` 并设置 newapi Key： + +```powershell +. .\tests\scripts\load-env.ps1 +$env:NEWAPI_KEY = "sk-..." +``` + +`.env` 需要包含可写的 `CLICKHOUSE_URL`，例如 `clickhouse://tokenthief:tokenthief@localhost:9000/tokenthief`。 + +运行聊天端点冒烟测试： + +```powershell +.\tests\scripts\smoke.ps1 -Model "gpt-5.4-mini" +``` + +按 request_id 打印数据库里的流式 `response_body`： + +```powershell +.\tests\scripts\dump-stream-body.ps1 -RequestID "" +``` + +## 构建排除 + +- 测试目录中的 `_test.go` 文件不会参与 `go build`。 +- `tests/scripts/` 只放手工测试脚本，不被主程序 import。 +- `Dockerfile` 仅构建 main 包（`./`），不会触及 `tests/`。 +- 仓库根目录的 `.dockerignore` 把 `tests/` 整体排除在 build context 之外，镜像中不会包含测试代码或测试脚本。 diff --git a/tests/config/config_test.go b/tests/config/config_test.go new file mode 100644 index 0000000..21379a2 --- /dev/null +++ b/tests/config/config_test.go @@ -0,0 +1,283 @@ +package config_test + +import ( + "net/netip" + "strings" + "testing" + "time" + + "git.misaka.ren/M1saka/token_thief/config" +) + +func TestLoadTimeoutDefaults(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + + if cfg.ReadTimeout != 30*time.Second { + t.Fatalf("ReadTimeout=%s want 30s", cfg.ReadTimeout) + } + if cfg.WriteTimeout != 10*time.Minute { + t.Fatalf("WriteTimeout=%s want 10m", cfg.WriteTimeout) + } + if cfg.IdleTimeout != 5*time.Minute { + t.Fatalf("IdleTimeout=%s want 5m", cfg.IdleTimeout) + } + if cfg.UpstreamTimeout != 30*time.Second { + t.Fatalf("UpstreamTimeout=%s want 30s", cfg.UpstreamTimeout) + } + if cfg.UpstreamResponseTimeout != 30*time.Second { + t.Fatalf("UpstreamResponseTimeout=%s want 30s", cfg.UpstreamResponseTimeout) + } + if cfg.UpstreamStreamIdleTimeout != 2*time.Minute { + t.Fatalf("UpstreamStreamIdleTimeout=%s want 2m", cfg.UpstreamStreamIdleTimeout) + } +} + +func TestLoadQueueDefaults(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + + if cfg.MaxBodyBytes != 1<<20 { + t.Fatalf("MaxBodyBytes=%d want %d", cfg.MaxBodyBytes, 1<<20) + } + if cfg.LogQueueSize != 256 { + t.Fatalf("LogQueueSize=%d want 256", cfg.LogQueueSize) + } + if cfg.LogQueueBytes != 64<<20 { + t.Fatalf("LogQueueBytes=%d want %d", cfg.LogQueueBytes, 64<<20) + } +} + +func TestLoadTimeoutOverrides(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + t.Setenv("READ_TIMEOUT", "10s") + t.Setenv("WRITE_TIMEOUT", "20s") + t.Setenv("IDLE_TIMEOUT", "30s") + t.Setenv("UPSTREAM_TIMEOUT", "40s") + t.Setenv("UPSTREAM_RESPONSE_TIMEOUT", "50s") + t.Setenv("UPSTREAM_STREAM_IDLE_TIMEOUT", "60s") + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + + if cfg.ReadTimeout != 10*time.Second { + t.Fatalf("ReadTimeout=%s want 10s", cfg.ReadTimeout) + } + if cfg.WriteTimeout != 20*time.Second { + t.Fatalf("WriteTimeout=%s want 20s", cfg.WriteTimeout) + } + if cfg.IdleTimeout != 30*time.Second { + t.Fatalf("IdleTimeout=%s want 30s", cfg.IdleTimeout) + } + if cfg.UpstreamTimeout != 40*time.Second { + t.Fatalf("UpstreamTimeout=%s want 40s", cfg.UpstreamTimeout) + } + if cfg.UpstreamResponseTimeout != 50*time.Second { + t.Fatalf("UpstreamResponseTimeout=%s want 50s", cfg.UpstreamResponseTimeout) + } + if cfg.UpstreamStreamIdleTimeout != 60*time.Second { + t.Fatalf("UpstreamStreamIdleTimeout=%s want 60s", cfg.UpstreamStreamIdleTimeout) + } +} + +func TestLoadRejectsInvalidTypedValues(t *testing.T) { + tests := []struct { + name string + key string + value string + }{ + {name: "int", key: "LOG_QUEUE_SIZE", value: "not-an-int"}, + {name: "int64", key: "MAX_BODY_BYTES", value: "1.5"}, + {name: "duration", key: "READ_TIMEOUT", value: "30"}, + {name: "bool", key: "UPSTREAM_TLS_INSECURE_SKIP_VERIFY", value: "yes"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + t.Setenv(tt.key, tt.value) + + _, err := config.Load() + if err == nil { + t.Fatalf("Load succeeded with %s=%q", tt.key, tt.value) + } + if !strings.Contains(err.Error(), tt.key) { + t.Fatalf("err=%q does not identify %s", err, tt.key) + } + }) + } +} + +func TestLoadRejectsNonPositiveNumericAndDurationValues(t *testing.T) { + tests := []struct { + key string + value string + }{ + {key: "MAX_BODY_BYTES", value: "0"}, + {key: "LOG_QUEUE_SIZE", value: "-1"}, + {key: "LOG_QUEUE_BYTES", value: "0"}, + {key: "LOG_BATCH_SIZE", value: "0"}, + {key: "LOG_BATCH_INTERVAL", value: "-1s"}, + {key: "LOG_WORKERS", value: "0"}, + {key: "DB_RECONNECT_INTERVAL", value: "0s"}, + {key: "READ_TIMEOUT", value: "0s"}, + {key: "WRITE_TIMEOUT", value: "0s"}, + {key: "IDLE_TIMEOUT", value: "0s"}, + {key: "UPSTREAM_TIMEOUT", value: "0s"}, + {key: "UPSTREAM_RESPONSE_TIMEOUT", value: "0s"}, + {key: "UPSTREAM_STREAM_IDLE_TIMEOUT", value: "0s"}, + } + + for _, tt := range tests { + t.Run(tt.key, func(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + t.Setenv(tt.key, tt.value) + + _, err := config.Load() + if err == nil { + t.Fatalf("Load succeeded with %s=%q", tt.key, tt.value) + } + if !strings.Contains(err.Error(), tt.key) { + t.Fatalf("err=%q does not identify %s", err, tt.key) + } + }) + } +} + +func TestLoadTrustedProxies(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + t.Setenv("TRUSTED_PROXIES", "10.0.0.0/8, 192.0.2.1,2001:db8::/32, 2001:db8::1") + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + + want := []netip.Prefix{ + netip.MustParsePrefix("10.0.0.0/8"), + netip.MustParsePrefix("192.0.2.1/32"), + netip.MustParsePrefix("2001:db8::/32"), + netip.MustParsePrefix("2001:db8::1/128"), + } + if len(cfg.TrustedProxies) != len(want) { + t.Fatalf("TrustedProxies=%v want %v", cfg.TrustedProxies, want) + } + for i := range want { + if cfg.TrustedProxies[i] != want[i] { + t.Fatalf("TrustedProxies[%d]=%v want %v", i, cfg.TrustedProxies[i], want[i]) + } + } +} + +func TestLoadRejectsInvalidTrustedProxies(t *testing.T) { + for _, value := range []string{"not-an-ip", "10.0.0.0/33", "10.0.0.1/8", "10.0.0.1,,192.0.2.1"} { + t.Run(value, func(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + t.Setenv("TRUSTED_PROXIES", value) + + _, err := config.Load() + if err == nil { + t.Fatalf("Load succeeded with TRUSTED_PROXIES=%q", value) + } + }) + } +} + +func TestLoadUpstreamTLSInsecureSkipVerifyDefault(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + + if cfg.UpstreamTLSInsecureSkipVerify { + t.Fatal("UpstreamTLSInsecureSkipVerify=true want false") + } +} + +func TestLoadUpstreamTLSInsecureSkipVerifyOverride(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db") + t.Setenv("UPSTREAM_TLS_INSECURE_SKIP_VERIFY", "true") + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + + if !cfg.UpstreamTLSInsecureSkipVerify { + t.Fatal("UpstreamTLSInsecureSkipVerify=false want true") + } +} + +func TestLoadRequiresClickHouseURL(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("DATABASE_URL", "postgres://user:pass@localhost/db") + + _, err := config.Load() + if err == nil { + t.Fatal("Load succeeded without CLICKHOUSE_URL") + } + if err.Error() != "CLICKHOUSE_URL is required" { + t.Fatalf("err=%q want CLICKHOUSE_URL is required", err.Error()) + } +} + +func TestLoadValidatesClickHouseURL(t *testing.T) { + tests := []struct { + name string + dsn string + }{ + {name: "missing port", dsn: "clickhouse://user:pass@localhost/db"}, + {name: "unsupported scheme", dsn: "https://localhost:9440/db"}, + {name: "plaintext skip verify", dsn: "clickhouse://localhost:9000/db?skip_verify=true"}, + {name: "conflicting secure parameter", dsn: "clickhouses://localhost:9440/db?secure=false"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", tt.dsn) + + _, err := config.Load() + if err == nil { + t.Fatalf("Load succeeded with CLICKHOUSE_URL=%q", tt.dsn) + } + if !strings.Contains(err.Error(), "CLICKHOUSE_URL") { + t.Fatalf("err=%q does not identify CLICKHOUSE_URL", err) + } + }) + } +} + +func TestLoadAcceptsSecureClickHouseURL(t *testing.T) { + t.Setenv("UPSTREAM_URL", "https://example.com") + t.Setenv("CLICKHOUSE_URL", "clickhouses://user:pass@localhost:9440/db?skip_verify=false") + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.ClickHouseURL != "clickhouses://user:pass@localhost:9440/db?skip_verify=false" { + t.Fatalf("ClickHouseURL=%q", cfg.ClickHouseURL) + } +} diff --git a/tests/config/filter_test.go b/tests/config/filter_test.go new file mode 100644 index 0000000..064c023 --- /dev/null +++ b/tests/config/filter_test.go @@ -0,0 +1,164 @@ +package config_test + +import ( + "os" + "path/filepath" + "testing" + + "git.misaka.ren/M1saka/token_thief/config" +) + +func TestCompileGlobMatch(t *testing.T) { + cases := []struct { + pattern string + target string + want bool + }{ + // 单段 * + {"/v1/audio/*", "/v1/audio/speech", true}, + {"/v1/audio/*", "/v1/audio/transcriptions", true}, + {"/v1/audio/*", "/v1/audio/sub/x", false}, + + // 字面匹配 + {"/v1/chat/completions", "/v1/chat/completions", true}, + {"/v1/chat/completions", "/v1/chat/completions/extra", false}, + + // 跨段 ** + {"/v1/videos/**", "/v1/videos/", true}, + {"/v1/videos/**", "/v1/videos/abc", true}, + {"/v1/videos/**", "/v1/videos/abc/content", true}, + + // Gemini 风格 :generateContent + {"/v1beta/models/*:generateContent", "/v1beta/models/gemini-pro:generateContent", true}, + {"/v1beta/models/*:generateContent", "/v1beta/models/gemini-1.5-flash:generateContent", true}, + {"/v1beta/models/*:generateContent", "/v1beta/models/x/y:generateContent", false}, + + // engines 嵌套 + {"/v1/engines/*/embeddings", "/v1/engines/text-embedding-ada-002/embeddings", true}, + {"/v1/engines/*/embeddings", "/v1/engines/a/b/embeddings", false}, + + // 含正则元字符 + {"/v1/models/*", "/v1/models/gpt-4.1", true}, + } + for _, c := range cases { + re, err := config.CompileGlob(c.pattern) + if err != nil { + t.Fatalf("compile %q: %v", c.pattern, err) + } + got := re.MatchString(c.target) + if got != c.want { + t.Errorf("%q vs %q: got %v want %v (regex=%s)", c.pattern, c.target, got, c.want, re.String()) + } + } +} + +func TestFilterShouldLog(t *testing.T) { + f, err := config.NewFilter(config.FilterWhitelist, []string{"/v1/chat/completions", "/v1/videos/*"}) + if err != nil { + t.Fatalf("NewFilter: %v", err) + } + if !f.ShouldLog("/v1/chat/completions") { + t.Error("whitelist must allow /v1/chat/completions") + } + if !f.ShouldLog("/v1/videos/abc") { + t.Error("whitelist must allow /v1/videos/abc") + } + if f.ShouldLog("/v1/embeddings") { + t.Error("whitelist must reject /v1/embeddings") + } + + fb, err := config.NewFilter(config.FilterBlacklist, []string{"/v1/chat/completions", "/v1/videos/*"}) + if err != nil { + t.Fatalf("NewFilter blacklist: %v", err) + } + if fb.ShouldLog("/v1/chat/completions") { + t.Error("blacklist must reject /v1/chat/completions") + } + if !fb.ShouldLog("/v1/embeddings") { + t.Error("blacklist must allow /v1/embeddings") + } + + fd, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatalf("NewFilter disabled: %v", err) + } + if !fd.ShouldLog("/anything") { + t.Error("disabled must allow everything") + } +} + +// TestRepoFilterOnlyAllowsChatAndCompletions 加载仓库根目录的 filter.yaml， +// 验证默认过滤器只记录聊天与补全端点。 +func TestRepoFilterOnlyAllowsChatAndCompletions(t *testing.T) { + cwd, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + // 测试目录位于 /tests/config，filter.yaml 在 /filter.yaml + path := filepath.Join(cwd, "..", "..", "filter.yaml") + if _, err := os.Stat(path); err != nil { + t.Fatalf("filter.yaml not found: %v", err) + } + f, err := config.LoadFilter(path) + if err != nil { + t.Fatalf("load filter: %v", err) + } + + allowed := []string{ + // Chat + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1beta/models/gemini-1.5-pro:generateContent", + "/v1beta/models/gemini-1.5-pro:generateContent/", + "/v1beta/models/gemini-1.5-pro:streamGenerateContent", + // Completions + "/v1/completions", + } + for _, p := range allowed { + if !f.ShouldLog(p) { + t.Errorf("filter.yaml should match chat/completion endpoint %q", p) + } + } + + blocked := []string{ + // Models + "/v1/models", + "/v1beta/models", + // Embeddings + "/v1/embeddings", + "/v1/engines/text-embedding-ada-002/embeddings", + // Moderations / Rerank / Realtime + "/v1/moderations", + "/v1/rerank", + "/v1/realtime", + // Audio + "/v1/audio/speech", + "/v1/audio/transcriptions", + "/v1/audio/translations", + // Images + "/v1/images/generations", + "/v1/images/generations/", + "/v1/images/edits", + "/v1/images/edits/", + // Videos - 通用 + "/v1/video/generations", + "/v1/video/generations/task_abc", + // Videos - Sora + "/v1/videos", + "/v1/videos/task_abc", + "/v1/videos/task_abc/content", + // Videos - 即梦 + "/jimeng/", + // Videos - Kling + "/kling/v1/videos/text2video", + "/kling/v1/videos/text2video/task_abc", + "/kling/v1/videos/image2video", + "/kling/v1/videos/image2video/task_abc", + } + for _, p := range blocked { + if f.ShouldLog(p) { + t.Errorf("filter.yaml should not match non-chat/completion endpoint %q", p) + } + } +} diff --git a/tests/deployment/compose_test.go b/tests/deployment/compose_test.go new file mode 100644 index 0000000..f5e4cc5 --- /dev/null +++ b/tests/deployment/compose_test.go @@ -0,0 +1,66 @@ +package deployment_test + +import ( + "os" + "path/filepath" + "regexp" + "strings" + "testing" +) + +func TestComposeUsesSafeDeploymentDefaults(t *testing.T) { + root := filepath.Join("..", "..") + compose := readFile(t, filepath.Join(root, "compose.yml")) + envExample := readFile(t, filepath.Join(root, ".env.example")) + + for _, required := range []string{ + `${UPSTREAM_URL:?set UPSTREAM_URL}`, + `${CLICKHOUSE_URL:?set CLICKHOUSE_URL with URL-encoded credentials}`, + `${CLICKHOUSE_PASSWORD:?set a strong CLICKHOUSE_PASSWORD}`, + } { + if !strings.Contains(compose, required) { + t.Errorf("compose.yml must reject an empty required setting with %q", required) + } + } + + if !regexp.MustCompile(`(?m)^\s+image: clickhouse/clickhouse-server:\d+\.\d+\.\d+\.\d+-alpine\s*$`).MatchString(compose) { + t.Error("ClickHouse image must use an exact version tag") + } + + serviceStart := strings.LastIndex(compose, " thief_clickhouse:") + if serviceStart < 0 { + t.Fatal("compose.yml is missing the thief_clickhouse service") + } + clickhouseService := compose[serviceStart:] + if strings.Contains(clickhouseService, "\n ports:") { + t.Error("ClickHouse must not publish host ports by default") + } + + for _, emptySecret := range []string{"CLICKHOUSE_URL=\n", "CLICKHOUSE_PASSWORD=\n"} { + if !strings.Contains(envExample, emptySecret) { + t.Errorf(".env.example must leave %q empty", strings.TrimSpace(emptySecret)) + } + } +} + +func TestReviewDocumentsUseStablePaths(t *testing.T) { + root := filepath.Join("..", "..") + for _, path := range []string{ + filepath.Join(root, "docs", "compose", "specs", "reliability-security-fixes.md"), + filepath.Join(root, "docs", "compose", "plans", "reliability-security-fixes.md"), + } { + content := readFile(t, path) + if !strings.Contains(content, "2026-07-09-clickhouse-migration.md") { + t.Errorf("%s must link to the dated source document", path) + } + } +} + +func readFile(t *testing.T, path string) string { + t.Helper() + b, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + return string(b) +} diff --git a/tests/logger/queue_test.go b/tests/logger/queue_test.go new file mode 100644 index 0000000..6d77587 --- /dev/null +++ b/tests/logger/queue_test.go @@ -0,0 +1,323 @@ +package logger_test + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "git.misaka.ren/M1saka/token_thief/logger" +) + +type batchScript struct { + prepareErr error + appendAt int + appendErr error + sendErr error +} + +type fakeBatch struct { + script batchScript + rows [][]any + appendCall int + abortCalls int +} + +func (f *fakeBatch) Append(v ...any) error { + f.appendCall++ + if f.script.appendErr != nil && f.appendCall == f.script.appendAt { + return f.script.appendErr + } + f.rows = append(f.rows, append([]any(nil), v...)) + return nil +} + +func (f *fakeBatch) Send() error { return f.script.sendErr } +func (f *fakeBatch) Abort() error { f.abortCalls++; return nil } + +type fakePool struct { + calls int + queries []string + batches []*fakeBatch + scripts []batchScript +} + +func (p *fakePool) PrepareBatch(_ context.Context, query string) (logger.Batch, error) { + idx := p.calls + p.calls++ + p.queries = append(p.queries, query) + var script batchScript + if idx < len(p.scripts) { + script = p.scripts[idx] + } + if script.prepareErr != nil { + return nil, script.prepareErr + } + b := &fakeBatch{script: script} + p.batches = append(p.batches, b) + return b, nil +} + +func TestFlushSuccessUsesClickHouseTypes(t *testing.T) { + p := &fakePool{} + started := time.Date(2026, 7, 9, 1, 2, 3, 4_000_000, time.UTC) + entry := &logger.LogEntry{ + RequestID: "rid", Method: "POST", Path: "/v1", Query: "a=b", ClientIP: "127.0.0.1", + RequestHeaders: []byte(`{"x":"y"}`), RequestBody: []byte("request"), RequestTruncated: true, + StatusCode: 201, ResponseHeaders: []byte(`{"h":"v"}`), ResponseBody: []byte("response"), + ResponseTruncated: true, IsStream: true, LatencyMS: 150, StartedAt: started, + FinishedAt: started.Add(150 * time.Millisecond), + } + + result := logger.Flush(context.Background(), p, []*logger.LogEntry{entry}) + if result.Err != nil || result.Failed != 0 || len(result.Retry) != 0 { + t.Fatalf("unexpected result: %+v", result) + } + if p.calls != 1 || len(p.batches) != 1 || len(p.batches[0].rows) != 1 { + t.Fatalf("calls=%d batches=%d rows=%d", p.calls, len(p.batches), len(p.batches[0].rows)) + } + if contains(p.queries[0], "VALUES") { + t.Fatalf("query contains VALUES: %s", p.queries[0]) + } + row := p.batches[0].rows[0] + if len(row) != 17 { + t.Fatalf("columns=%d want 17", len(row)) + } + if row[5] != string(entry.RequestHeaders) || row[6] != string(entry.RequestBody) || row[9] != string(entry.ResponseHeaders) || row[10] != string(entry.ResponseBody) { + t.Fatalf("byte fields were not converted to strings: %#v", row) + } + if _, ok := row[8].(int32); !ok { + t.Fatalf("status code type=%T want int32", row[8]) + } + if p.batches[0].abortCalls != 0 { + t.Fatalf("abort calls=%d want 0", p.batches[0].abortCalls) + } +} + +func TestFlushEmptyBatchDoesNotPrepare(t *testing.T) { + p := &fakePool{} + + result := logger.Flush(context.Background(), p, nil) + if result.Err != nil || result.Failed != 0 || result.Ambiguous != 0 || len(result.Retry) != 0 { + t.Fatalf("unexpected result: %+v", result) + } + if p.calls != 0 { + t.Fatalf("PrepareBatch calls=%d want 0", p.calls) + } +} + +func TestFlushPrepareFailureIsRetryable(t *testing.T) { + p := &fakePool{scripts: []batchScript{{prepareErr: errors.New("prepare")}}} + entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}} + result := logger.Flush(context.Background(), p, entries) + if result.Err == nil || result.Failed != 0 || len(result.Retry) != 2 { + t.Fatalf("unexpected result: %+v", result) + } + if len(p.batches) != 0 { + t.Fatalf("prepare failure created %d batches", len(p.batches)) + } +} + +func TestFlushAppendFailureIsolatesOnlyBadRow(t *testing.T) { + appendErr := errors.New("bad row") + p := &fakePool{scripts: []batchScript{ + {appendAt: 2, appendErr: appendErr}, + {}, + {appendAt: 1, appendErr: appendErr}, + {}, + }} + entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "bad"}, {RequestID: "c"}} + + result := logger.Flush(context.Background(), p, entries) + if result.Err == nil || result.Failed != 1 || len(result.Retry) != 0 { + t.Fatalf("unexpected result: %+v", result) + } + if p.calls != 4 { + t.Fatalf("PrepareBatch calls=%d want 4", p.calls) + } + if p.batches[0].abortCalls != 1 || p.batches[2].abortCalls != 1 { + t.Fatalf("abort calls initial=%d bad-row=%d want 1 each", p.batches[0].abortCalls, p.batches[2].abortCalls) + } + if p.batches[1].abortCalls != 0 || p.batches[3].abortCalls != 0 { + t.Fatalf("successful batches were aborted") + } +} + +func TestFlushAppendFailureCanFullyRecover(t *testing.T) { + p := &fakePool{scripts: []batchScript{ + {appendAt: 2, appendErr: errors.New("batch append")}, + {}, + {}, + }} + entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}} + + result := logger.Flush(context.Background(), p, entries) + if result.Err != nil || result.Failed != 0 || len(result.Retry) != 0 { + t.Fatalf("unexpected result: %+v", result) + } + if p.calls != 3 || p.batches[0].abortCalls != 1 { + t.Fatalf("calls=%d abort=%d", p.calls, p.batches[0].abortCalls) + } +} + +func TestFlushSendFailureIsAmbiguousAndNotRetried(t *testing.T) { + p := &fakePool{scripts: []batchScript{{sendErr: errors.New("connection lost")}}} + entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}} + + result := logger.Flush(context.Background(), p, entries) + if result.Err == nil || result.Failed != 2 || len(result.Retry) != 0 { + t.Fatalf("unexpected result: %+v", result) + } + if result.Ambiguous != 2 { + t.Fatalf("Ambiguous=%d want 2", result.Ambiguous) + } + if p.calls != 1 { + t.Fatalf("PrepareBatch calls=%d want 1", p.calls) + } + if p.batches[0].abortCalls != 1 { + t.Fatalf("abort calls=%d want 1", p.batches[0].abortCalls) + } +} + +func TestFlushAppendRecoveryCountsAmbiguousSingleSend(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"}} + + result := logger.Flush(context.Background(), p, entries) + if result.Failed != 1 || result.Ambiguous != 1 || len(result.Retry) != 0 { + t.Fatalf("unexpected result: %+v", result) + } +} + +func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) { + entry := &logger.LogEntry{ + RequestID: "1", Method: "22", Path: "333", Query: "4444", ClientIP: "55555", Error: "666666", + 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) + } +} + +func TestQueueByteBudgetReleasedAfterFinalDiscard(t *testing.T) { + entry := &logger.LogEntry{ + RequestID: "request", Method: "POST", Path: "/path", Query: "q=1", ClientIP: "ip", Error: "error", + RequestHeaders: []byte("rh"), RequestBody: []byte("rb"), ResponseHeaders: []byte("sh"), ResponseBody: []byte("sb"), + } + budget := logger.EstimatedBytes(entry) + q := logger.NewQueue(nil, 4, 2, 1, time.Hour, budget) + + q.Submit(entry) + q.Submit(entry) + stats := q.Stats() + if stats.Enqueued != 1 || stats.Dropped != 1 || stats.Bytes != budget { + t.Fatalf("before drain: %+v budget=%d", stats, budget) + } + q.Start(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := q.Stop(ctx); err != nil { + t.Fatalf("Stop: %v", err) + } + stats = q.Stats() + if stats.Bytes != 0 || stats.Failed != 1 { + t.Fatalf("after drain: %+v", stats) + } +} + +func TestQueueStopBeforeStartReleasesBudget(t *testing.T) { + entry := &logger.LogEntry{RequestID: "queued"} + q := logger.NewQueue(nil, 2, 2, 1, time.Hour, logger.EstimatedBytes(entry)) + q.Submit(entry) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := q.Stop(ctx); err != nil { + t.Fatalf("Stop: %v", err) + } + if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 1 { + t.Fatalf("stats=%+v", stats) + } +} + +func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) { + entry := &logger.LogEntry{RequestID: "queued"} + q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry)) + q.Start(context.Background()) + for i := 0; i < 32; i++ { + q.Submit(entry) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _ = 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 { + t.Fatalf("stats=%+v", stats) + } +} + +func TestQueueStopAndSubmitAreConcurrentAndRepeatSafe(t *testing.T) { + q := logger.NewQueue(nil, 32, 8, 2, time.Millisecond, 1<<20) + q.Start(context.Background()) + + var submitters sync.WaitGroup + for i := 0; i < 8; i++ { + submitters.Add(1) + go func() { + defer submitters.Done() + for j := 0; j < 2_000; j++ { + q.Submit(&logger.LogEntry{RequestID: "x"}) + } + }() + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := q.Stop(ctx); err != nil { + t.Fatalf("first Stop: %v", err) + } + submitters.Wait() + if err := q.Stop(ctx); err != nil { + t.Fatalf("second Stop: %v", err) + } + q.Submit(&logger.LogEntry{RequestID: "after-stop"}) + if stats := q.Stats(); stats.Bytes != 0 { + t.Fatalf("bytes after Stop=%d want 0", stats.Bytes) + } +} + +func TestStopIsIndependentFromStartContext(t *testing.T) { + root, cancelRoot := context.WithCancel(context.Background()) + q := logger.NewQueue(nil, 4, 4, 1, time.Hour, 1024) + q.Start(root) + q.Submit(&logger.LogEntry{RequestID: "x"}) + cancelRoot() + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := q.Stop(ctx); err != nil { + t.Fatalf("Stop after root cancellation: %v", err) + } + if stats := q.Stats(); stats.Failed != 1 || stats.Bytes != 0 { + t.Fatalf("stats=%+v", stats) + } +} + +func contains(s, substr string) bool { + for i := 0; i+len(substr) <= len(s); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/tests/proxy/error_test.go b/tests/proxy/error_test.go new file mode 100644 index 0000000..63c1a27 --- /dev/null +++ b/tests/proxy/error_test.go @@ -0,0 +1,139 @@ +package proxy_test + +import ( + "io" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" + + "git.misaka.ren/M1saka/token_thief/config" + "git.misaka.ren/M1saka/token_thief/logger" + "git.misaka.ren/M1saka/token_thief/proxy" +) + +type sliceSubmitter struct{ entries []*logger.LogEntry } + +func (s *sliceSubmitter) Submit(e *logger.LogEntry) { s.entries = append(s.entries, e) } + +// TestUpstreamErrorRecorded 验证上游不可达时 502 响应被记录、错误信息进入 LogEntry.Error。 +func TestUpstreamErrorRecorded(t *testing.T) { + // 指向一个一定不可用的端口 + badURL, _ := url.Parse("http://127.0.0.1:1") // port 1 几乎肯定 connection refused + + sub := &sliceSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(badURL, filter, sub, 1024) + + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/v1/chat/completions") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusBadGateway { + t.Errorf("status=%d want 502, body=%q", resp.StatusCode, body) + } + if resp.Header.Get("X-Request-Id") == "" { + t.Errorf("missing X-Request-Id header") + } + + deadline := time.Now().Add(2 * time.Second) + for len(sub.entries) == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if len(sub.entries) == 0 { + t.Fatal("no log entry captured") + } + e := sub.entries[0] + if e.StatusCode != http.StatusBadGateway { + t.Errorf("LogEntry.StatusCode=%d want 502", e.StatusCode) + } + if e.Error == "" { + t.Errorf("LogEntry.Error should be set, got empty") + } + if string(e.ResponseBody) != "bad gateway\n" { + t.Errorf("response_body should be fixed bad gateway, got %q", e.ResponseBody) + } + if e.RequestID == "" { + t.Errorf("LogEntry.RequestID empty") + } +} + +// TestModifyResponseSetsRequestID 验证正常上游响应也会带上 X-Request-Id。 +func TestModifyResponseSetsRequestID(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(200) + _, _ = w.Write([]byte(`{"ok":true}`)) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &sliceSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(u, filter, sub, 1024) + + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/v1/anything") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + rid := resp.Header.Get("X-Request-Id") + if len(rid) != 32 { + t.Errorf("X-Request-Id length=%d want 32, value=%q", len(rid), rid) + } + + deadline := time.Now().Add(2 * time.Second) + for len(sub.entries) == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if len(sub.entries) == 0 { + t.Fatal("no log entry") + } + if sub.entries[0].RequestID != rid { + t.Errorf("LogEntry.RequestID=%q response header=%q (should match)", sub.entries[0].RequestID, rid) + } +} + +func TestUpstreamTLSInsecureSkipVerifyAllowsSelfSigned(t *testing.T) { + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"ok":true}`)) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &sliceSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.NewWithOptions(u, filter, sub, 1024, proxy.Options{UpstreamTLSInsecureSkipVerify: true}) + + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/v1/anything") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + t.Fatalf("status=%d want 200, body=%q", resp.StatusCode, body) + } +} diff --git a/tests/proxy/robustness_test.go b/tests/proxy/robustness_test.go new file mode 100644 index 0000000..2db21d4 --- /dev/null +++ b/tests/proxy/robustness_test.go @@ -0,0 +1,336 @@ +package proxy_test + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "net/netip" + "net/url" + "strings" + "sync/atomic" + "testing" + "time" + + "git.misaka.ren/M1saka/token_thief/config" + "git.misaka.ren/M1saka/token_thief/proxy" +) + +type failingBody struct{ err error } + +func (b failingBody) Read([]byte) (int, error) { return 0, b.err } +func (failingBody) Close() error { return nil } + +type lateFailingBody struct { + remaining int + err error +} + +func (b *lateFailingBody) Read(p []byte) (int, error) { + if b.remaining == 0 { + return 0, b.err + } + n := min(len(p), b.remaining) + for i := range p[:n] { + p[i] = 'x' + } + b.remaining -= n + return n, nil +} + +func (*lateFailingBody) Close() error { return nil } + +type failingResponseWriter struct { + header http.Header + short bool +} + +func (w *failingResponseWriter) Header() http.Header { return w.header } +func (*failingResponseWriter) WriteHeader(int) {} +func (w *failingResponseWriter) Write(p []byte) (int, error) { + if w.short { + return len(p) - 1, nil + } + return 0, errors.New("write failed") +} + +func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler { + t.Helper() + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + return proxy.NewWithOptions(upstream, filter, sub, 1024*1024, opts) +} + +func TestRequestBodyReadFailureReturnsFixed400WithoutUpstream(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 = 1 + 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) { + var upstreamCalls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + })) + 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 = &lateFailingBody{remaining: 1024*1024 + 1, err: errors.New("late read failure")} + req.ContentLength = 1024*1024 + 2 + 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 TestFilteredRequestBodyReadFailureDoesNotReachUpstream(t *testing.T) { + var upstreamCalls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls.Add(1) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + filter, err := config.NewFilter(config.FilterBlacklist, []string{"/ignored"}) + if err != nil { + t.Fatal(err) + } + h := proxy.NewWithOptions(u, filter, &captureSubmitter{}, 1024, proxy.Options{}) + + req := httptest.NewRequest(http.MethodPost, "http://proxy.test/ignored", nil) + req.Body = failingBody{err: errors.New("read failure")} + req.ContentLength = 1 + 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 TestBadGatewayResponseDoesNotLeakUpstreamError(t *testing.T) { + badURL, _ := url.Parse("http://127.0.0.1:1") + h := newTestHandler(t, badURL, &captureSubmitter{}, proxy.Options{}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil)) + + if rec.Code != http.StatusBadGateway || rec.Body.String() != "bad gateway\n" { + t.Fatalf("response=(%d, %q), want fixed 502 bad gateway", rec.Code, rec.Body.String()) + } +} + +func TestResponseWriteFailurePreventsCommit(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "response") + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + for _, short := range []bool{false, true} { + sub := &captureSubmitter{} + h := newTestHandler(t, u, sub, proxy.Options{}) + w := &failingResponseWriter{header: make(http.Header), short: short} + h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil)) + if sub.Len() != 0 { + t.Fatalf("short=%v: failed response write was committed", short) + } + } +} + +func TestTrustedProxyClientIP(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "ok") + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + trusted := netip.MustParsePrefix("10.0.0.0/8") + + tests := []struct { + name string + peer string + xff string + wantIP string + }{ + {name: "untrusted peer ignores xff", peer: "203.0.113.9:1234", xff: "198.51.100.1", wantIP: "203.0.113.9"}, + {name: "strip trusted from right", peer: "10.0.0.2:1234", xff: "198.51.100.7, 10.0.0.3", wantIP: "198.51.100.7"}, + {name: "ipv6", peer: "[2001:db8::2]:1234", xff: "198.51.100.7", wantIP: "2001:db8::2"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + sub := &captureSubmitter{} + h := newTestHandler(t, u, sub, proxy.Options{TrustedProxies: []netip.Prefix{trusted}}) + req := httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil) + req.RemoteAddr = tc.peer + req.Header.Set("X-Forwarded-For", tc.xff) + h.ServeHTTP(httptest.NewRecorder(), req) + if sub.Len() != 1 || sub.Entry(0).ClientIP != tc.wantIP { + t.Fatalf("ClientIP=%q, want %q", sub.Entry(0).ClientIP, tc.wantIP) + } + }) + } +} + +func TestTrustedProxyUsesValidXRealIPWithoutXFF(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + sub := &captureSubmitter{} + trusted := netip.MustParsePrefix("192.0.2.0/24") + h := newTestHandler(t, u, sub, proxy.Options{TrustedProxies: []netip.Prefix{trusted}}) + req := httptest.NewRequest(http.MethodGet, "http://proxy.test/x", nil) + req.RemoteAddr = "192.0.2.10:1234" + req.Header.Set("X-Real-IP", "198.51.100.20") + h.ServeHTTP(httptest.NewRecorder(), req) + if sub.Len() != 1 || sub.Entry(0).ClientIP != "198.51.100.20" { + t.Fatalf("entries=%d", sub.Len()) + } +} + +func TestResponseBodyTimeoutPreventsCommit(t *testing.T) { + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-release + })) + u, _ := url.Parse(upstream.URL) + sub := &captureSubmitter{} + h := newTestHandler(t, u, sub, proxy.Options{ResponseTimeout: 50 * time.Millisecond}) + + srv := httptest.NewServer(h) + resp, err := http.Get(srv.URL + "/v1/test") + if err == nil { + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + } + time.Sleep(50 * time.Millisecond) + if sub.Len() != 0 { + t.Fatal("timed out response must not be committed") + } + close(release) + srv.Close() + upstream.Close() +} + +func TestSSEIdleTimeoutResetsAfterReads(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + f := w.(http.Flusher) + for _, event := range []string{"data: one\n\n", "data: two\n\n", "data: [DONE]\n\n"} { + _, _ = io.WriteString(w, event) + f.Flush() + time.Sleep(30 * time.Millisecond) + } + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + sub := &captureSubmitter{} + h := newTestHandler(t, u, sub, proxy.Options{SSEIdleTimeout: 60 * time.Millisecond}) + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/v1/test") + if err != nil { + t.Fatal(err) + } + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + if sub.Len() != 1 { + t.Fatalf("entries=%d, want completed SSE commit", sub.Len()) + } +} + +func TestIncompleteSSEEventPreventsCommit(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: partial") + })) + 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() != 0 { + t.Fatal("SSE ending mid-event must not be committed as complete") + } +} + +func TestIncompleteSSEEventAfterTerminalEventPreventsCommit(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + w.(http.Flusher).Flush() + _, _ = io.WriteString(w, "data: partial") + })) + 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() != 0 { + t.Fatal("SSE ending mid-event after a terminal event must not be committed") + } +} + +func TestShutdownReturnsWithoutHijackedConnections(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{}) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := h.Shutdown(ctx); err != nil { + t.Fatal(err) + } +} + +func TestTerminalTextInNonSSEBodyDoesNotAffectCommit(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"text":"data: [DONE]"}`) + })) + 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", strings.NewReader(""))) + if sub.Len() != 1 { + t.Fatalf("entries=%d, want 1", sub.Len()) + } +} diff --git a/tests/proxy/sse_capture_test.go b/tests/proxy/sse_capture_test.go new file mode 100644 index 0000000..51d7227 --- /dev/null +++ b/tests/proxy/sse_capture_test.go @@ -0,0 +1,824 @@ +package proxy_test + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "testing" + "time" + + "git.misaka.ren/M1saka/token_thief/config" + "git.misaka.ren/M1saka/token_thief/logger" + "git.misaka.ren/M1saka/token_thief/proxy" +) + +type captureSubmitter struct { + mu sync.Mutex + entries []*logger.LogEntry +} + +func (c *captureSubmitter) Submit(e *logger.LogEntry) { + c.mu.Lock() + defer c.mu.Unlock() + c.entries = append(c.entries, e) +} + +func (c *captureSubmitter) Len() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.entries) +} + +func (c *captureSubmitter) Entry(i int) *logger.LogEntry { + c.mu.Lock() + defer c.mu.Unlock() + return c.entries[i] +} + +// fakeOpenAIStreamUpstream 模拟一个 OpenAI 兼容的 SSE 上游： +// 分 5 次往响应里 write 一行 SSE 数据，每次都 Flush。 +func fakeOpenAIStreamUpstream() *httptest.Server { + chunks := []string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"你"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"好"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n", + "data: [DONE]\n\n", + } + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + for _, ch := range chunks { + _, _ = io.WriteString(w, ch) + flusher.Flush() + time.Sleep(5 * time.Millisecond) + } + })) +} + +func fakeAnthropicStreamUpstream() *httptest.Server { + chunks := []string{ + `event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-1","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n", + `event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n", + `event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"你"}}` + "\n\n", + `event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"好"}}` + "\n\n", + `event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n", + `event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n", + } + return newSSEUpstream(chunks) +} + +func fakeGeminiStreamUpstream() *httptest.Server { + chunks := []string{ + `data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"index":0}]}` + "\n\n", + `data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n", + } + return newSSEUpstream(chunks) +} + +func fakeUnknownStreamUpstream() *httptest.Server { + return newSSEUpstream([]string{ + `event: custom` + "\n" + `data: not-json` + "\n\n", + }) +} + +func fakeOpenAIStreamUpstreamThatStaysOpen(release <-chan struct{}) *httptest.Server { + chunks := []string{ + `data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n", + "data: [DONE]\n\n", + } + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + for _, ch := range chunks { + _, _ = io.WriteString(w, ch) + flusher.Flush() + } + <-release + })) +} + +func newSSEUpstream(chunks []string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + for _, ch := range chunks { + _, _ = io.WriteString(w, ch) + flusher.Flush() + time.Sleep(5 * time.Millisecond) + } + })) +} + +func TestSSEChunkAssembly(t *testing.T) { + upstream := fakeOpenAIStreamUpstream() + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(u, filter, sub, 1024*1024) + + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + // 客户端走原始 TCP，逐字节读 + 打印，验证流式实时到达 + pu, _ := url.Parse(proxySrv.URL) + conn, err := net.DialTimeout("tcp", pu.Host, 3*time.Second) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + fmt.Fprintf(conn, "POST /v1/chat/completions HTTP/1.1\r\nHost: %s\r\nContent-Length: 0\r\n\r\n", pu.Host) + _ = conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, nil) + if err != nil { + t.Fatal(err) + } + clientBody, _ := io.ReadAll(resp.Body) + + // 等待 proxy.ServeHTTP 返回并把 entry 提交 + deadline := time.Now().Add(2 * time.Second) + for sub.Len() == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if sub.Len() == 0 { + t.Fatal("no log entry captured") + } + e := sub.Entry(0) + + t.Logf("\n========== 客户端收到的字节 ==========\n%s", clientBody) + t.Logf("\n========== 数据库 response_body 字段（按字节原样存储）==========\n%s", e.ResponseBody) + t.Logf("\n========== 元数据 ==========") + t.Logf("is_stream = %v", e.IsStream) + t.Logf("status_code = %d", e.StatusCode) + t.Logf("len(body) = %d bytes", len(e.ResponseBody)) + t.Logf("response_truncated = %v", e.ResponseTruncated) + + if !strings.Contains(string(clientBody), `"content":"你"`) || + !strings.Contains(string(clientBody), `"content":"好"`) || + !strings.Contains(string(clientBody), "[DONE]") { + t.Errorf("client response body 缺少预期 chunk 内容") + } + + var captured struct { + Choices []struct { + Message struct { + Role string `json:"role"` + Content string `json:"content"` + } `json:"message"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("stream response_body should be assembled JSON: %v; body=%q", err, e.ResponseBody) + } + if len(captured.Choices) != 1 { + t.Fatalf("assembled JSON choices length=%d, want 1", len(captured.Choices)) + } + if captured.Choices[0].Message.Role != "assistant" { + t.Errorf("assembled role=%q, want assistant", captured.Choices[0].Message.Role) + } + if captured.Choices[0].Message.Content != "你好" { + t.Errorf("assembled content=%q, want 你好", captured.Choices[0].Message.Content) + } + if captured.Choices[0].FinishReason != "stop" { + t.Errorf("assembled finish_reason=%q, want stop", captured.Choices[0].FinishReason) + } + if !e.IsStream { + t.Errorf("is_stream 应为 true") + } +} + +func TestOpenAIStreamPreservesReasoningAndMetadata(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"think "},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n", + `data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"reasoning_content":"hard"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n", + `data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"content":"final"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n", + `data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{},"finish_reason":"stop","native_finish_reason":"stop"}]}` + "\n\n", + "data: [DONE]\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1/chat/completions") + + var captured struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + Index int `json:"index"` + Message struct { + Role string `json:"role"` + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + } `json:"message"` + FinishReason string `json:"finish_reason"` + NativeFinishReason string `json:"native_finish_reason"` + } `json:"choices"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("openai stream response_body should be JSON: %v; body=%q", err, e.ResponseBody) + } + if captured.ID != "resp-1" || captured.Object != "chat.completion" || captured.Created != 1779335544 || captured.Model != "gpt-5.4-mini-2026-03-17" { + t.Fatalf("unexpected metadata: %+v", captured) + } + if len(captured.Choices) != 1 { + t.Fatalf("choices length=%d, want 1", len(captured.Choices)) + } + choice := captured.Choices[0] + if choice.Message.Role != "assistant" || choice.Message.Content != "final" || choice.Message.ReasoningContent != "think hard" { + t.Fatalf("unexpected message: %+v", choice.Message) + } + if choice.FinishReason != "stop" || choice.NativeFinishReason != "stop" { + t.Fatalf("unexpected finish reasons: %+v", choice) + } +} + +func TestOpenAIStreamDoneDoesNotSubmitBeforeUpstreamCloses(t *testing.T) { + release := make(chan struct{}) + upstream := fakeOpenAIStreamUpstreamThatStaysOpen(release) + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(u, filter, sub, 1024*1024) + + proxySrv := httptest.NewServer(h) + released := false + defer func() { + if !released { + close(release) + } + proxySrv.Close() + upstream.Close() + }() + + clientDone := make(chan error, 1) + go func() { + resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`)) + if err != nil { + clientDone <- err + return + } + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + clientDone <- nil + }() + + time.Sleep(100 * time.Millisecond) + if sub.Len() != 0 { + t.Fatal("terminal SSE event must not submit before ReverseProxy returns") + } + close(release) + released = true + + select { + case err := <-clientDone: + if err != nil { + t.Fatal(err) + } + if sub.Len() != 1 { + t.Fatalf("stream should be submitted once, got %d entries", sub.Len()) + } + case <-time.After(2 * time.Second): + t.Fatal("client did not finish after upstream closed") + } +} + +func TestOpenAIStreamPreservesUsageChunk(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}` + "\n\n", + `data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}` + "\n\n", + "data: [DONE]\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1/chat/completions") + + var captured struct { + ID string `json:"id"` + SystemFingerprint string `json:"system_fingerprint"` + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("openai stream with usage should be assembled JSON: %v; body=%q", err, e.ResponseBody) + } + if captured.ID != "resp-usage" || captured.SystemFingerprint != "fp_123" { + t.Fatalf("metadata not preserved: %+v", captured) + } + if len(captured.Choices) != 1 || captured.Choices[0].Message.Content != "ok" { + t.Fatalf("choices not assembled: %+v", captured.Choices) + } + if captured.Usage.PromptTokens != 5 || captured.Usage.CompletionTokens != 2 || captured.Usage.TotalTokens != 7 { + t.Fatalf("usage not preserved: %+v", captured.Usage) + } +} + +func TestOpenAIStreamAssemblesToolCalls(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{\"q\":"}}]},"finish_reason":null}]}` + "\n\n", + `data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"weather\"}"}}]},"finish_reason":null}]}` + "\n\n", + `data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}` + "\n\n", + "data: [DONE]\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1/chat/completions") + + var captured struct { + Choices []struct { + Message struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []struct { + ID string `json:"id"` + Type string `json:"type"` + Function struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"message"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("openai tool stream should be assembled JSON: %v; body=%q", err, e.ResponseBody) + } + if len(captured.Choices) != 1 || captured.Choices[0].FinishReason != "tool_calls" { + t.Fatalf("unexpected choices: %+v", captured.Choices) + } + message := captured.Choices[0].Message + if message.Role != "assistant" || message.Content != "" { + t.Fatalf("unexpected message basics: %+v", message) + } + if len(message.ToolCalls) != 1 { + t.Fatalf("tool_calls length=%d, want 1; body=%s", len(message.ToolCalls), e.ResponseBody) + } + tool := message.ToolCalls[0] + if tool.ID != "call_1" || tool.Type != "function" || tool.Function.Name != "lookup" || tool.Function.Arguments != `{"q":"weather"}` { + t.Fatalf("unexpected tool call: %+v", tool) + } +} + +func TestOpenAIResponsesStreamAssemblesCompletedResponse(t *testing.T) { + upstream := newSSEUpstream([]string{ + `event: response.created` + "\n" + `data: {"type":"response.created","response":{"id":"resp-1","object":"response","status":"in_progress","model":"gpt-5.4-mini","output":[]}}` + "\n\n", + `event: response.output_text.delta` + "\n" + `data: {"type":"response.output_text.delta","item_id":"msg-1","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n", + `event: response.completed` + "\n" + `data: {"type":"response.completed","response":{"id":"resp-1","object":"response","status":"completed","model":"gpt-5.4-mini","output":[{"id":"msg-1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}` + "\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1/responses") + + var captured struct { + ID string `json:"id"` + Object string `json:"object"` + Status string `json:"status"` + Model string `json:"model"` + Output []struct { + Type string `json:"type"` + Role string `json:"role"` + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + } `json:"output"` + Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("responses stream should be completed response JSON: %v; body=%q", err, e.ResponseBody) + } + if captured.ID != "resp-1" || captured.Object != "response" || captured.Status != "completed" || captured.Model != "gpt-5.4-mini" { + t.Fatalf("unexpected response metadata: %+v", captured) + } + if len(captured.Output) != 1 || len(captured.Output[0].Content) != 1 || captured.Output[0].Content[0].Text != "hello" { + t.Fatalf("unexpected response output: %+v", captured.Output) + } + if captured.Usage.TotalTokens != 4 { + t.Fatalf("usage not preserved: %+v", captured.Usage) + } +} + +func TestOpenAICompletionsStreamAssemblesText(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"hello","finish_reason":null}]}` + "\n\n", + `data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":" world","finish_reason":null}]}` + "\n\n", + `data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"","finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":2,"total_tokens":4}}` + "\n\n", + "data: [DONE]\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1/completions") + + var captured struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + Index int `json:"index"` + Text string `json:"text"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + Usage struct { + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("completions stream should be assembled JSON: %v; body=%q", err, e.ResponseBody) + } + if captured.ID != "cmpl-1" || captured.Object != "text_completion" || captured.Model != "gpt-5.4-mini" || captured.Created != 1779335544 { + t.Fatalf("unexpected metadata: %+v", captured) + } + if len(captured.Choices) != 1 || captured.Choices[0].Text != "hello world" || captured.Choices[0].FinishReason != "stop" { + t.Fatalf("unexpected choices: %+v", captured.Choices) + } + if captured.Usage.TotalTokens != 4 { + t.Fatalf("usage not preserved: %+v", captured.Usage) + } +} + +func TestAnthropicSSEAssemblesNativeMessageJSON(t *testing.T) { + e := requestStreamEntry(t, fakeAnthropicStreamUpstream(), "/v1/messages") + + var captured struct { + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role"` + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + StopReason string `json:"stop_reason"` + Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } `json:"usage"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("anthropic stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody) + } + if captured.ID != "msg-1" || captured.Type != "message" || captured.Role != "assistant" { + t.Fatalf("unexpected anthropic message metadata: %+v", captured) + } + if len(captured.Content) != 1 || captured.Content[0].Type != "text" || captured.Content[0].Text != "你好" { + t.Fatalf("unexpected anthropic content: %+v", captured.Content) + } + if captured.StopReason != "end_turn" { + t.Errorf("stop_reason=%q, want end_turn", captured.StopReason) + } + if captured.Usage.InputTokens != 10 || captured.Usage.OutputTokens != 3 { + t.Errorf("usage=%+v, want input=10 output=3", captured.Usage) + } +} + +func TestAnthropicSSEAssemblesToolUseContent(t *testing.T) { + upstream := newSSEUpstream([]string{ + `event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-tool","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n", + `event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"lookup","input":{}}}` + "\n\n", + `event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}` + "\n\n", + `event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"weather\"}"}}` + "\n\n", + `event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":8}}` + "\n\n", + `event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1/messages") + + var captured struct { + Content []struct { + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + Input map[string]any `json:"input"` + } `json:"content"` + StopReason string `json:"stop_reason"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("anthropic tool stream should be JSON: %v; body=%q", err, e.ResponseBody) + } + if len(captured.Content) != 1 { + t.Fatalf("content length=%d, want 1; body=%s", len(captured.Content), e.ResponseBody) + } + tool := captured.Content[0] + if tool.Type != "tool_use" || tool.ID != "toolu_1" || tool.Name != "lookup" || tool.Input["q"] != "weather" { + t.Fatalf("unexpected tool content: %+v", tool) + } + if captured.StopReason != "tool_use" { + t.Fatalf("stop_reason=%q, want tool_use", captured.StopReason) + } +} + +func TestGeminiSSEAssemblesNativeGenerateContentJSON(t *testing.T) { + e := requestStreamEntry(t, fakeGeminiStreamUpstream(), "/v1beta/models/gemini-1.5-pro:generateContent") + + var captured struct { + Candidates []struct { + Content struct { + Role string `json:"role"` + Parts []struct { + Text string `json:"text"` + } `json:"parts"` + } `json:"content"` + FinishReason string `json:"finishReason"` + Index int `json:"index"` + } `json:"candidates"` + UsageMetadata struct { + PromptTokenCount int `json:"promptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + TotalTokenCount int `json:"totalTokenCount"` + } `json:"usageMetadata"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("gemini stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody) + } + if len(captured.Candidates) != 1 { + t.Fatalf("candidates length=%d, want 1", len(captured.Candidates)) + } + candidate := captured.Candidates[0] + if candidate.Content.Role != "model" || len(candidate.Content.Parts) != 1 || candidate.Content.Parts[0].Text != "你好" { + t.Fatalf("unexpected gemini content: %+v", candidate.Content) + } + if candidate.FinishReason != "STOP" { + t.Errorf("finishReason=%q, want STOP", candidate.FinishReason) + } + if captured.UsageMetadata.TotalTokenCount != 4 { + t.Errorf("usageMetadata=%+v, want totalTokenCount=4", captured.UsageMetadata) + } +} + +func TestGeminiStreamPreservesSafetyRatingsAndUsageMetadata(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"finishReason":null,"index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":8,"thoughtsTokenCount":10}}` + "\n\n", + `data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":2,"totalTokenCount":20,"thoughtsTokenCount":10}}` + "\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent") + + var captured struct { + Candidates []struct { + Content struct { + Parts []struct { + Text string `json:"text"` + } `json:"parts"` + } `json:"content"` + SafetyRatings []struct { + Category string `json:"category"` + Probability string `json:"probability"` + } `json:"safetyRatings"` + } `json:"candidates"` + UsageMetadata struct { + PromptTokenCount int `json:"promptTokenCount"` + ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + TotalTokenCount int `json:"totalTokenCount"` + ThoughtsTokenCount int `json:"thoughtsTokenCount"` + } `json:"usageMetadata"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("gemini stream response_body should be JSON: %v; body=%q", err, e.ResponseBody) + } + if len(captured.Candidates) != 1 || len(captured.Candidates[0].SafetyRatings) != 1 { + t.Fatalf("expected safetyRatings to be preserved, got %+v", captured.Candidates) + } + if captured.Candidates[0].SafetyRatings[0].Category != "HARM_CATEGORY_HARASSMENT" { + t.Fatalf("unexpected safetyRatings: %+v", captured.Candidates[0].SafetyRatings) + } + if captured.UsageMetadata.ToolUsePromptTokenCount != 0 || captured.UsageMetadata.ThoughtsTokenCount != 10 || captured.UsageMetadata.TotalTokenCount != 20 { + t.Fatalf("usageMetadata fields not preserved: %+v", captured.UsageMetadata) + } +} + +func TestGeminiStreamPreservesFunctionCallParts(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n", + }) + e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent") + + var captured struct { + Candidates []struct { + Content struct { + Parts []struct { + FunctionCall struct { + Name string `json:"name"` + Args map[string]any `json:"args"` + } `json:"functionCall"` + } `json:"parts"` + } `json:"content"` + } `json:"candidates"` + } + if err := json.Unmarshal(e.ResponseBody, &captured); err != nil { + t.Fatalf("gemini functionCall stream should be JSON: %v; body=%q", err, e.ResponseBody) + } + call := captured.Candidates[0].Content.Parts[0].FunctionCall + if call.Name != "lookup" || call.Args["q"] != "weather" { + t.Fatalf("functionCall not preserved: %+v; body=%s", call, e.ResponseBody) + } +} + +func TestUnknownSSEKeepsRawBody(t *testing.T) { + e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions") + + if string(e.ResponseBody) != "event: custom\ndata: not-json\n\n" { + t.Fatalf("unknown stream should keep raw body, got %q", e.ResponseBody) + } +} + +func TestTruncatedSSEKeepsCapturedRawBody(t *testing.T) { + upstream := fakeOpenAIStreamUpstream() + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(u, filter, sub, 520) + + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", nil) + if err != nil { + t.Fatal(err) + } + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + + deadline := time.Now().Add(2 * time.Second) + for sub.Len() == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if sub.Len() == 0 { + t.Fatal("no log entry captured") + } + e := sub.Entry(0) + + if !e.ResponseTruncated { + t.Fatal("response should be marked truncated") + } + if !strings.HasPrefix(string(e.ResponseBody), "data: ") { + t.Fatalf("truncated stream should keep captured raw body, got %q", e.ResponseBody) + } + if json.Valid(e.ResponseBody) { + t.Fatalf("truncated stream should not be assembled as JSON, got %q", e.ResponseBody) + } +} + +func TestMultimodalStreamRequestBodyIsCaptured(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n", + `data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n", + "data: [DONE]\n\n", + }) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(u, filter, sub, 1024*1024) + + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}` + resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(reqBody)) + if err != nil { + t.Fatal(err) + } + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + + deadline := time.Now().Add(2 * time.Second) + for sub.Len() == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if sub.Len() == 0 { + t.Fatal("no log entry captured") + } + e := sub.Entry(0) + + if e.RequestTruncated { + t.Fatal("multimodal request should not be truncated") + } + requestText := string(e.RequestBody) + if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) { + t.Fatalf("request_body should contain image input, got %q", requestText) + } + if !e.IsStream { + t.Fatal("response should still be marked stream") + } +} + +func TestChunkedMultimodalRequestBodyIsCaptured(t *testing.T) { + upstream := newSSEUpstream([]string{ + `data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}` + "\n\n", + "data: [DONE]\n\n", + }) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(u, filter, sub, 1024*1024) + + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}` + req, err := http.NewRequest(http.MethodPost, proxySrv.URL+"/v1/chat/completions", strings.NewReader(reqBody)) + if err != nil { + t.Fatal(err) + } + req.ContentLength = -1 + req.Header.Set("Content-Type", "application/json") + req.TransferEncoding = []string{"chunked"} + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + + deadline := time.Now().Add(2 * time.Second) + for sub.Len() == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if sub.Len() == 0 { + t.Fatal("no log entry captured") + } + e := sub.Entry(0) + + requestText := string(e.RequestBody) + if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) { + t.Fatalf("chunked request_body should contain image input, got %q", requestText) + } +} + +func requestStreamEntry(t *testing.T, upstream *httptest.Server, path string) *logger.LogEntry { + t.Helper() + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + h := proxy.New(u, filter, sub, 1024*1024) + + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + resp, err := http.Post(proxySrv.URL+path, "application/json", nil) + if err != nil { + t.Fatal(err) + } + _, _ = io.ReadAll(resp.Body) + _ = resp.Body.Close() + + deadline := time.Now().Add(2 * time.Second) + for sub.Len() == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if sub.Len() == 0 { + t.Fatal("no log entry captured") + } + return sub.Entry(0) +} diff --git a/tests/proxy/websocket_test.go b/tests/proxy/websocket_test.go new file mode 100644 index 0000000..e780e03 --- /dev/null +++ b/tests/proxy/websocket_test.go @@ -0,0 +1,299 @@ +package proxy_test + +import ( + "bufio" + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "git.misaka.ren/M1saka/token_thief/config" + "git.misaka.ren/M1saka/token_thief/logger" + "git.misaka.ren/M1saka/token_thief/proxy" +) + +type noopSubmitter struct{} + +func (noopSubmitter) Submit(*logger.LogEntry) {} + +type chanSubmitter chan *logger.LogEntry + +func (c chanSubmitter) Submit(e *logger.LogEntry) { c <- e } + +// fakeUpstream 模拟一个最简 WebSocket 升级： +// 收到 GET + Upgrade: websocket 后回 101，然后做字节回声直到对端关闭。 +func fakeUpstream(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.ToLower(r.Header.Get("Upgrade")) != "websocket" { + http.Error(w, "expected websocket upgrade", http.StatusBadRequest) + return + } + hj, ok := w.(http.Hijacker) + if !ok { + http.Error(w, "no hijack", http.StatusInternalServerError) + return + } + conn, brw, err := hj.Hijack() + if err != nil { + t.Errorf("upstream hijack: %v", err) + return + } + defer conn.Close() + // 直接回 101 握手响应（简化版，不做真正 Sec-WebSocket-Accept 计算） + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\n" + + "Upgrade: websocket\r\n" + + "Connection: Upgrade\r\n" + + "\r\n") + _ = brw.Flush() + // echo + buf := make([]byte, 1024) + for { + n, err := conn.Read(buf) + if err != nil { + return + } + if _, err := conn.Write(buf[:n]); err != nil { + return + } + } + })) + return srv +} + +func TestWebSocketHandshakeCapturedAndShutdownClosesConnection(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) + } + entries := make(chanSubmitter, 1) + h := proxy.New(u, filter, entries, 1024) + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + pu, _ := url.Parse(proxySrv.URL) + conn, err := net.DialTimeout("tcp", pu.Host, time.Second) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + _, _ = io.WriteString(conn, "GET /v1/realtime HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n") + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, nil) + if err != nil { + t.Fatal(err) + } + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Fatalf("status=%d", resp.StatusCode) + } + select { + case entry := <-entries: + if entry.StatusCode != http.StatusSwitchingProtocols || len(entry.ResponseBody) != 0 { + t.Fatalf("handshake entry status=%d body=%q", entry.StatusCode, entry.ResponseBody) + } + case <-time.After(time.Second): + t.Fatal("websocket handshake was not captured") + } + + payload := []byte("frame-data-must-not-be-logged") + if _, err := conn.Write(payload); err != nil { + t.Fatal(err) + } + echo := make([]byte, len(payload)) + if _, err := io.ReadFull(br, echo); err != nil { + t.Fatalf("read websocket payload: %v", err) + } + if string(echo) != string(payload) { + t.Fatalf("echo mismatch: got %q want %q", echo, payload) + } + select { + case entry := <-entries: + t.Fatalf("websocket frame produced an extra log entry: %+v", entry) + case <-time.After(50 * time.Millisecond): + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := h.Shutdown(ctx); err != nil { + t.Fatal(err) + } + _ = conn.SetReadDeadline(time.Now().Add(time.Second)) + if _, err := conn.Read(make([]byte, 1)); err == nil { + t.Fatal("connection remains open after Shutdown") + } + if err := h.Shutdown(ctx); err != nil { + t.Fatalf("second Shutdown: %v", err) + } +} + +func TestWebSocketNaturalCloseUnregistersConnection(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) + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + pu, _ := url.Parse(proxySrv.URL) + conn, err := net.DialTimeout("tcp", pu.Host, time.Second) + if err != nil { + t.Fatal(err) + } + _, _ = io.WriteString(conn, "GET /v1/realtime HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n") + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + t.Fatal(err) + } + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Fatalf("status=%d", resp.StatusCode) + } + if err := conn.Close(); err != nil { + t.Fatal(err) + } + + deadline := time.Now().Add(time.Second) + for { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + err := h.Shutdown(ctx) + if err == nil { + break + } + if time.Now().After(deadline) { + t.Fatalf("connection was not unregistered after natural close: %v", err) + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestNonWebSocketUpgradeIsManagedButNotLogged(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hj := w.(http.Hijacker) + conn, brw, err := hj.Hijack() + if err != nil { + t.Errorf("upstream hijack: %v", err) + return + } + defer conn.Close() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: test-protocol\r\nConnection: Upgrade\r\n\r\n") + _ = brw.Flush() + _, _ = io.Copy(io.Discard, conn) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + entries := make(chanSubmitter, 1) + h := proxy.New(u, filter, entries, 1024) + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + pu, _ := url.Parse(proxySrv.URL) + conn, err := net.DialTimeout("tcp", pu.Host, time.Second) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + _, _ = io.WriteString(conn, "GET /upgrade HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: test-protocol\r\nConnection: Upgrade\r\n\r\n") + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + t.Fatal(err) + } + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Fatalf("status=%d", resp.StatusCode) + } + select { + case entry := <-entries: + t.Fatalf("non-WebSocket upgrade was logged: %+v", entry) + case <-time.After(50 * time.Millisecond): + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := h.Shutdown(ctx); err != nil { + t.Fatal(err) + } +} + +// TestWebSocketProxyPassthrough 验证 Upgrade 请求能正确透传， +// 确认 captureWriter 的 Hijacker 实现没破坏 ReverseProxy 的 WS 行为。 +func TestWebSocketProxyPassthrough(t *testing.T) { + upstream := fakeUpstream(t) + defer upstream.Close() + + u, err := url.Parse(upstream.URL) + if err != nil { + t.Fatal(err) + } + + // disabled 模式下 ShouldLog 返回 true（全量记录），会进入捕获分支。 + filter, err := config.NewFilter(config.FilterDisabled, nil) + if err != nil { + t.Fatal(err) + } + + h := proxy.New(u, filter, noopSubmitter{}, 1024) + + proxySrv := httptest.NewServer(h) + defer proxySrv.Close() + + // 建立到代理的 TCP 连接，手写 Upgrade 请求 + pu, _ := url.Parse(proxySrv.URL) + d := net.Dialer{Timeout: 3 * time.Second} + conn, err := d.DialContext(context.Background(), "tcp", pu.Host) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + req := "GET /v1/realtime HTTP/1.1\r\n" + + "Host: " + pu.Host + "\r\n" + + "Upgrade: websocket\r\n" + + "Connection: Upgrade\r\n" + + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" + + "Sec-WebSocket-Version: 13\r\n" + + "\r\n" + if _, err := io.WriteString(conn, req); err != nil { + t.Fatal(err) + } + + _ = conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, nil) + if err != nil { + t.Fatalf("read upgrade response: %v", err) + } + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Fatalf("expected 101, got %d", resp.StatusCode) + } + if !strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") { + t.Fatalf("expected Upgrade: websocket, got %q", resp.Header.Get("Upgrade")) + } + + // echo 测试 + payload := "hello-websocket" + if _, err := io.WriteString(conn, payload); err != nil { + t.Fatal(err) + } + got := make([]byte, len(payload)) + if _, err := io.ReadFull(br, got); err != nil { + t.Fatalf("read echo: %v", err) + } + if string(got) != payload { + t.Fatalf("echo mismatch: got %q want %q", got, payload) + } +} diff --git a/tests/scripts/dbcheck/main.go b/tests/scripts/dbcheck/main.go new file mode 100644 index 0000000..6dca2ea --- /dev/null +++ b/tests/scripts/dbcheck/main.go @@ -0,0 +1,109 @@ +// Standalone helper for smoke.ps1. It reads CHECK_RIDS (JSON array) and verifies rows in proxy_logs. +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "os" + "time" + + "github.com/ClickHouse/clickhouse-go/v2" + + "git.misaka.ren/M1saka/token_thief/db" +) + +type row struct { + RequestID string + Method string + Path string + StatusCode int + RequestTruncated bool + ResponseTruncated bool + IsStream bool + LatencyMS int64 + ReqBodyLen int64 + RespBodyLen int64 + ReqHeaders string + ErrorMsg string +} + +func main() { + log.SetFlags(0) + + dsn := os.Getenv("CLICKHOUSE_URL") + if dsn == "" { + log.Fatalf("CLICKHOUSE_URL not set") + } + ridsRaw := os.Getenv("CHECK_RIDS") + if ridsRaw == "" { + log.Fatalf("CHECK_RIDS not set") + } + var rids []string + if err := json.Unmarshal([]byte(ridsRaw), &rids); err != nil { + log.Fatalf("parse CHECK_RIDS: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + opts, err := db.ClickHouseOptions(dsn) + if err != nil { + log.Fatalf("connect: %v", err) + } + conn, err := clickhouse.Open(opts) + if err != nil { + log.Fatalf("connect: %v", err) + } + defer conn.Close() + + const q = ` +SELECT request_id, method, path, status_code, + request_truncated, response_truncated, is_stream, latency_ms, + toInt64(length(request_body)), + toInt64(length(response_body)), + request_headers, + error +FROM proxy_logs WHERE request_id = ? +ORDER BY started_at DESC LIMIT 1` + + ok := 0 + for _, rid := range rids { + var r row + err := conn.QueryRow(ctx, q, rid).Scan( + &r.RequestID, &r.Method, &r.Path, &r.StatusCode, + &r.RequestTruncated, &r.ResponseTruncated, &r.IsStream, &r.LatencyMS, + &r.ReqBodyLen, &r.RespBodyLen, &r.ReqHeaders, &r.ErrorMsg, + ) + if err != nil { + fmt.Printf(" FAIL rid=%s: not found in db (%v)\n", rid, err) + continue + } + ok++ + fmt.Printf(" OK rid=%s\n", rid) + fmt.Printf(" method=%s path=%s status=%d latency_ms=%d\n", r.Method, r.Path, r.StatusCode, r.LatencyMS) + fmt.Printf(" req_body=%d B (truncated=%v) resp_body=%d B (truncated=%v) is_stream=%v\n", + r.ReqBodyLen, r.RequestTruncated, r.RespBodyLen, r.ResponseTruncated, r.IsStream) + hasAuth := false + var hm map[string][]string + if err := json.Unmarshal([]byte(r.ReqHeaders), &hm); err == nil { + _, hasAuth = hm["Authorization"] + } + fmt.Printf(" headers has Authorization=%v\n", hasAuth) + if r.ErrorMsg != "" { + fmt.Printf(" error=%s\n", trunc(r.ErrorMsg, 200)) + } + } + fmt.Printf("\n found %d / %d\n", ok, len(rids)) + if ok != len(rids) { + os.Exit(1) + } +} + +func trunc(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "..." +} diff --git a/tests/scripts/dump-stream-body.ps1 b/tests/scripts/dump-stream-body.ps1 new file mode 100644 index 0000000..d28882a --- /dev/null +++ b/tests/scripts/dump-stream-body.ps1 @@ -0,0 +1,54 @@ +param([string]$RequestID) + +if (-not $RequestID) { throw "RequestID is required" } + +. .\tests\scripts\load-env.ps1 | Out-Null + +$env:DUMP_RID = $RequestID +@' +package main + +import ( + "context" + "fmt" + "log" + "os" + "time" + + "github.com/ClickHouse/clickhouse-go/v2" + + "git.misaka.ren/M1saka/token_thief/db" +) + +func main() { + log.SetFlags(0) + dsn := os.Getenv("CLICKHOUSE_URL") + if dsn == "" { + log.Fatal("CLICKHOUSE_URL not set") + } + rid := os.Getenv("DUMP_RID") + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + opts, err := db.ClickHouseOptions(dsn) + if err != nil { + log.Fatal(err) + } + conn, err := clickhouse.Open(opts) + if err != nil { + log.Fatal(err) + } + defer conn.Close() + var body string + if err := conn.QueryRow(ctx, + `SELECT response_body FROM proxy_logs WHERE request_id = ? ORDER BY started_at DESC LIMIT 1`, + rid, + ).Scan(&body); err != nil { + log.Fatal(err) + } + fmt.Println(body) +} +'@ | Set-Content -Path tmp_dump.go -Encoding UTF8 -NoNewline + +go run tmp_dump.go +Remove-Item tmp_dump.go -Force +Remove-Item Env:DUMP_RID -ErrorAction SilentlyContinue diff --git a/tests/scripts/load-env.ps1 b/tests/scripts/load-env.ps1 new file mode 100644 index 0000000..84ea785 --- /dev/null +++ b/tests/scripts/load-env.ps1 @@ -0,0 +1,25 @@ +# Load KEY=VALUE pairs from .env into the current PowerShell process. +# Usage: . .\tests\scripts\load-env.ps1 +param( + [string]$Path = ".env" +) + +if (-not (Test-Path $Path)) { + Write-Error "env file not found: $Path" + return +} + +Get-Content $Path | ForEach-Object { + $line = $_.Trim() + if ($line -eq "" -or $line.StartsWith("#")) { return } + $idx = $line.IndexOf("=") + if ($idx -lt 1) { return } + $key = $line.Substring(0, $idx).Trim() + $val = $line.Substring($idx + 1).Trim() + if (($val.StartsWith('"') -and $val.EndsWith('"')) -or + ($val.StartsWith("'") -and $val.EndsWith("'"))) { + $val = $val.Substring(1, $val.Length - 2) + } + [Environment]::SetEnvironmentVariable($key, $val, "Process") + Write-Host " loaded $key" +} diff --git a/tests/scripts/smoke.ps1 b/tests/scripts/smoke.ps1 new file mode 100644 index 0000000..aa3bc7c --- /dev/null +++ b/tests/scripts/smoke.ps1 @@ -0,0 +1,135 @@ +# End-to-end smoke test: +# start proxy -> run chat cases -> wait batch flush -> verify ClickHouse -> stop proxy. +param( + [string]$Model = "gpt-5.4-mini", + [string]$ApiKey = $env:NEWAPI_KEY, + [string]$ProxyBase = "http://127.0.0.1:8080" +) + +if (-not $ApiKey) { throw "NEWAPI_KEY not set" } + +$ErrorActionPreference = "Stop" + +function Section($name) { + Write-Host "" + Write-Host ("=" * 70) -ForegroundColor Cyan + Write-Host $name -ForegroundColor Cyan + Write-Host ("=" * 70) -ForegroundColor Cyan +} + +Section "Start proxy" +if (Test-Path proxy.log) { Remove-Item proxy.log -Force } +if (Test-Path proxy.err.log) { Remove-Item proxy.err.log -Force } +$proxy = Start-Process -FilePath .\TokenThief.exe -PassThru -RedirectStandardOutput proxy.log -RedirectStandardError proxy.err.log -WindowStyle Hidden +Write-Host "proxy pid=$($proxy.Id)" +Start-Sleep -Seconds 2 +$health = & curl.exe -s -o NUL -w "%{http_code}" "$ProxyBase/healthz" +Write-Host "healthz: $health" +if ($health -ne "200") { + Get-Content proxy.log -Tail 30 | Write-Host + Get-Content proxy.err.log -Tail 30 | Write-Host + throw "proxy did not start" +} + +$results = @{} + +function CallChat { + param([string]$Url, [string]$JsonBody) + + $bodyTmp = New-TemporaryFile + $headTmp = New-TemporaryFile + $outTmp = New-TemporaryFile + # PowerShell 5.1 Set-Content -Encoding UTF8 writes a BOM; newapi rejects BOM-prefixed JSON. + $utf8NoBom = New-Object System.Text.UTF8Encoding($false) + [System.IO.File]::WriteAllText($bodyTmp.FullName, $JsonBody, $utf8NoBom) + + $code = & curl.exe -s -X POST $Url ` + -H "Content-Type: application/json" ` + -H "Authorization: Bearer $ApiKey" ` + -D $headTmp.FullName ` + -o $outTmp.FullName ` + --data-binary "@$($bodyTmp.FullName)" ` + -w "%{http_code}" + + $rid = "" + foreach ($line in Get-Content $headTmp.FullName) { + if ($line -match '^X-Request-Id:\s*(.+)$') { + $rid = $matches[1].Trim() + break + } + } + $body = Get-Content $outTmp.FullName -Raw -ErrorAction SilentlyContinue + if (-not $body) { $body = "" } + + Remove-Item $bodyTmp.FullName, $headTmp.FullName, $outTmp.FullName -Force -ErrorAction SilentlyContinue + + return @{ StatusCode = $code; RequestID = $rid; Body = $body } +} + +try { + Section "T1: non-stream chat completions" + $t1Body = '{"model":"' + $Model + '","messages":[{"role":"user","content":"Say hello in one short sentence."}],"stream":false}' + $r = CallChat "$ProxyBase/v1/chat/completions" $t1Body + Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)" + $preview = if ($r.Body.Length -gt 200) { $r.Body.Substring(0, 200) } else { $r.Body } + Write-Host " body(first 200): $preview" + $results.T1 = $r + + Section "T2: stream chat completions" + $t2Body = '{"model":"' + $Model + '","messages":[{"role":"user","content":"Say exactly: one two three four five"}],"stream":true}' + $r = CallChat "$ProxyBase/v1/chat/completions" $t2Body + Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)" + $chunkCount = ([regex]::Matches($r.Body, "^data:", "Multiline")).Count + $hasDone = $r.Body.Contains("[DONE]") + Write-Host " SSE chunk lines=$chunkCount contains [DONE]=$hasDone body_len=$($r.Body.Length)" + $results.T2 = $r + + Section "T3: large body (request_truncated should be true)" + $bigContent = "x" * (12 * 1024 * 1024) + $t3Body = '{"model":"' + $Model + '","messages":[{"role":"user","content":"' + $bigContent + '"}],"stream":false,"max_tokens":5}' + Write-Host " request body size: $($t3Body.Length) bytes" + $r = CallChat "$ProxyBase/v1/chat/completions" $t3Body + Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)" + $preview = if ($r.Body.Length -gt 200) { $r.Body.Substring(0, 200) } else { $r.Body } + Write-Host " body(first 200): $preview" + $results.T3 = $r + + Section "T4: nonexistent model (upstream error response should be logged)" + $t4Body = '{"model":"definitely-not-a-real-model-xyz","messages":[{"role":"user","content":"hi"}]}' + $r = CallChat "$ProxyBase/v1/chat/completions" $t4Body + Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)" + $preview = if ($r.Body.Length -gt 200) { $r.Body.Substring(0, 200) } else { $r.Body } + Write-Host " body(first 200): $preview" + $results.T4 = $r + + Section "T5: healthz (should not be logged)" + $code = & curl.exe -s -o NUL -w "%{http_code}" "$ProxyBase/healthz" + Write-Host " /healthz status=$code" + + Section "Wait batch flush (4s)" + Start-Sleep -Seconds 4 + + Section "T6: ClickHouse verification" + $rids = @() + foreach ($k in @("T1","T2","T3","T4")) { + if ($results[$k].RequestID) { $rids += $results[$k].RequestID } + } + Write-Host " request_ids: $($rids -join ', ')" + $env:CHECK_RIDS = ($rids | ConvertTo-Json -Compress) + if ($rids.Count -eq 1) { $env:CHECK_RIDS = "[`"$($rids[0])`"]" } + & go run .\tests\scripts\dbcheck\main.go + Remove-Item Env:CHECK_RIDS -ErrorAction SilentlyContinue + +} finally { + Section "Stop proxy" + Stop-Process -Id $proxy.Id -Force + Write-Host "tail proxy.log:" + Get-Content proxy.log -Tail 30 | Write-Host + if (Test-Path proxy.err.log) { + $err = Get-Content proxy.err.log -ErrorAction SilentlyContinue + if ($err) { + Write-Host "proxy.err.log:" + $err | Write-Host + } + } +}