fix: 修掉模型调用链路上的四处失效逻辑
- 超时：NewGateway 里 gw_timeout*time.Second 把已是 Duration 的值二次换算， int64 溢出后无论 OT_PROXY_TIMEOUT 填多少都等于没有超时，渠道 timeout_ms 也从没被读过。现在非流式用 context deadline 约束整段时长，流式只约束首包 （Transport.ResponseHeaderTimeout），正文不限时长，长回答不会被砍。 - 断连：上游请求改用 NewRequestWithContext 绑定客户端 context，客户端断开即 取消上游并记 canceled，不再空转生成继续付费；PATCH 已写 200 头后不再往 响应里追加错误体。 - 白名单：APIKey.AllowedModels 以前只存不用，现在 Auth 注入 context 并参与 判定（用户黑名单 > 密钥白名单 > 用户白名单 > 全局）。 - 选路：模型已有渠道绑定时只用绑定渠道，都不可用直接 503 no_channel，不再 把请求发给不认识它的渠道；新增运行期健康（连续 3 次网络错误/超时/429/5xx 进 30s 冷却，成功清零），与 60s 一次的 /models 探测互补。 - 密钥 API：PATCH 回显更新后的完整对象；限额填 0 统一视为不限制（原来 token 配额存 0 会让密钥第一次请求就被 429 打死，而请求配额存 0 又等于 不限制，两个字段语义相反）；前端编辑框总是提交白名单，清空能真的取消。 测试：新增 proxy 链路集成测试（超时生效、流式不被截断、断连取消并记账、 密钥白名单 403、连续 429 后候选剔除）与 channel 用例（绑定模型不退全局、 冷却阈值与到期恢复、成功清零）。
This commit is contained in:
Sakurasan committed
1 parent
d911881499
commit
e865f69a41
10 files changed
+659
-69
No files matched your search
@@ -6,6 +6,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -23,9 +24,10 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
CtxUserID = "proxy_user_id"
|
||||
CtxKeyID = "proxy_key_id"
|
||||
CtxTrace = "proxy_trace_id"
|
||||
CtxUserID = "proxy_user_id"
|
||||
CtxKeyID = "proxy_key_id"
|
||||
CtxKeyModel = "proxy_key_models"
|
||||
CtxTrace = "proxy_trace_id"
|
||||
)
|
||||
|
||||
type Gateway struct {
|
||||
@@ -36,7 +38,11 @@ type Gateway struct {
|
||||
lim *ratelimit.Limiter
|
||||
userRPS int
|
||||
logRaw bool
|
||||
hc *http.Client
|
||||
|
||||
// 上游超时兜底值（渠道 timeout_ms 非法时使用）；按超时值缓存 Transport 复用连接池。
|
||||
gwTimeout time.Duration
|
||||
tpMu sync.Mutex
|
||||
tps map[time.Duration]*http.Transport
|
||||
|
||||
policyMu sync.Mutex
|
||||
policy modelPolicy
|
||||
@@ -58,7 +64,7 @@ func (g *Gateway) ResetModelPolicy() {
|
||||
g.policyMu.Unlock()
|
||||
}
|
||||
|
||||
// globalModelRestrictions 读取全局模型允许/禁止列表（缓存 30s）。
|
||||
// globalModelRestrictions 读取全局模型允许/禁止列表（缓存 modelPolicyTTL）。
|
||||
func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
|
||||
g.policyMu.Lock()
|
||||
defer g.policyMu.Unlock()
|
||||
@@ -75,18 +81,24 @@ func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
|
||||
return
|
||||
}
|
||||
|
||||
// checkModelAllowed 模型访问控制：用户级 > 全局。
|
||||
// checkModelAllowed 模型访问控制：用户级 > 密钥级 > 全局。
|
||||
// 1. 用户禁止列表命中 → 拒绝
|
||||
// 2. 用户允许列表非空 → 仅列表内可访问（不再看全局）
|
||||
// 3. 全局禁止命中 → 拒绝
|
||||
// 4. 全局允许列表非空 → 仅列表内可访问
|
||||
func (g *Gateway) checkModelAllowed(u *store.User, model string) bool {
|
||||
// 2. 密钥允许列表非空且不含该模型 → 拒绝（密钥级只收窄，不放宽）
|
||||
// 3. 用户允许列表非空 → 仅列表内可访问，不再看全局（管理员显式授权该用户，短路全局）
|
||||
// 4. 全局禁止命中 → 拒绝
|
||||
// 5. 全局允许列表非空 → 仅列表内可访问
|
||||
func (g *Gateway) checkModelAllowed(c *gin.Context, u *store.User, model string) bool {
|
||||
if model == "" {
|
||||
return true
|
||||
}
|
||||
if contains(u.DeniedModels, model) {
|
||||
return false
|
||||
}
|
||||
if keyAllow, ok := c.Get(CtxKeyModel); ok {
|
||||
if list, _ := keyAllow.([]string); len(list) > 0 && !contains(list, model) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
if len(u.AllowedModels) > 0 {
|
||||
return contains(u.AllowedModels, model)
|
||||
}
|
||||
@@ -111,17 +123,57 @@ func contains(list []string, s string) bool {
|
||||
|
||||
func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ratelimit.Limiter, userRPS int, logRaw bool, gw_timeout time.Duration) *Gateway {
|
||||
return &Gateway{
|
||||
db: db,
|
||||
ch: channel.NewService(db, enc),
|
||||
rec: rec,
|
||||
enc: enc,
|
||||
lim: lim,
|
||||
userRPS: userRPS,
|
||||
logRaw: logRaw,
|
||||
hc: &http.Client{Timeout: gw_timeout * time.Second},
|
||||
db: db,
|
||||
ch: channel.NewService(db, enc),
|
||||
rec: rec,
|
||||
enc: enc,
|
||||
lim: lim,
|
||||
userRPS: userRPS,
|
||||
logRaw: logRaw,
|
||||
gwTimeout: gw_timeout,
|
||||
tps: map[time.Duration]*http.Transport{},
|
||||
}
|
||||
}
|
||||
|
||||
// channelTimeout 该渠道的上游超时：timeout_ms 优先，未配置则用全局兜底值。
|
||||
func (g *Gateway) channelTimeout(ch *store.Channel) time.Duration {
|
||||
if ch.TimeoutMS > 0 {
|
||||
return time.Duration(ch.TimeoutMS) * time.Millisecond
|
||||
}
|
||||
return g.gwTimeout
|
||||
}
|
||||
|
||||
// clientFor 返回带该超时的上游 client。Transport 按超时值缓存以复用连接池，
|
||||
// ResponseHeaderTimeout 只约束"连上之后多久要回首包"——流式正文不限时长，长回答不会被砍。
|
||||
// 整段时长（含读 body）由请求 context 的 deadline 控制，见 proxyOne。
|
||||
func (g *Gateway) clientFor(d time.Duration) *http.Client {
|
||||
if d < 0 {
|
||||
d = 0
|
||||
}
|
||||
g.tpMu.Lock()
|
||||
defer g.tpMu.Unlock()
|
||||
tp, ok := g.tps[d]
|
||||
if !ok {
|
||||
tp = &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
ForceAttemptHTTP2: true,
|
||||
DialContext: (&net.Dialer{
|
||||
Timeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
}).DialContext,
|
||||
MaxIdleConns: 100,
|
||||
MaxIdleConnsPerHost: 32,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
}
|
||||
if d > 0 {
|
||||
tp.ResponseHeaderTimeout = d
|
||||
}
|
||||
g.tps[d] = tp
|
||||
}
|
||||
return &http.Client{Transport: tp}
|
||||
}
|
||||
|
||||
// Auth 代理鉴权中间件：Bearer sk-xxx → 哈希查表 → 校验状态/过期/模型白名单。
|
||||
func (g *Gateway) Auth(c *gin.Context) {
|
||||
// 先按路径确定客户端协议，保证 Auth 阶段错误也按协议格式返回
|
||||
@@ -186,6 +238,7 @@ func (g *Gateway) Auth(c *gin.Context) {
|
||||
|
||||
c.Set(CtxUserID, u.ID)
|
||||
c.Set(CtxKeyID, k.ID)
|
||||
c.Set(CtxKeyModel, k.AllowedModels)
|
||||
c.Set(CtxTrace, newTraceID())
|
||||
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
|
||||
c.Next()
|
||||
|
||||
Reference in new issue
Block a user