Sakurasan
e865f69a41
- 超时：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 用例（绑定模型不退全局、 冷却阈值与到期恢复、成功清零）。
393 lines
12 KiB
Go
393 lines
12 KiB
Go
// Package proxy API 网关核心：代理 /v1/chat/completions、/v1/responses、/v1/messages、/v1/models。
|
||
// M1 直通 OpenAI 渠道；M4 起按客户端协议 × 渠道协议自动转换（见 convert）。
|
||
package proxy
|
||
|
||
import (
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"net"
|
||
"net/http"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/openteam/server/internal/channel"
|
||
"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/proxy/convert"
|
||
"github.com/openteam/server/internal/store"
|
||
"github.com/openteam/server/internal/usage"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
const (
|
||
CtxUserID = "proxy_user_id"
|
||
CtxKeyID = "proxy_key_id"
|
||
CtxKeyModel = "proxy_key_models"
|
||
CtxTrace = "proxy_trace_id"
|
||
)
|
||
|
||
type Gateway struct {
|
||
db *gorm.DB
|
||
ch *channel.Service
|
||
rec *usage.Recorder
|
||
enc *crypto.Encryptor
|
||
lim *ratelimit.Limiter
|
||
userRPS int
|
||
logRaw bool
|
||
|
||
// 上游超时兜底值（渠道 timeout_ms 非法时使用）；按超时值缓存 Transport 复用连接池。
|
||
gwTimeout time.Duration
|
||
tpMu sync.Mutex
|
||
tps map[time.Duration]*http.Transport
|
||
|
||
policyMu sync.Mutex
|
||
policy modelPolicy
|
||
}
|
||
|
||
// modelPolicy 全局模型限制策略（来自系统配置，短时缓存）。
|
||
type modelPolicy struct {
|
||
allow []string
|
||
deny []string
|
||
at time.Time
|
||
}
|
||
|
||
const modelPolicyTTL = 5 * time.Second
|
||
|
||
// ResetModelPolicy 清空全局模型限制缓存（系统配置保存后调用）。
|
||
func (g *Gateway) ResetModelPolicy() {
|
||
g.policyMu.Lock()
|
||
g.policy = modelPolicy{}
|
||
g.policyMu.Unlock()
|
||
}
|
||
|
||
// globalModelRestrictions 读取全局模型允许/禁止列表（缓存 modelPolicyTTL）。
|
||
func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
|
||
g.policyMu.Lock()
|
||
defer g.policyMu.Unlock()
|
||
if time.Since(g.policy.at) < modelPolicyTTL {
|
||
return g.policy.allow, g.policy.deny
|
||
}
|
||
var raw string
|
||
g.db.Model(&store.SystemConfig{}).Where("key = ?", "model_allowlist").Pluck("value", &raw)
|
||
_ = json.Unmarshal([]byte(raw), &allow)
|
||
raw = ""
|
||
g.db.Model(&store.SystemConfig{}).Where("key = ?", "model_denylist").Pluck("value", &raw)
|
||
_ = json.Unmarshal([]byte(raw), &deny)
|
||
g.policy = modelPolicy{allow: allow, deny: deny, at: time.Now()}
|
||
return
|
||
}
|
||
|
||
// checkModelAllowed 模型访问控制：用户级 > 密钥级 > 全局。
|
||
// 1. 用户禁止列表命中 → 拒绝
|
||
// 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)
|
||
}
|
||
allow, deny := g.globalModelRestrictions()
|
||
if contains(deny, model) {
|
||
return false
|
||
}
|
||
if len(allow) > 0 {
|
||
return contains(allow, model)
|
||
}
|
||
return true
|
||
}
|
||
|
||
func contains(list []string, s string) bool {
|
||
for _, v := range list {
|
||
if v == s {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
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,
|
||
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 阶段错误也按协议格式返回
|
||
switch c.Request.URL.Path {
|
||
case "/v1/messages":
|
||
c.Set("protocol", convert.ProtoMessages)
|
||
case "/v1/responses":
|
||
c.Set("protocol", convert.ProtoResponses)
|
||
default:
|
||
c.Set("protocol", convert.ProtoChat)
|
||
}
|
||
|
||
auth := c.GetHeader("Authorization")
|
||
key := strings.TrimPrefix(auth, "Bearer ")
|
||
key = strings.TrimSpace(key)
|
||
if key == "" {
|
||
// Anthropic 客户端（Claude Code / SDK）用 x-api-key 头而不是 Authorization
|
||
key = strings.TrimSpace(c.GetHeader("x-api-key"))
|
||
}
|
||
if !apikey.Valid(key) {
|
||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key format. Expected: Authorization: Bearer sk-ot-... or x-api-key: sk-ot-...")
|
||
c.Abort()
|
||
return
|
||
}
|
||
hash := apikey.Hash(key)
|
||
var k store.APIKey
|
||
if err := g.db.Where("key_hash = ? AND status = ?", hash, store.KeyStatusActive).First(&k).Error; err != nil {
|
||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||
c.Abort()
|
||
return
|
||
}
|
||
var u store.User
|
||
if err := g.db.First(&u, k.UserID).Error; err != nil || u.Status != store.UserStatusActive {
|
||
apiError(c, http.StatusForbidden, "user_disabled", "User account is disabled")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if k.ExpiresAt != nil && time.Now().After(*k.ExpiresAt) {
|
||
apiError(c, http.StatusUnauthorized, "key_expired", "API key has expired")
|
||
c.Abort()
|
||
return
|
||
}
|
||
|
||
// 限流与配额（内存计数）
|
||
if g.lim != nil {
|
||
if k.QuotaRequestsPerDay != nil && !g.lim.AllowRequestDaily(k.ID, *k.QuotaRequestsPerDay) {
|
||
apiError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Daily request quota exceeded for this API key")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if k.QuotaTokensPerDay != nil && g.lim.TokensUsed(k.ID) >= *k.QuotaTokensPerDay {
|
||
apiError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Daily token quota exceeded for this API key")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if !g.lim.AllowUserRate(u.ID, g.userRPS) {
|
||
apiError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Too many requests. Please slow down.")
|
||
c.Abort()
|
||
return
|
||
}
|
||
}
|
||
|
||
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()
|
||
}
|
||
|
||
// Handle 路由到对应协议处理器。
|
||
func (g *Gateway) Handle(c *gin.Context) {
|
||
switch {
|
||
case c.Request.URL.Path == "/v1/chat/completions":
|
||
g.chatCompletions(c)
|
||
case c.Request.URL.Path == "/v1/responses":
|
||
g.responses(c)
|
||
case c.Request.URL.Path == "/v1/messages":
|
||
g.messages(c)
|
||
case c.Request.URL.Path == "/v1/models" && c.Request.Method == http.MethodGet:
|
||
g.models(c)
|
||
default:
|
||
apiError(c, http.StatusNotFound, "not_found", "Unknown endpoint: "+c.Request.URL.Path)
|
||
}
|
||
}
|
||
|
||
// models GET /v1/models：返回对外可见的模型（启用且至少绑定到一个启用且健康的渠道；OpenAI 风格）。
|
||
func (g *Gateway) models(c *gin.Context) {
|
||
ids := g.ch.AvailableModelIDs()
|
||
if len(ids) == 0 {
|
||
c.JSON(http.StatusOK, gin.H{"object": "list", "data": []gin.H{}})
|
||
return
|
||
}
|
||
var ms []store.Model
|
||
if err := g.db.Where("id IN ? AND enabled = ?", ids, true).Order("sort ASC, id ASC").Find(&ms).Error; err != nil {
|
||
apiError(c, http.StatusInternalServerError, "internal_error", "failed to load models")
|
||
return
|
||
}
|
||
data := make([]gin.H, 0, len(ms))
|
||
for _, m := range ms {
|
||
data = append(data, gin.H{
|
||
"id": m.Name,
|
||
"object": "model",
|
||
"created": m.CreatedAt.Unix(),
|
||
"owned_by": "openteam",
|
||
})
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{"object": "list", "data": data})
|
||
}
|
||
|
||
// candidateChannels 返回可用渠道候选（按模型绑定优先，退化全局；携带模型映射）。
|
||
func (g *Gateway) candidateChannels(model string) []channel.Candidate {
|
||
return g.ch.Candidates(model)
|
||
}
|
||
|
||
// resolveUser 取当前用户（含余额）。
|
||
func (g *Gateway) resolveUser(c *gin.Context) (*store.User, bool) {
|
||
uid, _ := c.Get(CtxUserID)
|
||
var u store.User
|
||
if err := g.db.First(&u, uid).Error; err != nil {
|
||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||
return nil, false
|
||
}
|
||
return &u, true
|
||
}
|
||
|
||
// checkBalance 余额不足返回 402（PLANNING §4.4.3）。
|
||
func (g *Gateway) checkBalance(c *gin.Context, u *store.User) bool {
|
||
if u.Balance <= 0 {
|
||
apiError(c, http.StatusPaymentRequired, "insufficient_balance", "Insufficient balance. Please recharge or contact admin.")
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 协议分派
|
||
|
||
func upstreamPath(proto string) string {
|
||
switch proto {
|
||
case convert.ProtoMessages:
|
||
return "/messages"
|
||
case convert.ProtoResponses:
|
||
return "/responses"
|
||
default:
|
||
return "/chat/completions"
|
||
}
|
||
}
|
||
|
||
// upstreamPlan 描述一次代理请求的上游访问方式。
|
||
type upstreamPlan struct {
|
||
proto string // 上游协议（用于分协议 base_url）
|
||
path string // 上游路径
|
||
body []byte // 已转换的请求体
|
||
lineConv func([]byte) []byte // 流式逐行转换（nil=直通）
|
||
bodyConv func([]byte) ([]byte, error) // 非流式响应体转换（nil=直通）
|
||
}
|
||
|
||
// conversionTarget 决定客户端协议在渠道上的处理方式：
|
||
// 渠道声明支持该协议则直通（返回原协议）；否则转为其首选支持协议（chat > messages > responses）。
|
||
func conversionTarget(formats []string, clientProto string) string {
|
||
for _, f := range formats {
|
||
if f == clientProto {
|
||
return clientProto
|
||
}
|
||
}
|
||
for _, p := range []string{convert.ProtoChat, convert.ProtoMessages, convert.ProtoResponses} {
|
||
for _, f := range formats {
|
||
if f == p {
|
||
return p
|
||
}
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// prepareUpstream 计算上游访问计划：渠道声明支持客户端协议则直通，否则转换；
|
||
// 应用模型名称映射（upstream_model）。
|
||
func prepareUpstream(ch *store.Channel, clientProto string, body []byte, upstreamModel string) (*upstreamPlan, error) {
|
||
target := conversionTarget(ch.FormatsEffective(), clientProto)
|
||
if target == "" {
|
||
return nil, fmt.Errorf("channel %q declares no supported protocol format", ch.Name)
|
||
}
|
||
plan := &upstreamPlan{proto: target, path: upstreamPath(target), body: body}
|
||
if target != clientProto {
|
||
converted, err := convert.ConvertRequest(body, clientProto, target)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
plan.body = converted
|
||
plan.lineConv = convert.NewStreamTransformer(target, clientProto)
|
||
plan.bodyConv = func(b []byte) ([]byte, error) { return convert.ConvertResponse(b, target, clientProto) }
|
||
}
|
||
// 模型名称映射：把请求体 model 字段改写为渠道侧的 upstream_model
|
||
if upstreamModel != "" {
|
||
if out, err := rewriteModel(plan.body, upstreamModel); err == nil {
|
||
plan.body = out
|
||
}
|
||
}
|
||
return plan, nil
|
||
}
|
||
|
||
// rewriteModel 改写请求体中的 model 字段（三种协议 model 都在顶层）。
|
||
func rewriteModel(body []byte, model string) ([]byte, error) {
|
||
var m map[string]any
|
||
if err := json.Unmarshal(body, &m); err != nil {
|
||
return body, nil
|
||
}
|
||
if cur, _ := m["model"].(string); cur == model {
|
||
return body, nil
|
||
}
|
||
m["model"] = model
|
||
return json.Marshal(m)
|
||
}
|
||
|
||
var errNoChannel = errors.New("no available channel")
|