* fix(ops): 修复运维监控系统的关键安全和稳定性问题
## 修复内容
### P0 严重问题
1. **DNS Rebinding防护** (ops_alert_service.go)
- 实现IP钉住机制防止验证后的DNS rebinding攻击
- 自定义Transport.DialContext强制只允许拨号到验证过的公网IP
- 扩展IP黑名单,包括云metadata地址(169.254.169.254)
- 添加完整的单元测试覆盖
2. **OpsAlertService生命周期管理** (wire.go)
- 在ProvideOpsMetricsCollector中添加opsAlertService.Start()调用
- 确保stopCtx正确初始化,避免nil指针问题
- 实现防御式启动,保证服务启动顺序
3. **数据库查询排序** (ops_repo.go)
- 在ListRecentSystemMetrics中添加显式ORDER BY updated_at DESC, id DESC
- 在GetLatestSystemMetric中添加排序保证
- 避免数据库返回顺序不确定导致告警误判
### P1 重要问题
4. **并发安全** (ops_metrics_collector.go)
- 为lastGCPauseTotal字段添加sync.Mutex保护
- 防止数据竞争
5. **Goroutine泄漏** (ops_error_logger.go)
- 实现worker pool模式限制并发goroutine数量
- 使用256容量缓冲队列和10个固定worker
- 非阻塞投递,队列满时丢弃任务
6. **生命周期控制** (ops_alert_service.go)
- 添加Start/Stop方法实现优雅关闭
- 使用context控制goroutine生命周期
- 实现WaitGroup等待后台任务完成
7. **Webhook URL验证** (ops_alert_service.go)
- 防止SSRF攻击:验证scheme、禁止内网IP
- DNS解析验证,拒绝解析到私有IP的域名
- 添加8个单元测试覆盖各种攻击场景
8. **资源泄漏** (ops_repo.go)
- 修复多处defer rows.Close()问题
- 简化冗余的defer func()包装
9. **HTTP超时控制** (ops_alert_service.go)
- 创建带10秒超时的http.Client
- 添加buildWebhookHTTPClient辅助函数
- 防止HTTP请求无限期挂起
10. **数据库查询优化** (ops_repo.go)
- 将GetWindowStats的4次独立查询合并为1次CTE查询
- 减少网络往返和表扫描次数
- 显著提升性能
11. **重试机制** (ops_alert_service.go)
- 实现邮件发送重试:最多3次,指数退避(1s/2s/4s)
- 添加webhook备用通道
- 实现完整的错误处理和日志记录
12. **魔法数字** (ops_repo.go, ops_metrics_collector.go)
- 提取硬编码数字为有意义的常量
- 提高代码可读性和可维护性
## 测试验证
- ✅ go test ./internal/service -tags opsalert_unit 通过
- ✅ 所有webhook验证测试通过
- ✅ 重试机制测试通过
## 影响范围
- 运维监控系统安全性显著提升
- 系统稳定性和性能优化
- 无破坏性变更,向后兼容
* feat(ops): 运维监控系统V2 - 完整实现
## 核心功能
- 运维监控仪表盘V2(实时监控、历史趋势、告警管理)
- WebSocket实时QPS/TPS监控(30s心跳,自动重连)
- 系统指标采集(CPU、内存、延迟、错误率等)
- 多维度统计分析(按provider、model、user等维度)
- 告警规则管理(阈值配置、通知渠道)
- 错误日志追踪(详细错误信息、堆栈跟踪)
## 数据库Schema (Migration 025)
### 扩展现有表
- ops_system_metrics: 新增RED指标、错误分类、延迟指标、资源指标、业务指标
- ops_alert_rules: 新增JSONB字段(dimension_filters, notify_channels, notify_config)
### 新增表
- ops_dimension_stats: 多维度统计数据
- ops_data_retention_config: 数据保留策略配置
### 新增视图和函数
- ops_latest_metrics: 最新1分钟窗口指标(已修复字段名和window过滤)
- ops_active_alerts: 当前活跃告警(已修复字段名和状态值)
- calculate_health_score: 健康分数计算函数
## 一致性修复(98/100分)
### P0级别(阻塞Migration)
- ✅ 修复ops_latest_metrics视图字段名(latency_p99→p99_latency_ms, cpu_usage→cpu_usage_percent)
- ✅ 修复ops_active_alerts视图字段名(metric→metric_type, triggered_at→fired_at, trigger_value→metric_value, threshold→threshold_value)
- ✅ 统一告警历史表名(删除ops_alert_history,使用ops_alert_events)
- ✅ 统一API参数限制(ListMetricsHistory和ListErrorLogs的limit改为5000)
### P1级别(功能完整性)
- ✅ 修复ops_latest_metrics视图未过滤window_minutes(添加WHERE m.window_minutes = 1)
- ✅ 修复数据回填UPDATE逻辑(QPS计算改为request_count/(window_minutes*60.0))
- ✅ 添加ops_alert_rules JSONB字段后端支持(Go结构体+序列化)
### P2级别(优化)
- ✅ 前端WebSocket自动重连(指数退避1s→2s→4s→8s→16s,最大5次)
- ✅ 后端WebSocket心跳检测(30s ping,60s pong超时)
## 技术实现
### 后端 (Go)
- Handler层: ops_handler.go(REST API), ops_ws_handler.go(WebSocket)
- Service层: ops_service.go(核心逻辑), ops_cache.go(缓存), ops_alerts.go(告警)
- Repository层: ops_repo.go(数据访问), ops.go(模型定义)
- 路由: admin.go(新增ops相关路由)
- 依赖注入: wire_gen.go(自动生成)
### 前端 (Vue3 + TypeScript)
- 组件: OpsDashboardV2.vue(仪表盘主组件)
- API: ops.ts(REST API + WebSocket封装)
- 路由: index.ts(新增/admin/ops路由)
- 国际化: en.ts, zh.ts(中英文支持)
## 测试验证
- ✅ 所有Go测试通过
- ✅ Migration可正常执行
- ✅ WebSocket连接稳定
- ✅ 前后端数据结构对齐
* refactor: 代码清理和测试优化
## 测试文件优化
- 简化integration test fixtures和断言
- 优化test helper函数
- 统一测试数据格式
## 代码清理
- 移除未使用的代码和注释
- 简化concurrency_cache实现
- 优化middleware错误处理
## 小修复
- 修复gateway_handler和openai_gateway_handler的小问题
- 统一代码风格和格式
变更统计: 27个文件,292行新增,322行删除(净减少30行)
* fix(ops): 运维监控系统安全加固和功能优化
## 安全增强
- feat(security): WebSocket日志脱敏机制,防止token/api_key泄露
- feat(security): X-Forwarded-Host白名单验证,防止CSRF绕过
- feat(security): Origin策略配置化,支持strict/permissive模式
- feat(auth): WebSocket认证支持query参数传递token
## 配置优化
- feat(config): 支持环境变量配置代理信任和Origin策略
- OPS_WS_TRUST_PROXY
- OPS_WS_TRUSTED_PROXIES
- OPS_WS_ORIGIN_POLICY
- fix(ops): 错误日志查询限流从5000降至500,优化内存使用
## 架构改进
- refactor(ops): 告警服务解耦,独立运行评估定时器
- refactor(ops): OpsDashboard统一版本,移除V2分离
## 测试和文档
- test(ops): 添加WebSocket安全验证单元测试(8个测试用例)
- test(ops): 添加告警服务集成测试
- docs(api): 更新API文档,标注限流变更
- docs: 添加CHANGELOG记录breaking changes
## 修复文件
Backend:
- backend/internal/server/middleware/logger.go
- backend/internal/handler/admin/ops_handler.go
- backend/internal/handler/admin/ops_ws_handler.go
- backend/internal/server/middleware/admin_auth.go
- backend/internal/service/ops_alert_service.go
- backend/internal/service/ops_metrics_collector.go
- backend/internal/service/wire.go
Frontend:
- frontend/src/views/admin/ops/OpsDashboard.vue
- frontend/src/router/index.ts
- frontend/src/api/admin/ops.ts
Tests:
- backend/internal/handler/admin/ops_ws_handler_test.go (新增)
- backend/internal/service/ops_alert_service_integration_test.go (新增)
Docs:
- CHANGELOG.md (新增)
- docs/API-运维监控中心2.0.md (更新)
* fix(migrations): 修复calculate_health_score函数类型匹配问题
在ops_latest_metrics视图中添加显式类型转换,确保参数类型与函数签名匹配
* fix(lint): 修复golangci-lint检查发现的所有问题
- 将Redis依赖从service层移到repository层
- 添加错误检查(WebSocket连接和读取超时)
- 运行gofmt格式化代码
- 添加nil指针检查
- 删除未使用的alertService字段
修复问题:
- depguard: 3个(service层不应直接import redis)
- errcheck: 3个(未检查错误返回值)
- gofmt: 2个(代码格式问题)
- staticcheck: 4个(nil指针解引用)
- unused: 1个(未使用字段)
代码统计:
- 修改文件:11个
- 删除代码:490行
- 新增代码:105行
- 净减少:385行
801 lines
26 KiB
Go
801 lines
26 KiB
Go
// Package handler provides HTTP request handlers for the API gateway.
|
||
// It handles authentication, request routing, concurrency control, and billing validation.
|
||
package handler
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// GatewayHandler handles API gateway requests
|
||
type GatewayHandler struct {
|
||
gatewayService *service.GatewayService
|
||
geminiCompatService *service.GeminiMessagesCompatService
|
||
antigravityGatewayService *service.AntigravityGatewayService
|
||
userService *service.UserService
|
||
billingCacheService *service.BillingCacheService
|
||
concurrencyHelper *ConcurrencyHelper
|
||
opsService *service.OpsService
|
||
}
|
||
|
||
// NewGatewayHandler creates a new GatewayHandler
|
||
func NewGatewayHandler(
|
||
gatewayService *service.GatewayService,
|
||
geminiCompatService *service.GeminiMessagesCompatService,
|
||
antigravityGatewayService *service.AntigravityGatewayService,
|
||
userService *service.UserService,
|
||
concurrencyService *service.ConcurrencyService,
|
||
billingCacheService *service.BillingCacheService,
|
||
opsService *service.OpsService,
|
||
) *GatewayHandler {
|
||
return &GatewayHandler{
|
||
gatewayService: gatewayService,
|
||
geminiCompatService: geminiCompatService,
|
||
antigravityGatewayService: antigravityGatewayService,
|
||
userService: userService,
|
||
billingCacheService: billingCacheService,
|
||
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude),
|
||
opsService: opsService,
|
||
}
|
||
}
|
||
|
||
// Messages handles Claude API compatible messages endpoint
|
||
// POST /v1/messages
|
||
func (h *GatewayHandler) Messages(c *gin.Context) {
|
||
// 从context获取apiKey和user(APIKeyAuth中间件已设置)
|
||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||
return
|
||
}
|
||
|
||
// 读取请求体
|
||
body, err := io.ReadAll(c.Request.Body)
|
||
if err != nil {
|
||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||
return
|
||
}
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||
return
|
||
}
|
||
|
||
if len(body) == 0 {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
|
||
return
|
||
}
|
||
|
||
parsedReq, err := service.ParseGatewayRequest(body)
|
||
if err != nil {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||
return
|
||
}
|
||
reqModel := parsedReq.Model
|
||
reqStream := parsedReq.Stream
|
||
setOpsRequestContext(c, reqModel, reqStream)
|
||
|
||
// 验证 model 必填
|
||
if reqModel == "" {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||
return
|
||
}
|
||
|
||
// Track if we've started streaming (for error handling)
|
||
streamStarted := false
|
||
|
||
// 获取订阅信息(可能为nil)- 提前获取用于后续检查
|
||
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
||
|
||
// 0. 检查wait队列是否已满
|
||
maxWait := service.CalculateMaxWait(subject.Concurrency)
|
||
canWait, err := h.concurrencyHelper.IncrementWaitCount(c.Request.Context(), subject.UserID, maxWait)
|
||
if err != nil {
|
||
log.Printf("Increment wait count failed: %v", err)
|
||
// On error, allow request to proceed
|
||
} else if !canWait {
|
||
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later")
|
||
return
|
||
}
|
||
// 确保在函数退出时减少wait计数
|
||
defer h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
|
||
|
||
// 1. 首先获取用户并发槽位
|
||
userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted)
|
||
if err != nil {
|
||
log.Printf("User concurrency acquire failed: %v", err)
|
||
h.handleConcurrencyError(c, err, "user", streamStarted)
|
||
return
|
||
}
|
||
if userReleaseFunc != nil {
|
||
defer userReleaseFunc()
|
||
}
|
||
|
||
// 2. 【新增】Wait后二次检查余额/订阅
|
||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription); err != nil {
|
||
log.Printf("Billing eligibility check failed after wait: %v", err)
|
||
h.handleStreamingAwareError(c, http.StatusForbidden, "billing_error", err.Error(), streamStarted)
|
||
return
|
||
}
|
||
|
||
// 计算粘性会话hash
|
||
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
|
||
|
||
// 获取平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context),否则使用分组平台
|
||
platform := ""
|
||
if forcePlatform, ok := middleware2.GetForcePlatformFromContext(c); ok {
|
||
platform = forcePlatform
|
||
} else if apiKey.Group != nil {
|
||
platform = apiKey.Group.Platform
|
||
}
|
||
sessionKey := sessionHash
|
||
if platform == service.PlatformGemini && sessionHash != "" {
|
||
sessionKey = "gemini:" + sessionHash
|
||
}
|
||
|
||
if platform == service.PlatformGemini {
|
||
const maxAccountSwitches = 3
|
||
switchCount := 0
|
||
failedAccountIDs := make(map[int64]struct{})
|
||
lastFailoverStatus := 0
|
||
|
||
for {
|
||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs)
|
||
if err != nil {
|
||
if len(failedAccountIDs) == 0 {
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||
return
|
||
}
|
||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||
return
|
||
}
|
||
account := selection.Account
|
||
|
||
// 检查预热请求拦截(在账号选择后、转发前检查)
|
||
if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
|
||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
if reqStream {
|
||
sendMockWarmupStream(c, reqModel)
|
||
} else {
|
||
sendMockWarmupResponse(c, reqModel)
|
||
}
|
||
return
|
||
}
|
||
|
||
// 3. 获取账号并发槽位
|
||
accountReleaseFunc := selection.ReleaseFunc
|
||
var accountWaitRelease func()
|
||
if !selection.Acquired {
|
||
if selection.WaitPlan == nil {
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", streamStarted)
|
||
return
|
||
}
|
||
canWait, err := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
||
if err != nil {
|
||
log.Printf("Increment account wait count failed: %v", err)
|
||
} else if !canWait {
|
||
log.Printf("Account wait queue full: account=%d", account.ID)
|
||
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
||
return
|
||
} else {
|
||
// Only set release function if increment succeeded
|
||
accountWaitRelease = func() {
|
||
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||
}
|
||
}
|
||
|
||
accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
|
||
c,
|
||
account.ID,
|
||
selection.WaitPlan.MaxConcurrency,
|
||
selection.WaitPlan.Timeout,
|
||
reqStream,
|
||
&streamStarted,
|
||
)
|
||
if err != nil {
|
||
if accountWaitRelease != nil {
|
||
accountWaitRelease()
|
||
}
|
||
log.Printf("Account concurrency acquire failed: %v", err)
|
||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||
return
|
||
}
|
||
if err := h.gatewayService.BindStickySession(c.Request.Context(), sessionKey, account.ID); err != nil {
|
||
log.Printf("Bind sticky session failed: %v", err)
|
||
}
|
||
}
|
||
|
||
// 转发请求 - 根据账号平台分流
|
||
var result *service.ForwardResult
|
||
if account.Platform == service.PlatformAntigravity {
|
||
result, err = h.antigravityGatewayService.ForwardGemini(c.Request.Context(), c, account, reqModel, "generateContent", reqStream, body)
|
||
} else {
|
||
result, err = h.geminiCompatService.Forward(c.Request.Context(), c, account, body)
|
||
}
|
||
if accountReleaseFunc != nil {
|
||
accountReleaseFunc()
|
||
}
|
||
if accountWaitRelease != nil {
|
||
accountWaitRelease()
|
||
}
|
||
if err != nil {
|
||
var failoverErr *service.UpstreamFailoverError
|
||
if errors.As(err, &failoverErr) {
|
||
failedAccountIDs[account.ID] = struct{}{}
|
||
if switchCount >= maxAccountSwitches {
|
||
lastFailoverStatus = failoverErr.StatusCode
|
||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||
return
|
||
}
|
||
lastFailoverStatus = failoverErr.StatusCode
|
||
switchCount++
|
||
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
||
continue
|
||
}
|
||
// 错误响应已在Forward中处理,这里只记录日志
|
||
log.Printf("Forward request failed: %v", err)
|
||
return
|
||
}
|
||
|
||
// 异步记录使用量(subscription已在函数开头获取)
|
||
go func(result *service.ForwardResult, usedAccount *service.Account) {
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||
Result: result,
|
||
APIKey: apiKey,
|
||
User: apiKey.User,
|
||
Account: usedAccount,
|
||
Subscription: subscription,
|
||
}); err != nil {
|
||
log.Printf("Record usage failed: %v", err)
|
||
}
|
||
}(result, account)
|
||
return
|
||
}
|
||
}
|
||
|
||
const maxAccountSwitches = 10
|
||
switchCount := 0
|
||
failedAccountIDs := make(map[int64]struct{})
|
||
lastFailoverStatus := 0
|
||
|
||
for {
|
||
// 选择支持该模型的账号
|
||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs)
|
||
if err != nil {
|
||
if len(failedAccountIDs) == 0 {
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||
return
|
||
}
|
||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||
return
|
||
}
|
||
account := selection.Account
|
||
|
||
// 检查预热请求拦截(在账号选择后、转发前检查)
|
||
if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
|
||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
if reqStream {
|
||
sendMockWarmupStream(c, reqModel)
|
||
} else {
|
||
sendMockWarmupResponse(c, reqModel)
|
||
}
|
||
return
|
||
}
|
||
|
||
// 3. 获取账号并发槽位
|
||
accountReleaseFunc := selection.ReleaseFunc
|
||
var accountWaitRelease func()
|
||
if !selection.Acquired {
|
||
if selection.WaitPlan == nil {
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", streamStarted)
|
||
return
|
||
}
|
||
canWait, err := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
||
if err != nil {
|
||
log.Printf("Increment account wait count failed: %v", err)
|
||
} else if !canWait {
|
||
log.Printf("Account wait queue full: account=%d", account.ID)
|
||
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
||
return
|
||
} else {
|
||
// Only set release function if increment succeeded
|
||
accountWaitRelease = func() {
|
||
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||
}
|
||
}
|
||
|
||
accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
|
||
c,
|
||
account.ID,
|
||
selection.WaitPlan.MaxConcurrency,
|
||
selection.WaitPlan.Timeout,
|
||
reqStream,
|
||
&streamStarted,
|
||
)
|
||
if err != nil {
|
||
if accountWaitRelease != nil {
|
||
accountWaitRelease()
|
||
}
|
||
log.Printf("Account concurrency acquire failed: %v", err)
|
||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||
return
|
||
}
|
||
if err := h.gatewayService.BindStickySession(c.Request.Context(), sessionKey, account.ID); err != nil {
|
||
log.Printf("Bind sticky session failed: %v", err)
|
||
}
|
||
}
|
||
|
||
// 转发请求 - 根据账号平台分流
|
||
var result *service.ForwardResult
|
||
if account.Platform == service.PlatformAntigravity {
|
||
result, err = h.antigravityGatewayService.Forward(c.Request.Context(), c, account, body)
|
||
} else {
|
||
result, err = h.gatewayService.Forward(c.Request.Context(), c, account, parsedReq)
|
||
}
|
||
if accountReleaseFunc != nil {
|
||
accountReleaseFunc()
|
||
}
|
||
if accountWaitRelease != nil {
|
||
accountWaitRelease()
|
||
}
|
||
if err != nil {
|
||
var failoverErr *service.UpstreamFailoverError
|
||
if errors.As(err, &failoverErr) {
|
||
failedAccountIDs[account.ID] = struct{}{}
|
||
if switchCount >= maxAccountSwitches {
|
||
lastFailoverStatus = failoverErr.StatusCode
|
||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||
return
|
||
}
|
||
lastFailoverStatus = failoverErr.StatusCode
|
||
switchCount++
|
||
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
||
continue
|
||
}
|
||
// 错误响应已在Forward中处理,这里只记录日志
|
||
log.Printf("Forward request failed: %v", err)
|
||
return
|
||
}
|
||
|
||
// 异步记录使用量(subscription已在函数开头获取)
|
||
go func(result *service.ForwardResult, usedAccount *service.Account) {
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||
Result: result,
|
||
APIKey: apiKey,
|
||
User: apiKey.User,
|
||
Account: usedAccount,
|
||
Subscription: subscription,
|
||
}); err != nil {
|
||
log.Printf("Record usage failed: %v", err)
|
||
}
|
||
}(result, account)
|
||
return
|
||
}
|
||
}
|
||
|
||
// Models handles listing available models
|
||
// GET /v1/models
|
||
// Returns models based on account configurations (model_mapping whitelist)
|
||
// Falls back to default models if no whitelist is configured
|
||
func (h *GatewayHandler) Models(c *gin.Context) {
|
||
apiKey, _ := middleware2.GetAPIKeyFromContext(c)
|
||
|
||
var groupID *int64
|
||
var platform string
|
||
|
||
if apiKey != nil && apiKey.Group != nil {
|
||
groupID = &apiKey.Group.ID
|
||
platform = apiKey.Group.Platform
|
||
}
|
||
|
||
// Get available models from account configurations (without platform filter)
|
||
availableModels := h.gatewayService.GetAvailableModels(c.Request.Context(), groupID, "")
|
||
|
||
if len(availableModels) > 0 {
|
||
// Build model list from whitelist
|
||
models := make([]claude.Model, 0, len(availableModels))
|
||
for _, modelID := range availableModels {
|
||
models = append(models, claude.Model{
|
||
ID: modelID,
|
||
Type: "model",
|
||
DisplayName: modelID,
|
||
CreatedAt: "2024-01-01T00:00:00Z",
|
||
})
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": models,
|
||
})
|
||
return
|
||
}
|
||
|
||
// Fallback to default models
|
||
if platform == "openai" {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": openai.DefaultModels,
|
||
})
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": claude.DefaultModels,
|
||
})
|
||
}
|
||
|
||
// Usage handles getting account balance for CC Switch integration
|
||
// GET /v1/usage
|
||
func (h *GatewayHandler) Usage(c *gin.Context) {
|
||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
// 订阅模式:返回订阅限额信息
|
||
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() {
|
||
subscription, ok := middleware2.GetSubscriptionFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusForbidden, "subscription_error", "No active subscription")
|
||
return
|
||
}
|
||
|
||
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"isValid": true,
|
||
"planName": apiKey.Group.Name,
|
||
"remaining": remaining,
|
||
"unit": "USD",
|
||
})
|
||
return
|
||
}
|
||
|
||
// 余额模式:返回钱包余额
|
||
latestUser, err := h.userService.GetByID(c.Request.Context(), subject.UserID)
|
||
if err != nil {
|
||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to get user info")
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"isValid": true,
|
||
"planName": "钱包余额",
|
||
"remaining": latestUser.Balance,
|
||
"unit": "USD",
|
||
})
|
||
}
|
||
|
||
// calculateSubscriptionRemaining 计算订阅剩余可用额度
|
||
// 逻辑:
|
||
// 1. 如果日/周/月任一限额达到100%,返回0
|
||
// 2. 否则返回所有已配置周期中剩余额度的最小值
|
||
func (h *GatewayHandler) calculateSubscriptionRemaining(group *service.Group, sub *service.UserSubscription) float64 {
|
||
var remainingValues []float64
|
||
|
||
// 检查日限额
|
||
if group.HasDailyLimit() {
|
||
remaining := *group.DailyLimitUSD - sub.DailyUsageUSD
|
||
if remaining <= 0 {
|
||
return 0
|
||
}
|
||
remainingValues = append(remainingValues, remaining)
|
||
}
|
||
|
||
// 检查周限额
|
||
if group.HasWeeklyLimit() {
|
||
remaining := *group.WeeklyLimitUSD - sub.WeeklyUsageUSD
|
||
if remaining <= 0 {
|
||
return 0
|
||
}
|
||
remainingValues = append(remainingValues, remaining)
|
||
}
|
||
|
||
// 检查月限额
|
||
if group.HasMonthlyLimit() {
|
||
remaining := *group.MonthlyLimitUSD - sub.MonthlyUsageUSD
|
||
if remaining <= 0 {
|
||
return 0
|
||
}
|
||
remainingValues = append(remainingValues, remaining)
|
||
}
|
||
|
||
// 如果没有配置任何限额,返回-1表示无限制
|
||
if len(remainingValues) == 0 {
|
||
return -1
|
||
}
|
||
|
||
// 返回最小值
|
||
min := remainingValues[0]
|
||
for _, v := range remainingValues[1:] {
|
||
if v < min {
|
||
min = v
|
||
}
|
||
}
|
||
return min
|
||
}
|
||
|
||
// handleConcurrencyError handles concurrency-related errors with proper 429 response
|
||
func (h *GatewayHandler) handleConcurrencyError(c *gin.Context, err error, slotType string, streamStarted bool) {
|
||
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error",
|
||
fmt.Sprintf("Concurrency limit exceeded for %s, please retry later", slotType), streamStarted)
|
||
}
|
||
|
||
func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, statusCode int, streamStarted bool) {
|
||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||
}
|
||
|
||
func (h *GatewayHandler) mapUpstreamError(statusCode int) (int, string, string) {
|
||
switch statusCode {
|
||
case 401:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream authentication failed, please contact administrator"
|
||
case 403:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream access forbidden, please contact administrator"
|
||
case 429:
|
||
return http.StatusTooManyRequests, "rate_limit_error", "Upstream rate limit exceeded, please retry later"
|
||
case 529:
|
||
return http.StatusServiceUnavailable, "overloaded_error", "Upstream service overloaded, please retry later"
|
||
case 500, 502, 503, 504:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream service temporarily unavailable"
|
||
default:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream request failed"
|
||
}
|
||
}
|
||
|
||
// handleStreamingAwareError handles errors that may occur after streaming has started
|
||
func (h *GatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
||
if streamStarted {
|
||
recordOpsError(c, h.opsService, status, errType, message, "")
|
||
// Stream already started, send error as SSE event then close
|
||
flusher, ok := c.Writer.(http.Flusher)
|
||
if ok {
|
||
// Send error event in SSE format with proper JSON marshaling
|
||
errorData := map[string]any{
|
||
"type": "error",
|
||
"error": map[string]string{
|
||
"type": errType,
|
||
"message": message,
|
||
},
|
||
}
|
||
jsonBytes, err := json.Marshal(errorData)
|
||
if err != nil {
|
||
_ = c.Error(err)
|
||
return
|
||
}
|
||
errorEvent := fmt.Sprintf("data: %s\n\n", string(jsonBytes))
|
||
if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
|
||
_ = c.Error(err)
|
||
}
|
||
flusher.Flush()
|
||
}
|
||
return
|
||
}
|
||
|
||
// Normal case: return JSON response with proper status code
|
||
h.errorResponse(c, status, errType, message)
|
||
}
|
||
|
||
// errorResponse 返回Claude API格式的错误响应
|
||
func (h *GatewayHandler) errorResponse(c *gin.Context, status int, errType, message string) {
|
||
recordOpsError(c, h.opsService, status, errType, message, "")
|
||
c.JSON(status, gin.H{
|
||
"type": "error",
|
||
"error": gin.H{
|
||
"type": errType,
|
||
"message": message,
|
||
},
|
||
})
|
||
}
|
||
|
||
// CountTokens handles token counting endpoint
|
||
// POST /v1/messages/count_tokens
|
||
// 特点:校验订阅/余额,但不计算并发、不记录使用量
|
||
func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
||
// 从context获取apiKey和user(APIKeyAuth中间件已设置)
|
||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
_, ok = middleware2.GetAuthSubjectFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||
return
|
||
}
|
||
|
||
// 读取请求体
|
||
body, err := io.ReadAll(c.Request.Body)
|
||
if err != nil {
|
||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||
return
|
||
}
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||
return
|
||
}
|
||
|
||
if len(body) == 0 {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
|
||
return
|
||
}
|
||
|
||
parsedReq, err := service.ParseGatewayRequest(body)
|
||
if err != nil {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||
return
|
||
}
|
||
|
||
// 验证 model 必填
|
||
if parsedReq.Model == "" {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||
return
|
||
}
|
||
|
||
// 获取订阅信息(可能为nil)
|
||
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
||
|
||
// 校验 billing eligibility(订阅/余额)
|
||
// 【注意】不计算并发,但需要校验订阅/余额
|
||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription); err != nil {
|
||
h.errorResponse(c, http.StatusForbidden, "billing_error", err.Error())
|
||
return
|
||
}
|
||
|
||
// 计算粘性会话 hash
|
||
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
|
||
|
||
// 选择支持该模型的账号
|
||
account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, parsedReq.Model)
|
||
if err != nil {
|
||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error())
|
||
return
|
||
}
|
||
|
||
// 转发请求(不记录使用量)
|
||
if err := h.gatewayService.ForwardCountTokens(c.Request.Context(), c, account, parsedReq); err != nil {
|
||
log.Printf("Forward count_tokens request failed: %v", err)
|
||
// 错误响应已在 ForwardCountTokens 中处理
|
||
return
|
||
}
|
||
}
|
||
|
||
// isWarmupRequest 检测是否为预热请求(标题生成、Warmup等)
|
||
func isWarmupRequest(body []byte) bool {
|
||
// 快速检查:如果body不包含关键字,直接返回false
|
||
bodyStr := string(body)
|
||
if !strings.Contains(bodyStr, "title") && !strings.Contains(bodyStr, "Warmup") {
|
||
return false
|
||
}
|
||
|
||
// 解析完整请求
|
||
var req struct {
|
||
Messages []struct {
|
||
Content []struct {
|
||
Type string `json:"type"`
|
||
Text string `json:"text"`
|
||
} `json:"content"`
|
||
} `json:"messages"`
|
||
System []struct {
|
||
Text string `json:"text"`
|
||
} `json:"system"`
|
||
}
|
||
if err := json.Unmarshal(body, &req); err != nil {
|
||
return false
|
||
}
|
||
|
||
// 检查 messages 中的标题提示模式
|
||
for _, msg := range req.Messages {
|
||
for _, content := range msg.Content {
|
||
if content.Type == "text" {
|
||
if strings.Contains(content.Text, "Please write a 5-10 word title for the following conversation:") ||
|
||
content.Text == "Warmup" {
|
||
return true
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 检查 system 中的标题提取模式
|
||
for _, system := range req.System {
|
||
if strings.Contains(system.Text, "nalyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title") {
|
||
return true
|
||
}
|
||
}
|
||
|
||
return false
|
||
}
|
||
|
||
// sendMockWarmupStream 发送流式 mock 响应(用于预热请求拦截)
|
||
func sendMockWarmupStream(c *gin.Context, model string) {
|
||
c.Header("Content-Type", "text/event-stream")
|
||
c.Header("Cache-Control", "no-cache")
|
||
c.Header("Connection", "keep-alive")
|
||
c.Header("X-Accel-Buffering", "no")
|
||
|
||
// Build message_start event with proper JSON marshaling
|
||
messageStart := map[string]any{
|
||
"type": "message_start",
|
||
"message": map[string]any{
|
||
"id": "msg_mock_warmup",
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"model": model,
|
||
"content": []any{},
|
||
"stop_reason": nil,
|
||
"stop_sequence": nil,
|
||
"usage": map[string]int{
|
||
"input_tokens": 10,
|
||
"output_tokens": 0,
|
||
},
|
||
},
|
||
}
|
||
messageStartJSON, _ := json.Marshal(messageStart)
|
||
|
||
events := []string{
|
||
`event: message_start` + "\n" + `data: ` + string(messageStartJSON),
|
||
`event: content_block_start` + "\n" + `data: {"content_block":{"text":"","type":"text"},"index":0,"type":"content_block_start"}`,
|
||
`event: content_block_delta` + "\n" + `data: {"delta":{"text":"New","type":"text_delta"},"index":0,"type":"content_block_delta"}`,
|
||
`event: content_block_delta` + "\n" + `data: {"delta":{"text":" Conversation","type":"text_delta"},"index":0,"type":"content_block_delta"}`,
|
||
`event: content_block_stop` + "\n" + `data: {"index":0,"type":"content_block_stop"}`,
|
||
`event: message_delta` + "\n" + `data: {"delta":{"stop_reason":"end_turn","stop_sequence":null},"type":"message_delta","usage":{"input_tokens":10,"output_tokens":2}}`,
|
||
`event: message_stop` + "\n" + `data: {"type":"message_stop"}`,
|
||
}
|
||
|
||
for _, event := range events {
|
||
_, _ = c.Writer.WriteString(event + "\n\n")
|
||
c.Writer.Flush()
|
||
time.Sleep(20 * time.Millisecond)
|
||
}
|
||
}
|
||
|
||
// sendMockWarmupResponse 发送非流式 mock 响应(用于预热请求拦截)
|
||
func sendMockWarmupResponse(c *gin.Context, model string) {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"id": "msg_mock_warmup",
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"model": model,
|
||
"content": []gin.H{{"type": "text", "text": "New Conversation"}},
|
||
"stop_reason": "end_turn",
|
||
"usage": gin.H{
|
||
"input_tokens": 10,
|
||
"output_tokens": 2,
|
||
},
|
||
})
|
||
}
|