Sakurasan
ad10bdc400
链路功能已经跑通，缺的是边界。这几处默认值都是仓库里公开已知的，忘配就等于没有认证。
- 认证端点限速：/auth/login、/auth/register 此前完全没接限速器（Limiter 只挂在代理
Auth 上），爆破与批量建号零成本。新增 AllowWindow 滑动窗口，登录按「单 IP」和
「单 IP×账号」两个桶（默认 30/min、10/min），注册按 IP（5/min）且在 argon2 之前拦，
避免光靠建号请求打满 CPU；超限返回 429 + Retry-After。
- 生产启动断言：env=production 时拒绝默认 OT_JWT_SECRET、缺失或短于 16 字节的
OT_MASTER_KEY（缺失时 NewEncryptor 会用 sha256("") 派生，渠道密文可被公开还原）、
默认管理员密码；一次列全所有问题，不再改一个试一个。development 不拦。
- 请求体上限：parseBody 原来全量 io.ReadAll 且全项目没有一处 MaxBytesReader，
合法密钥可把任意大 body 整份压进内存再转发。/v1 组挂 BodyLimit（默认 32MB，
OT_PROXY_MAX_BODY_MB），超限按客户端协议返回 413 request_too_large 并点名上限。
- CORS：不再无条件反射 Origin + Allow-Credentials: true——那唯一挡住跨站凭据读取的
是 refresh cookie 的 SameSite=Lax，改成 None 就立刻变成窃取面。改为显式白名单，
空白名单完全不发跨域头（同源部署与 vite proxy 本来就不需要），通配 * 不配凭据；
允许头补上 X-Api-Key / Anthropic-Version / OpenAI-*，否则浏览器里跑官方 SDK 预检会挂。
- 配套 OT_SERVER_TRUSTED_PROXIES：不配置就不信 X-Forwarded-For（防伪造头绕过 IP 限速）。
顺手删掉 .env.example 末尾重复的 OT_PROXY_LOG_RAW 段。
336 lines
11 KiB
Go
336 lines
11 KiB
Go
package proxy
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/openteam/server/internal/pkg/apikey"
|
|
"github.com/openteam/server/internal/pkg/crypto"
|
|
"github.com/openteam/server/internal/pkg/ratelimit"
|
|
"github.com/openteam/server/internal/store"
|
|
"github.com/openteam/server/internal/usage"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type harness struct {
|
|
gw *Gateway
|
|
db *gorm.DB
|
|
rec *usage.Recorder
|
|
key string
|
|
}
|
|
|
|
// newHarness 起一个内存 SQLite 网关：1 用户 + 1 密钥 + 1 模型 + 1 渠道（可选绑定）。
|
|
func newHarness(t *testing.T, baseURL string, timeoutMS int, bind bool, keyModels []string) *harness {
|
|
t.Helper()
|
|
dsn := "/tmp/proxy_fix_test_" + time.Now().Format("150405.000000000") + ".db"
|
|
db, err := store.Open("sqlite", dsn)
|
|
if err != nil {
|
|
t.Fatalf("open db: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
sqlDB, _ := db.DB()
|
|
sqlDB.Close()
|
|
os.Remove(dsn)
|
|
})
|
|
|
|
enc := crypto.NewEncryptor("test-master-key-1234567890")
|
|
encKey, _ := enc.Encrypt("upstream-key")
|
|
u := store.User{Username: "u1", Email: "u1@x.com", PasswordHash: "x", Role: store.RoleUser,
|
|
Balance: 100, Status: store.UserStatusActive}
|
|
if err := db.Create(&u).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
plain, hash, prefix, _ := apikey.Generate()
|
|
k := store.APIKey{UserID: u.ID, Name: "t", KeyHash: hash, KeyPrefix: prefix,
|
|
Status: store.KeyStatusActive, AllowedModels: keyModels}
|
|
if err := db.Create(&k).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
m := store.Model{Name: "m1", InputPrice: 0.15, OutputPrice: 0.6, Enabled: true}
|
|
if err := db.Create(&m).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ch := store.Channel{Name: "c1", Provider: store.ChannelProviderOpenAI,
|
|
Formats: []string{store.FormatChat}, BaseURL: baseURL, APIKeyEnc: encKey,
|
|
Weight: 1, TimeoutMS: timeoutMS, MaxConcurrency: 4,
|
|
HealthStatus: store.ChannelHealthHealthy, Enabled: true}
|
|
if err := db.Create(&ch).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if bind {
|
|
if err := db.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
rec := usage.NewRecorder(db)
|
|
t.Cleanup(rec.Close)
|
|
return &harness{
|
|
gw: NewGateway(db, enc, rec, ratelimit.New(), 0, false, 300*time.Second, 32),
|
|
db: db, rec: rec, key: plain,
|
|
}
|
|
}
|
|
|
|
// run 走完整链路：Auth（含新的密钥白名单注入）→ Handle。
|
|
func (h *harness) run(t *testing.T, ctx context.Context, body string) *httptest.ResponseRecorder {
|
|
t.Helper()
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
|
|
req.Header.Set("Authorization", "Bearer "+h.key)
|
|
if ctx != nil {
|
|
req = req.WithContext(ctx)
|
|
}
|
|
c.Request = req
|
|
h.gw.Auth(c)
|
|
if !c.IsAborted() {
|
|
h.gw.Handle(c)
|
|
}
|
|
return w
|
|
}
|
|
|
|
// 修复①：渠道 timeout_ms 真正生效（此前被 int64 溢出成"永不超时"）。
|
|
func TestChannelTimeoutIsEnforced(t *testing.T) {
|
|
var hits atomic.Int64
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
hits.Add(1)
|
|
select {
|
|
case <-r.Context().Done():
|
|
case <-time.After(1 * time.Second):
|
|
}
|
|
_, _ = io.WriteString(w, `{"id":"late"}`)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
h := newHarness(t, srv.URL, 200, true, nil)
|
|
start := time.Now()
|
|
w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
|
|
elapsed := time.Since(start)
|
|
|
|
if elapsed > 600*time.Millisecond {
|
|
t.Fatalf("渠道 timeout_ms=200 没生效（上游 1s 后才返回），实际耗时 %s", elapsed)
|
|
}
|
|
if w.Code != http.StatusBadGateway {
|
|
t.Fatalf("超时后应报 upstream_error，实际 %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if hits.Load() == 0 {
|
|
t.Fatal("上游没被调用")
|
|
}
|
|
}
|
|
|
|
// 修复①延伸：流式不受整段超时限制，但首包超时由 transport 兜住。
|
|
func TestStreamHasNoWholeBodyDeadline(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(200)
|
|
f := w.(http.Flusher)
|
|
for i := 0; i < 3; i++ {
|
|
_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"x"}}]}`+"\n\n")
|
|
f.Flush()
|
|
time.Sleep(150 * time.Millisecond) // 总计 450ms > 渠道 timeout_ms
|
|
}
|
|
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
|
}))
|
|
defer srv.Close()
|
|
|
|
h := newHarness(t, srv.URL, 200, true, nil)
|
|
w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("stream code=%d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if n := strings.Count(w.Body.String(), `"content":"x"`); n != 3 {
|
|
t.Fatalf("流被超时截断，只收到 %d 块，期望 3", n)
|
|
}
|
|
}
|
|
|
|
// 修复②：客户端断连必须把取消传到上游，两侧都停止，而不是让上游继续生成。
|
|
func TestClientDisconnectCancelsUpstream(t *testing.T) {
|
|
started := make(chan struct{})
|
|
upstreamStop := make(chan error, 1)
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f := w.(http.Flusher)
|
|
close(started)
|
|
// 模拟持续生成：客户端断开后写会失败、request context 会被取消
|
|
for i := 0; i < 40; i++ {
|
|
if _, err := io.WriteString(w, `data: {"choices":[{"delta":{"content":"x"}}]}`+"\n\n"); err != nil {
|
|
upstreamStop <- err
|
|
return
|
|
}
|
|
f.Flush()
|
|
if err := r.Context().Err(); err != nil {
|
|
upstreamStop <- err
|
|
return
|
|
}
|
|
time.Sleep(50 * time.Millisecond)
|
|
}
|
|
upstreamStop <- errors.New("上游跑完了 2s 都没察觉客户端断开")
|
|
}))
|
|
defer srv.Close()
|
|
|
|
h := newHarness(t, srv.URL, 0, true, nil)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
done := make(chan *httptest.ResponseRecorder, 1)
|
|
go func() {
|
|
done <- h.run(t, ctx, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
|
}()
|
|
|
|
<-started
|
|
time.Sleep(120 * time.Millisecond) // 先收几块，保证已进入流式转发
|
|
cancel()
|
|
|
|
select {
|
|
case err := <-upstreamStop:
|
|
t.Logf("上游因客户端断开而停止：%v", err)
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("客户端断开后上游仍在生成：context 没有传播")
|
|
}
|
|
<-done
|
|
h.rec.Close()
|
|
|
|
var n int64
|
|
h.db.Model(&store.UsageLog{}).Where("status = ? AND error_code = ?",
|
|
store.UsageStatusCanceled, "client_disconnect").Count(&n)
|
|
if n != 1 {
|
|
t.Fatalf("canceled 记账条数 = %d，期望 1", n)
|
|
}
|
|
}
|
|
|
|
// 修复③：APIKey.AllowedModels 现在真的是道闸门。
|
|
func TestKeyModelWhitelistIsEnforced(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
h := newHarness(t, srv.URL, 0, true, []string{"gpt-4o"})
|
|
w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
|
|
if w.Code != http.StatusForbidden {
|
|
t.Fatalf("密钥白名单外的模型应 403，实际 %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var body map[string]any
|
|
_ = json.Unmarshal(w.Body.Bytes(), &body)
|
|
errObj, _ := body["error"].(map[string]any)
|
|
if errObj == nil || errObj["code"] != "model_not_allowed" {
|
|
t.Fatalf("错误体不符：%v", body)
|
|
}
|
|
}
|
|
|
|
// 修复④（渠道侧）：运行期连续失败让坏渠道从候选里消失。
|
|
func TestRuntimeCooldownEvictsFailingChannel(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"rate limited"}}`)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
h := newHarness(t, srv.URL, 0, true, nil)
|
|
var ch store.Channel
|
|
h.db.First(&ch)
|
|
|
|
// 3 次 429：每次换不到别的渠道，最终 502，同时累计运行期失败
|
|
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
|
for i := 0; i < 3; i++ {
|
|
if got := len(h.gw.candidateChannels("m1")); got != 1 {
|
|
t.Fatalf("第 %d 次失败前候选应为 1，实际 %d", i+1, got)
|
|
}
|
|
if w := h.run(t, nil, body); w.Code != http.StatusTooManyRequests {
|
|
t.Fatalf("429 应作为最后状态返回，实际 %d: %s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
if !h.gw.ch.InCooldown(ch.ID) {
|
|
t.Fatal("3 次连续 429 后渠道应进入运行期冷却")
|
|
}
|
|
if got := h.gw.candidateChannels("m1"); len(got) != 0 {
|
|
t.Fatalf("冷却中不应出现在候选集，实际 %d", len(got))
|
|
}
|
|
if w := h.run(t, nil, body); w.Code != http.StatusServiceUnavailable {
|
|
t.Fatalf("冷却后应 503 no_channel，实际 %d: %s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
|
|
// 修复④延伸：上游恢复正常后清空失败计数，不会带着旧账被误踢。
|
|
func TestRuntimeFailureCounterResetsOnSuccess(t *testing.T) {
|
|
fail := true
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if fail {
|
|
w.WriteHeader(http.StatusBadGateway)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"boom"}}`)
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
h := newHarness(t, srv.URL, 0, true, nil)
|
|
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
|
h.run(t, nil, body)
|
|
h.run(t, nil, body)
|
|
fail = false
|
|
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
|
t.Fatalf("应成功，实际 %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var ch store.Channel
|
|
h.db.First(&ch)
|
|
// 成功清零后再失败两次仍不该进冷却（阈值是"连续"3 次）
|
|
fail = true
|
|
h.run(t, nil, body)
|
|
h.run(t, nil, body)
|
|
if h.gw.ch.InCooldown(ch.ID) {
|
|
t.Fatal("失败计数没有被成功请求清零")
|
|
}
|
|
}
|
|
|
|
// 回归：全局兜底超时值不再被二次换算。
|
|
func TestGatewayTimeoutIsNotMultiplied(t *testing.T) {
|
|
gw := NewGateway(nil, nil, nil, nil, 0, false, 30*time.Second, 8)
|
|
if got := gw.channelTimeout(&store.Channel{TimeoutMS: 0}); got != 30*time.Second {
|
|
t.Fatalf("兜底超时 = %s，期望 30s", got)
|
|
}
|
|
if got := gw.channelTimeout(&store.Channel{TimeoutMS: 1500}); got != 1500*time.Millisecond {
|
|
t.Fatalf("渠道超时 = %s，期望 1.5s", got)
|
|
}
|
|
if c := gw.clientFor(30 * time.Second); c.Timeout != 0 {
|
|
t.Fatalf("client.Timeout 应为 0（由 context 控制整段时长），实际 %s", c.Timeout)
|
|
}
|
|
}
|
|
|
|
// 请求体超过上限时应按客户端协议返回 413（上限由 /v1 组的 BodyLimit 中间件把住）。
|
|
func TestRequestBodyTooLarge(t *testing.T) {
|
|
h := newHarness(t, "http://127.0.0.1:1", 0, true, nil)
|
|
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions",
|
|
strings.NewReader(`{"model":"m1","messages":[{"role":"user","content":"`+strings.Repeat("x", 4096)+`"}]}`))
|
|
req.Header.Set("Authorization", "Bearer "+h.key)
|
|
// 模拟 middleware.BodyLimit：超限后 ReadAll 拿到 *http.MaxBytesError
|
|
req.Body = http.MaxBytesReader(w, req.Body, 64)
|
|
c.Request = req
|
|
|
|
h.gw.Auth(c)
|
|
if !c.IsAborted() {
|
|
h.gw.Handle(c)
|
|
}
|
|
if w.Code != http.StatusRequestEntityTooLarge {
|
|
t.Fatalf("应返回 413，实际 %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if !strings.Contains(w.Body.String(), "request_too_large") {
|
|
t.Fatalf("错误码不符：%s", w.Body.String())
|
|
}
|
|
if !strings.Contains(w.Body.String(), "32 MB") {
|
|
t.Fatalf("报错应点名上限，实际：%s", w.Body.String())
|
|
}
|
|
}
|