实现 Master-Slave 架构的跨区域用户和配额同步功能,支持国内主节点与 海外从节点之间的用户创建推送、余额同步、预扣费批量回传。 主要变更: - 新增 region_sync 服务层(SyncManager/SyncClient/推送逻辑) - User 模型扩展 Source/RemoteUserId/SyncedQuota/LastSyncAt 字段 - BillingSession 支持 SyncedUserFunding 资金来源 - 新增 SyncAuth 中间件(X-Sync-API-Key + 常量时间比较) - 新增 PendingSyncRecord/QuotaSyncLog 数据模型 - 前端新增区域同步配置面板 - 支持 SESSION_NAME 环境变量避免多节点 cookie 冲突 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>feat/alipay-payment
| @@ -0,0 +1,33 @@ | |||
| package common | |||
| import "errors" | |||
| const ( | |||
| ForeignUserIDStart = 10000000 | |||
| UserSourceLocal = "local" | |||
| UserSourceSynced = "synced" | |||
| ) | |||
| func IsSyncedUser(source string, id int) bool { | |||
| return source == UserSourceSynced && id > 0 && id < ForeignUserIDStart | |||
| } | |||
| func IsLocalUser(source string, id int) bool { | |||
| if source == UserSourceLocal { | |||
| return true | |||
| } | |||
| if id >= ForeignUserIDStart { | |||
| return true | |||
| } | |||
| return false | |||
| } | |||
| func ValidateUserSource(source string, id int) error { | |||
| if source == UserSourceSynced && id >= ForeignUserIDStart { | |||
| return errors.New("synced 用户 ID 必须小于 1千万") | |||
| } | |||
| if source == UserSourceLocal && id > 0 && id < ForeignUserIDStart { | |||
| return errors.New("local 用户 ID 必须大于等于 1千万") | |||
| } | |||
| return nil | |||
| } | |||
| @@ -0,0 +1,93 @@ | |||
| package common | |||
| import ( | |||
| "testing" | |||
| ) | |||
| func TestForeignUserIDStart(t *testing.T) { | |||
| if ForeignUserIDStart != 10000000 { | |||
| t.Errorf("ForeignUserIDStart = %d, want 10000000", ForeignUserIDStart) | |||
| } | |||
| } | |||
| func TestUserSourceConstants(t *testing.T) { | |||
| if UserSourceLocal != "local" { | |||
| t.Errorf("UserSourceLocal = %q, want %q", UserSourceLocal, "local") | |||
| } | |||
| if UserSourceSynced != "synced" { | |||
| t.Errorf("UserSourceSynced = %q, want %q", UserSourceSynced, "synced") | |||
| } | |||
| } | |||
| func TestIsSyncedUser(t *testing.T) { | |||
| tests := []struct { | |||
| name string | |||
| source string | |||
| id int | |||
| want bool | |||
| }{ | |||
| {"synced with normal id", UserSourceSynced, 123, true}, | |||
| {"synced with id just below threshold", UserSourceSynced, ForeignUserIDStart - 1, true}, | |||
| {"synced with id 0", UserSourceSynced, 0, false}, | |||
| {"synced with id at threshold", UserSourceSynced, ForeignUserIDStart, false}, | |||
| {"local with normal id", UserSourceLocal, 123, false}, | |||
| {"empty source with normal id", "", 123, false}, | |||
| } | |||
| for _, tt := range tests { | |||
| t.Run(tt.name, func(t *testing.T) { | |||
| if got := IsSyncedUser(tt.source, tt.id); got != tt.want { | |||
| t.Errorf("IsSyncedUser(%q, %d) = %v, want %v", tt.source, tt.id, got, tt.want) | |||
| } | |||
| }) | |||
| } | |||
| } | |||
| func TestIsLocalUser(t *testing.T) { | |||
| tests := []struct { | |||
| name string | |||
| source string | |||
| id int | |||
| want bool | |||
| }{ | |||
| {"local at threshold", UserSourceLocal, ForeignUserIDStart, true}, | |||
| {"local above threshold", UserSourceLocal, ForeignUserIDStart + 100, true}, | |||
| {"local with id 0", UserSourceLocal, 0, true}, | |||
| {"synced with small id", UserSourceSynced, 123, false}, | |||
| {"empty source at threshold", "", ForeignUserIDStart, true}, | |||
| {"empty source with small id", "", 500, false}, | |||
| } | |||
| for _, tt := range tests { | |||
| t.Run(tt.name, func(t *testing.T) { | |||
| if got := IsLocalUser(tt.source, tt.id); got != tt.want { | |||
| t.Errorf("IsLocalUser(%q, %d) = %v, want %v", tt.source, tt.id, got, tt.want) | |||
| } | |||
| }) | |||
| } | |||
| } | |||
| func TestValidateUserSource(t *testing.T) { | |||
| tests := []struct { | |||
| name string | |||
| source string | |||
| id int | |||
| wantErr bool | |||
| }{ | |||
| {"synced with valid id", UserSourceSynced, 100, false}, | |||
| {"synced at threshold", UserSourceSynced, ForeignUserIDStart, true}, | |||
| {"local at threshold", UserSourceLocal, ForeignUserIDStart, false}, | |||
| {"local with small non-zero id", UserSourceLocal, 500, true}, | |||
| {"local with id 0", UserSourceLocal, 0, false}, | |||
| {"empty source with id 0", "", 0, false}, | |||
| } | |||
| for _, tt := range tests { | |||
| t.Run(tt.name, func(t *testing.T) { | |||
| err := ValidateUserSource(tt.source, tt.id) | |||
| if (err != nil) != tt.wantErr { | |||
| t.Errorf("ValidateUserSource(%q, %d) error = %v, wantErr %v", tt.source, tt.id, err, tt.wantErr) | |||
| } | |||
| }) | |||
| } | |||
| } | |||
| @@ -51,6 +51,7 @@ const ( | |||
| ContextKeyUserGroup ContextKey = "user_group" | |||
| ContextKeyUsingGroup ContextKey = "group" | |||
| ContextKeyUserName ContextKey = "username" | |||
| ContextKeyUserSource ContextKey = "user_source" | |||
| ContextKeyLocalCountTokens ContextKey = "local_count_tokens" | |||
| @@ -0,0 +1,273 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service/region_sync" | |||
| "github.com/gin-gonic/gin" | |||
| "gorm.io/gorm" | |||
| ) | |||
| // 统一错误响应 | |||
| func syncErrorResponse(c *gin.Context, status int, err string) { | |||
| c.JSON(status, gin.H{"success": false, "error": err}) | |||
| } | |||
| // 统一成功响应 | |||
| func syncSuccessResponse(c *gin.Context, message string) { | |||
| c.JSON(http.StatusOK, gin.H{"success": true, "message": message}) | |||
| } | |||
| // ReceiveSyncedUserCreate 接收从 master 节点推送过来的用户创建请求 | |||
| func ReceiveSyncedUserCreate(c *gin.Context) { | |||
| var req region_sync.SyncUserRequest | |||
| if err := c.ShouldBindJSON(&req); err != nil { | |||
| syncErrorResponse(c, http.StatusBadRequest, "invalid request body") | |||
| return | |||
| } | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveSyncedUserCreate: username=%s, remoteUserId=%d, quota=%d", req.Username, req.RemoteUserId, req.Quota) | |||
| // 检查用户是否已存在 | |||
| var existingUser model.User | |||
| if err := model.DB.Where("remote_user_id = ?", req.RemoteUserId).First(&existingUser).Error; err == nil { | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveSyncedUserCreate: user already exists, remoteUserId=%d", req.RemoteUserId) | |||
| syncSuccessResponse(c, "user already exists") | |||
| return | |||
| } | |||
| // 创建新用户 | |||
| user := createSyncedUser(&req) | |||
| if err := model.DB.Create(&user).Error; err != nil { | |||
| logger.LogError(c.Request.Context(), "failed to create synced user: "+err.Error()) | |||
| syncErrorResponse(c, http.StatusInternalServerError, "failed to create user") | |||
| return | |||
| } | |||
| // 记录同步日志 | |||
| recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypeUserCreate, "", 0, 0, model.SyncDirectionCnToOv) | |||
| logger.LogInfo(c.Request.Context(), "created synced user "+user.Username) | |||
| syncSuccessResponse(c, "user created successfully") | |||
| } | |||
| // createSyncedUser 根据请求创建同步用户 | |||
| func createSyncedUser(req *region_sync.SyncUserRequest) model.User { | |||
| affCode := req.AffCode | |||
| if affCode == "" { | |||
| affCode = common.GetRandomString(4) | |||
| } | |||
| displayName := req.DisplayName | |||
| if displayName == "" { | |||
| displayName = req.Username | |||
| } | |||
| group := req.Group | |||
| if group == "" { | |||
| group = "default" | |||
| } | |||
| return model.User{ | |||
| Username: req.Username, | |||
| Email: req.Email, | |||
| Password: req.PasswordHash, | |||
| DisplayName: displayName, | |||
| Quota: req.Quota, | |||
| RemoteUserId: req.RemoteUserId, | |||
| Source: common.UserSourceSynced, | |||
| Group: group, | |||
| AffCode: affCode, | |||
| } | |||
| } | |||
| // ReceiveQuotaUpdate 接收 master 节点的余额更新请求 | |||
| func ReceiveQuotaUpdate(c *gin.Context) { | |||
| var req region_sync.UpdateQuotaRequest | |||
| if err := c.ShouldBindJSON(&req); err != nil { | |||
| syncErrorResponse(c, http.StatusBadRequest, "invalid request body") | |||
| return | |||
| } | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] ReceiveQuotaUpdate: remoteUserId=%d, newQuota=%d", req.RemoteUserId, req.Quota) | |||
| // 查找同步用户 | |||
| var user model.User | |||
| if err := model.DB.Where("remote_user_id = ? AND source = ?", req.RemoteUserId, common.UserSourceSynced).First(&user).Error; err != nil { | |||
| logger.LogError(c.Request.Context(), "synced user not found: "+err.Error()) | |||
| syncErrorResponse(c, http.StatusNotFound, "user not found") | |||
| return | |||
| } | |||
| // 更新余额(选择性更新,避免覆盖其他字段) | |||
| if err := model.DB.Model(&user).Updates(map[string]interface{}{ | |||
| "synced_quota": req.Quota, | |||
| "last_sync_at": time.Now().Unix(), | |||
| }).Error; err != nil { | |||
| logger.LogError(c.Request.Context(), "failed to update synced quota: "+err.Error()) | |||
| syncErrorResponse(c, http.StatusInternalServerError, "failed to update quota") | |||
| return | |||
| } | |||
| // 记录同步日志 | |||
| recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypeQuotaChange, "", 0, req.Quota, model.SyncDirectionCnToOv) | |||
| logger.LogInfo(c.Request.Context(), "updated synced quota for user "+user.Username) | |||
| syncSuccessResponse(c, "quota updated successfully") | |||
| } | |||
| // QueryUserQuota 查询用户余额(供 slave 节点调用 master 节点) | |||
| func QueryUserQuota(c *gin.Context) { | |||
| var req region_sync.QueryQuotaRequest | |||
| if err := c.ShouldBindJSON(&req); err != nil { | |||
| syncErrorResponse(c, http.StatusBadRequest, "invalid request body") | |||
| return | |||
| } | |||
| // 查找用户 | |||
| var user model.User | |||
| if err := model.DB.Where("id = ?", req.UserId).First(&user).Error; err != nil { | |||
| syncErrorResponse(c, http.StatusNotFound, "user not found") | |||
| return | |||
| } | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] QueryUserQuota: userId=%d, quota=%d", user.Id, user.Quota) | |||
| // 记录同步日志 | |||
| recordSyncLog(c, user.Id, user.RemoteUserId, model.SyncTypePreConsumeQuery, "", user.Quota, 0, model.SyncDirectionOvToCn) | |||
| c.JSON(http.StatusOK, gin.H{"success": true, "quota": user.Quota}) | |||
| } | |||
| // BatchDeductQuota 批量扣费(供 slave 节点调用 master 节点) | |||
| func BatchDeductQuota(c *gin.Context) { | |||
| var req region_sync.BatchDeductRequest | |||
| if err := c.ShouldBindJSON(&req); err != nil { | |||
| syncErrorResponse(c, http.StatusBadRequest, "invalid request body") | |||
| return | |||
| } | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: received %d records", len(req.Records)) | |||
| results := make([]region_sync.DeductResult, 0, len(req.Records)) | |||
| // 使用事务保证批量扣费的原子性(含幂等性检查和日志记录) | |||
| err := model.DB.Transaction(func(tx *gorm.DB) error { | |||
| for _, record := range req.Records { | |||
| result := processDeductRecordWithTx(c, tx, record) | |||
| results = append(results, result) | |||
| // 在事务内记录同步日志,确保幂等性检查能查到 | |||
| if result.Success { | |||
| syncLog := &model.QuotaSyncLog{ | |||
| UserId: record.UserId, | |||
| SyncType: model.SyncTypeBatchSync, | |||
| Direction: model.SyncDirectionOvToCn, | |||
| Status: model.SyncStatusSuccess, | |||
| RequestId: record.RequestId, | |||
| ChangeAmount: record.Quota, | |||
| AfterQuota: result.RemainingQuota, | |||
| CreatedAt: time.Now().Unix(), | |||
| } | |||
| if err := tx.Create(syncLog).Error; err != nil { | |||
| return err | |||
| } | |||
| } | |||
| } | |||
| return nil | |||
| }) | |||
| if err != nil { | |||
| logger.LogError(c.Request.Context(), "batch deduct transaction failed: "+err.Error()) | |||
| syncErrorResponse(c, http.StatusInternalServerError, "batch deduct failed") | |||
| return | |||
| } | |||
| successCount := 0 | |||
| for _, r := range results { | |||
| if r.Success { | |||
| successCount++ | |||
| } | |||
| } | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] BatchDeductQuota: completed, success=%d, total=%d", successCount, len(req.Records)) | |||
| c.JSON(http.StatusOK, gin.H{"success": true, "results": results}) | |||
| } | |||
| // processDeductRecordWithTx 在事务中处理单条扣费记录 | |||
| func processDeductRecordWithTx(c *gin.Context, tx *gorm.DB, record region_sync.BatchDeductRecord) region_sync.DeductResult { | |||
| result := region_sync.DeductResult{UserId: record.UserId} | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: userId=%d, requestId=%s, quota=%d", record.UserId, record.RequestId, record.Quota) | |||
| // 幂等性检查:通过 quota_sync_logs 检查 request_id 是否已在 master 端处理过 | |||
| var existingLog model.QuotaSyncLog | |||
| if err := tx.Where("request_id = ? AND sync_type = ? AND status = ?", | |||
| record.RequestId, model.SyncTypeBatchSync, model.SyncStatusSuccess).First(&existingLog).Error; err == nil { | |||
| // 该 request_id 已处理过,直接返回成功(不重复扣费) | |||
| result.Success = true | |||
| result.Message = "already processed" | |||
| result.DeductedQuota = existingLog.ChangeAmount | |||
| // 查询当前余额 | |||
| var user model.User | |||
| if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil { | |||
| result.RemainingQuota = user.Quota | |||
| } | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: idempotent hit, userId=%d, requestId=%s, deductedQuota=%d", record.UserId, record.RequestId, result.DeductedQuota) | |||
| return result | |||
| } | |||
| // 原子扣费:使用 SQL 条件更新,确保余额充足时才扣减,避免并发覆盖 | |||
| dbResult := tx.Model(&model.User{}). | |||
| Where("id = ? AND quota >= ?", record.UserId, record.Quota). | |||
| Update("quota", gorm.Expr("quota - ?", record.Quota)) | |||
| if dbResult.Error != nil { | |||
| result.Error = "failed to deduct quota" | |||
| return result | |||
| } | |||
| if dbResult.RowsAffected == 0 { | |||
| // 没有匹配行,可能是用户不存在或余额不足 | |||
| var user model.User | |||
| if err := tx.Where("id = ?", record.UserId).First(&user).Error; err != nil { | |||
| result.Error = "user not found" | |||
| } else { | |||
| result.Error = "insufficient quota" | |||
| result.RemainingQuota = user.Quota | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: insufficient quota, userId=%d, quota=%d, need=%d", record.UserId, user.Quota, record.Quota) | |||
| } | |||
| return result | |||
| } | |||
| // 查询扣费后的余额 | |||
| var user model.User | |||
| if err := tx.Where("id = ?", record.UserId).First(&user).Error; err == nil { | |||
| result.RemainingQuota = user.Quota | |||
| } | |||
| result.Success = true | |||
| result.DeductedQuota = record.Quota | |||
| logger.LogDebug(c.Request.Context(), "[RegionSync] processDeductRecord: success, userId=%d, deducted=%d, remaining=%d", record.UserId, record.Quota, result.RemainingQuota) | |||
| return result | |||
| } | |||
| // recordSyncLog 记录同步日志 | |||
| func recordSyncLog(c *gin.Context, userId, remoteUserId int, syncType, requestId string, masterQuota, afterQuota int, direction string) { | |||
| syncLog := &model.QuotaSyncLog{ | |||
| UserId: userId, | |||
| RemoteUserId: remoteUserId, | |||
| SyncType: syncType, | |||
| Direction: direction, | |||
| Status: model.SyncStatusSuccess, | |||
| RequestId: requestId, | |||
| MasterQuota: masterQuota, | |||
| AfterQuota: afterQuota, | |||
| CreatedAt: time.Now().Unix(), | |||
| } | |||
| if err := model.CreateSyncLog(syncLog); err != nil { | |||
| logger.LogError(c.Request.Context(), "failed to create sync log: "+err.Error()) | |||
| } | |||
| } | |||
| @@ -0,0 +1,466 @@ | |||
| package controller | |||
| import ( | |||
| "bytes" | |||
| "encoding/json" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service/region_sync" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func TestReceiveSyncedUserCreate_Success(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| // 使用测试专用的认证中间件,绕过系统设置检查 | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate) | |||
| origDB := model.DB | |||
| defer func() { | |||
| model.DB = origDB | |||
| }() | |||
| // 设置测试数据库 | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| req := region_sync.SyncUserRequest{ | |||
| Username: "testuser", | |||
| Email: "test@example.com", | |||
| PasswordHash: "hashedpassword", | |||
| DisplayName: "Test User", | |||
| Quota: 10000, | |||
| RemoteUserId: 10000001, | |||
| Group: "default", | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp map[string]interface{} | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp["success"].(bool)) | |||
| // 验证用户已创建 | |||
| var user model.User | |||
| err := db.Where("username = ?", "testuser").First(&user).Error | |||
| require.NoError(t, err) | |||
| assert.Equal(t, "testuser", user.Username) | |||
| assert.Equal(t, 10000001, user.RemoteUserId) | |||
| assert.Equal(t, "synced", user.Source) | |||
| } | |||
| func TestReceiveQuotaUpdate_Success(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate) | |||
| origDB := model.DB | |||
| defer func() { | |||
| model.DB = origDB | |||
| }() | |||
| // 设置测试数据库 | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| // 创建测试用户 | |||
| user := model.User{ | |||
| Username: "testuser", | |||
| RemoteUserId: 10000001, | |||
| Source: "synced", | |||
| Quota: 5000, | |||
| SyncedQuota: 5000, | |||
| AffCode: "aff-qu-1", | |||
| } | |||
| db.Create(&user) | |||
| req := region_sync.UpdateQuotaRequest{ | |||
| RemoteUserId: 10000001, | |||
| Quota: 10000, | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp map[string]interface{} | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp["success"].(bool)) | |||
| // 验证余额已更新 | |||
| var updatedUser model.User | |||
| err := db.Where("remote_user_id = ?", 10000001).First(&updatedUser).Error | |||
| require.NoError(t, err) | |||
| assert.Equal(t, 10000, updatedUser.SyncedQuota) | |||
| } | |||
| func TestQueryUserQuota_Success(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/query", QueryUserQuota) | |||
| origDB := model.DB | |||
| defer func() { | |||
| model.DB = origDB | |||
| }() | |||
| // 设置测试数据库 | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| // 创建测试用户 | |||
| user := model.User{ | |||
| Username: "testuser", | |||
| Quota: 15000, | |||
| AffCode: "aff-qq-1", | |||
| } | |||
| db.Create(&user) | |||
| req := region_sync.QueryQuotaRequest{ | |||
| UserId: user.Id, | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/query", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp region_sync.QueryQuotaResponse | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp.Success) | |||
| assert.Equal(t, 15000, resp.Quota) | |||
| } | |||
| func TestBatchDeductQuota_Success(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) | |||
| origDB := model.DB | |||
| defer func() { | |||
| model.DB = origDB | |||
| }() | |||
| // 设置测试数据库 | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| // 创建测试用户 | |||
| user1 := model.User{Username: "user1", Quota: 1000, AffCode: "aff-bd-1"} | |||
| user2 := model.User{Username: "user2", Quota: 500, AffCode: "aff-bd-2"} | |||
| db.Create(&user1) | |||
| db.Create(&user2) | |||
| req := region_sync.BatchDeductRequest{ | |||
| Records: []region_sync.BatchDeductRecord{ | |||
| {UserId: user1.Id, RequestId: "req-1", Quota: 100}, | |||
| {UserId: user2.Id, RequestId: "req-2", Quota: 50}, | |||
| }, | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp region_sync.BatchDeductResponse | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp.Success) | |||
| assert.Len(t, resp.Results, 2) | |||
| // 验证用户余额已更新 | |||
| var updatedUser1 model.User | |||
| db.First(&updatedUser1, user1.Id) | |||
| assert.Equal(t, 900, updatedUser1.Quota) | |||
| var updatedUser2 model.User | |||
| db.First(&updatedUser2, user2.Id) | |||
| assert.Equal(t, 450, updatedUser2.Quota) | |||
| } | |||
| func TestBatchDeductQuota_InsufficientQuota(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) | |||
| origDB := model.DB | |||
| defer func() { | |||
| model.DB = origDB | |||
| }() | |||
| // 设置测试数据库 | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| // 创建余额不足的用户 | |||
| user := model.User{Username: "pooruser", Quota: 10, AffCode: "aff-poor-1"} | |||
| db.Create(&user) | |||
| req := region_sync.BatchDeductRequest{ | |||
| Records: []region_sync.BatchDeductRecord{ | |||
| {UserId: user.Id, RequestId: "req-1", Quota: 100}, | |||
| }, | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp region_sync.BatchDeductResponse | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp.Success) | |||
| assert.Len(t, resp.Results, 1) | |||
| assert.False(t, resp.Results[0].Success) | |||
| assert.Contains(t, resp.Results[0].Error, "insufficient quota") | |||
| } | |||
| // testSyncAuthMiddleware 测试专用的认证中间件,绕过系统设置 | |||
| func testSyncAuthMiddleware() gin.HandlerFunc { | |||
| return func(c *gin.Context) { | |||
| c.Set("sync_node", "test-node") | |||
| c.Next() | |||
| } | |||
| } | |||
| // setupSyncTestDB 设置同步测试数据库 | |||
| func setupSyncTestDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| sqlDB, _ := db.DB() | |||
| sqlDB.SetMaxOpenConns(1) | |||
| err = db.AutoMigrate(&model.User{}, &model.QuotaSyncLog{}) | |||
| require.NoError(t, err) | |||
| return db | |||
| } | |||
| // --------------------------------------------------------------------------- | |||
| // 阶段 4:Controller 边界条件测试 | |||
| // --------------------------------------------------------------------------- | |||
| func TestReceiveSyncedUserCreate_DuplicateUser(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate) | |||
| origDB := model.DB | |||
| defer func() { model.DB = origDB }() | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| // 先创建一个用户 | |||
| db.Create(&model.User{ | |||
| Username: "dupuser", RemoteUserId: 10000001, | |||
| Source: common.UserSourceSynced, Password: "hash", AffCode: "aff-dup-1", | |||
| }) | |||
| req := region_sync.SyncUserRequest{ | |||
| Username: "dupuser", | |||
| RemoteUserId: 10000001, | |||
| PasswordHash: "newhash", | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp map[string]interface{} | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp["success"].(bool)) | |||
| assert.Contains(t, resp["message"], "already exists") | |||
| } | |||
| func TestReceiveSyncedUserCreate_InvalidBody(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", strings.NewReader("not json")) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusBadRequest, w.Code) | |||
| } | |||
| func TestReceiveQuotaUpdate_UserNotFound(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate) | |||
| origDB := model.DB | |||
| defer func() { model.DB = origDB }() | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| req := region_sync.UpdateQuotaRequest{ | |||
| RemoteUserId: 99999999, | |||
| Quota: 5000, | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusNotFound, w.Code) | |||
| } | |||
| func TestReceiveQuotaUpdate_InvalidBody(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", strings.NewReader("invalid")) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusBadRequest, w.Code) | |||
| } | |||
| func TestQueryUserQuota_UserNotFound(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/query", QueryUserQuota) | |||
| origDB := model.DB | |||
| defer func() { model.DB = origDB }() | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| req := region_sync.QueryQuotaRequest{UserId: 99999} | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/query", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusNotFound, w.Code) | |||
| } | |||
| func TestBatchDeductQuota_EmptyRecords(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) | |||
| origDB := model.DB | |||
| defer func() { model.DB = origDB }() | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| req := region_sync.BatchDeductRequest{Records: []region_sync.BatchDeductRecord{}} | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp region_sync.BatchDeductResponse | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp.Success) | |||
| assert.Len(t, resp.Results, 0) | |||
| } | |||
| func TestBatchDeductQuota_UserNotFound(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) | |||
| origDB := model.DB | |||
| defer func() { model.DB = origDB }() | |||
| db := setupSyncTestDB(t) | |||
| model.DB = db | |||
| req := region_sync.BatchDeductRequest{ | |||
| Records: []region_sync.BatchDeductRecord{ | |||
| {UserId: 99999, RequestId: "req-nf", Quota: 100}, | |||
| }, | |||
| } | |||
| reqBody, _ := json.Marshal(req) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody)) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| var resp region_sync.BatchDeductResponse | |||
| json.Unmarshal(w.Body.Bytes(), &resp) | |||
| assert.True(t, resp.Success) | |||
| assert.Len(t, resp.Results, 1) | |||
| assert.False(t, resp.Results[0].Success) | |||
| assert.Contains(t, resp.Results[0].Error, "user not found") | |||
| } | |||
| func TestBatchDeductQuota_InvalidBody(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(testSyncAuthMiddleware()) | |||
| router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota) | |||
| reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", strings.NewReader("bad")) | |||
| reqHTTP.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, reqHTTP) | |||
| assert.Equal(t, http.StatusBadRequest, w.Code) | |||
| } | |||
| @@ -159,6 +159,8 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { | |||
| if priceData.FreeModel { | |||
| logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName)) | |||
| } else { | |||
| common.SysLog(fmt.Sprintf("[RegionSync] relay PreConsumeBilling: userId=%d, QuotaToPreConsume=%d, FreeModel=%v, UsePrice=%v, ModelRatio=%.4f, GroupRatio=%.4f, ModelPrice=%.4f, model=%s", | |||
| relayInfo.UserId, priceData.QuotaToPreConsume, priceData.FreeModel, priceData.UsePrice, priceData.ModelRatio, priceData.GroupRatioInfo.GroupRatio, priceData.ModelPrice, relayInfo.OriginModelName)) | |||
| newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo) | |||
| if newAPIError != nil { | |||
| return | |||
| @@ -16,6 +16,7 @@ import ( | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/service/region_sync" | |||
| "github.com/QuantumNous/new-api/setting" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| @@ -184,12 +185,8 @@ func Register(c *gin.Context) { | |||
| return | |||
| } | |||
| // 获取插入后的用户ID | |||
| var insertedUser model.User | |||
| if err := model.DB.Where("username = ?", cleanUser.Username).First(&insertedUser).Error; err != nil { | |||
| common.ApiErrorI18n(c, i18n.MsgUserRegisterFailed) | |||
| return | |||
| } | |||
| // 同步用户到海外节点(异步执行,不阻塞注册流程) | |||
| region_sync.PushUserCreateToSlave(&cleanUser) | |||
| // 生成默认令牌 | |||
| if constant.GenerateDefaultToken { | |||
| key, err := common.GenerateKey() | |||
| @@ -200,7 +197,7 @@ func Register(c *gin.Context) { | |||
| } | |||
| // 生成默认令牌 | |||
| token := model.Token{ | |||
| UserId: insertedUser.Id, // 使用插入后的用户ID | |||
| UserId: cleanUser.Id, // GORM Create 后已填充 ID | |||
| Name: cleanUser.Username + "的初始令牌", | |||
| Key: key, | |||
| CreatedTime: common.GetTimestamp(), | |||
| @@ -383,6 +380,11 @@ func GetSelf(c *gin.Context) { | |||
| userSetting := user.GetSetting() | |||
| // 构建响应数据,包含用户信息和权限 | |||
| // Slave 节点使用 SyncedQuota 作为用户额度 | |||
| quota := user.Quota | |||
| if user.IsSyncedUser() { | |||
| quota = user.SyncedQuota | |||
| } | |||
| responseData := map[string]interface{}{ | |||
| "id": user.Id, | |||
| "username": user.Username, | |||
| @@ -396,7 +398,7 @@ func GetSelf(c *gin.Context) { | |||
| "wechat_id": user.WeChatId, | |||
| "telegram_id": user.TelegramId, | |||
| "group": user.Group, | |||
| "quota": user.Quota, | |||
| "quota": quota, | |||
| "used_quota": user.UsedQuota, | |||
| "request_count": user.RequestCount, | |||
| "aff_code": user.AffCode, | |||
| @@ -22,8 +22,10 @@ import ( | |||
| "github.com/QuantumNous/new-api/relay" | |||
| "github.com/QuantumNous/new-api/router" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/service/region_sync" | |||
| _ "github.com/QuantumNous/new-api/setting/performance_setting" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/bytedance/gopkg/util/gopool" | |||
| "github.com/gin-contrib/sessions" | |||
| @@ -174,7 +176,12 @@ func main() { | |||
| Secure: false, | |||
| SameSite: http.SameSiteStrictMode, | |||
| }) | |||
| server.Use(sessions.Sessions("session", store)) | |||
| // 支持通过环境变量自定义 session 名称,用于多节点部署时避免 cookie 冲突 | |||
| sessionName := os.Getenv("SESSION_NAME") | |||
| if sessionName == "" { | |||
| sessionName = "session" | |||
| } | |||
| server.Use(sessions.Sessions(sessionName, store)) | |||
| InjectUmamiAnalytics() | |||
| InjectGoogleAnalytics() | |||
| @@ -309,5 +316,19 @@ func InitResources() error { | |||
| // Don't return error, custom OAuth is not critical | |||
| } | |||
| // 注册余额更新回调,Master 节点更新用户余额后推送到 Slave 节点 | |||
| model.SetQuotaUpdateCallback(func(userId int, quota int) { | |||
| if system_setting.GetRegionSyncSettings().IsMaster { | |||
| region_sync.PushQuotaUpdateToSlave(userId, quota) | |||
| } | |||
| }) | |||
| // Slave 节点启动后台同步任务(批量扣费同步 + 余额定时拉取 + 清理) | |||
| syncSettings := system_setting.GetRegionSyncSettings() | |||
| if syncSettings.Enabled && !syncSettings.IsMaster { | |||
| syncManager := region_sync.NewSyncManager() | |||
| syncManager.StartSyncWorkers() | |||
| } | |||
| return nil | |||
| } | |||
| @@ -22,7 +22,7 @@ DOCKER_TAG := $(DOCKER_REGISTRY):$(BUILD_TIME)-$(BRANCH_NAME) | |||
| docker-build: | |||
| @echo "Building Docker image with tag: $(DOCKER_TAG)" | |||
| docker build -t $(DOCKER_TAG) . | |||
| docker push $(DOCKER_TAG) | |||
| #docker push $(DOCKER_TAG) | |||
| docker-push: docker-build | |||
| @echo "Pushing Docker image: $(DOCKER_TAG)" | |||
| @@ -0,0 +1,96 @@ | |||
| package middleware | |||
| import ( | |||
| "crypto/subtle" | |||
| "net/http" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| // SyncAuth 同步 API 认证中间件 | |||
| // 通过 X-Sync-API-Key 请求头验证请求的合法性 | |||
| func SyncAuth() gin.HandlerFunc { | |||
| return func(c *gin.Context) { | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| // 如果未启用区域同步,拒绝请求 | |||
| if !settings.Enabled { | |||
| c.JSON(http.StatusForbidden, gin.H{ | |||
| "success": false, | |||
| "message": "region sync is not enabled", | |||
| }) | |||
| c.Abort() | |||
| return | |||
| } | |||
| // 检查 X-Sync-API-Key 请求头 | |||
| apiKey := c.GetHeader("X-Sync-API-Key") | |||
| if apiKey == "" { | |||
| c.JSON(http.StatusUnauthorized, gin.H{ | |||
| "success": false, | |||
| "message": "missing sync API key", | |||
| }) | |||
| c.Abort() | |||
| return | |||
| } | |||
| // 验证 API Key | |||
| if !validateSyncAPIKey(apiKey, settings.SyncApiKey) { | |||
| c.JSON(http.StatusUnauthorized, gin.H{ | |||
| "success": false, | |||
| "message": "invalid sync API key", | |||
| }) | |||
| c.Abort() | |||
| return | |||
| } | |||
| // 验证来源节点(从 X-Sync-Node 获取) | |||
| node := c.GetHeader("X-Sync-Node") | |||
| if node == "" { | |||
| c.JSON(http.StatusBadRequest, gin.H{ | |||
| "success": false, | |||
| "message": "missing sync node identifier", | |||
| }) | |||
| c.Abort() | |||
| return | |||
| } | |||
| // 将节点信息存储在上下文中 | |||
| c.Set("sync_node", node) | |||
| c.Next() | |||
| } | |||
| } | |||
| // validateSyncAPIKey 验证同步 API Key(使用常量时间比较防止 timing attack) | |||
| func validateSyncAPIKey(requestedKey, configuredKey string) bool { | |||
| if configuredKey == "" { | |||
| return false | |||
| } | |||
| // 支持 master 配置多个 slave 的 key,以逗号分隔 | |||
| if strings.Contains(configuredKey, ",") { | |||
| keys := strings.Split(configuredKey, ",") | |||
| for _, key := range keys { | |||
| if subtle.ConstantTimeCompare([]byte(strings.TrimSpace(key)), []byte(requestedKey)) == 1 { | |||
| return true | |||
| } | |||
| } | |||
| return false | |||
| } | |||
| return subtle.ConstantTimeCompare([]byte(strings.TrimSpace(configuredKey)), []byte(requestedKey)) == 1 | |||
| } | |||
| // GetSyncNode 从上下文获取同步节点标识 | |||
| func GetSyncNode(c *gin.Context) string { | |||
| if node, exists := c.Get("sync_node"); exists { | |||
| if nodeStr, ok := node.(string); ok { | |||
| return nodeStr | |||
| } | |||
| } | |||
| return "" | |||
| } | |||
| @@ -0,0 +1,149 @@ | |||
| package middleware | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/assert" | |||
| ) | |||
| func TestSyncAuth_Disabled(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(SyncAuth()) | |||
| router.GET("/test", func(c *gin.Context) { | |||
| c.JSON(http.StatusOK, gin.H{"success": true}) | |||
| }) | |||
| // 未启用区域同步 | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| settings.Enabled = false | |||
| req, _ := http.NewRequest("GET", "/test", nil) | |||
| req.Header.Set("X-Sync-API-Key", "test-key") | |||
| req.Header.Set("X-Sync-Node", "slave-1") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, req) | |||
| assert.Equal(t, http.StatusForbidden, w.Code) | |||
| assert.Contains(t, w.Body.String(), "region sync is not enabled") | |||
| } | |||
| func TestSyncAuth_MissingAPIKey(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(SyncAuth()) | |||
| router.GET("/test", func(c *gin.Context) { | |||
| c.JSON(http.StatusOK, gin.H{"success": true}) | |||
| }) | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| settings.Enabled = true | |||
| settings.SyncApiKey = "test-key" | |||
| req, _ := http.NewRequest("GET", "/test", nil) | |||
| req.Header.Set("X-Sync-Node", "slave-1") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, req) | |||
| assert.Equal(t, http.StatusUnauthorized, w.Code) | |||
| assert.Contains(t, w.Body.String(), "missing sync API key") | |||
| } | |||
| func TestSyncAuth_InvalidAPIKey(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(SyncAuth()) | |||
| router.GET("/test", func(c *gin.Context) { | |||
| c.JSON(http.StatusOK, gin.H{"success": true}) | |||
| }) | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| settings.Enabled = true | |||
| settings.SyncApiKey = "correct-key" | |||
| req, _ := http.NewRequest("GET", "/test", nil) | |||
| req.Header.Set("X-Sync-API-Key", "wrong-key") | |||
| req.Header.Set("X-Sync-Node", "slave-1") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, req) | |||
| assert.Equal(t, http.StatusUnauthorized, w.Code) | |||
| assert.Contains(t, w.Body.String(), "invalid sync API key") | |||
| } | |||
| func TestSyncAuth_MissingNode(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(SyncAuth()) | |||
| router.GET("/test", func(c *gin.Context) { | |||
| c.JSON(http.StatusOK, gin.H{"success": true}) | |||
| }) | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| settings.Enabled = true | |||
| settings.SyncApiKey = "test-key" | |||
| req, _ := http.NewRequest("GET", "/test", nil) | |||
| req.Header.Set("X-Sync-API-Key", "test-key") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, req) | |||
| assert.Equal(t, http.StatusBadRequest, w.Code) | |||
| assert.Contains(t, w.Body.String(), "missing sync node identifier") | |||
| } | |||
| func TestSyncAuth_Valid(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(SyncAuth()) | |||
| router.GET("/test", func(c *gin.Context) { | |||
| node := GetSyncNode(c) | |||
| c.JSON(http.StatusOK, gin.H{ | |||
| "success": true, | |||
| "node": node, | |||
| }) | |||
| }) | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| settings.Enabled = true | |||
| settings.SyncApiKey = "test-key" | |||
| req, _ := http.NewRequest("GET", "/test", nil) | |||
| req.Header.Set("X-Sync-API-Key", "test-key") | |||
| req.Header.Set("X-Sync-Node", "slave-1") | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, req) | |||
| assert.Equal(t, http.StatusOK, w.Code) | |||
| assert.Contains(t, w.Body.String(), "slave-1") | |||
| } | |||
| func TestValidateSyncAPIKey_Single(t *testing.T) { | |||
| assert.True(t, validateSyncAPIKey("test-key", "test-key")) | |||
| assert.False(t, validateSyncAPIKey("test-key", "different-key")) | |||
| assert.False(t, validateSyncAPIKey("", "")) | |||
| } | |||
| func TestValidateSyncAPIKey_Multiple(t *testing.T) { | |||
| // 支持逗号分隔的多个 key | |||
| assert.True(t, validateSyncAPIKey("key1", "key1,key2,key3")) | |||
| assert.True(t, validateSyncAPIKey("key2", "key1,key2,key3")) | |||
| assert.True(t, validateSyncAPIKey("key3", "key1,key2,key3")) | |||
| assert.False(t, validateSyncAPIKey("key4", "key1,key2,key3")) | |||
| } | |||
| func TestValidateSyncAPIKey_TrimSpace(t *testing.T) { | |||
| // 去除前后空格 | |||
| assert.True(t, validateSyncAPIKey("test-key", " test-key ")) | |||
| assert.True(t, validateSyncAPIKey("test-key", " test-key")) | |||
| assert.True(t, validateSyncAPIKey("test-key", "test-key ")) | |||
| } | |||
| @@ -280,6 +280,8 @@ func migrateDB() error { | |||
| &UserOAuthBinding{}, | |||
| &ChannelPricing{}, | |||
| &PricingTag{}, | |||
| &PendingSyncRecord{}, | |||
| &QuotaSyncLog{}, | |||
| ) | |||
| if err != nil { | |||
| return err | |||
| @@ -332,6 +334,8 @@ func migrateDBFast() error { | |||
| {&UserOAuthBinding{}, "UserOAuthBinding"}, | |||
| {&ChannelPricing{}, "ChannelPricing"}, | |||
| {&PricingTag{}, "PricingTag"}, | |||
| {&PendingSyncRecord{}, "PendingSyncRecord"}, | |||
| {&QuotaSyncLog{}, "QuotaSyncLog"}, | |||
| } | |||
| // 动态计算migration数量,确保errChan缓冲区足够大 | |||
| errChan := make(chan error, len(migrations)) | |||
| @@ -0,0 +1,126 @@ | |||
| package model | |||
| import ( | |||
| "time" | |||
| "gorm.io/gorm" | |||
| ) | |||
| const ( | |||
| PendingSyncStatusPending = "pending" | |||
| PendingSyncStatusSynced = "synced" | |||
| PendingSyncStatusFailed = "failed" | |||
| PendingSyncStatusArchived = "archived" | |||
| ) | |||
| // PendingSyncRecord 待同步消费记录 | |||
| type PendingSyncRecord struct { | |||
| Id int `json:"id" gorm:"primaryKey"` | |||
| UserId int `json:"user_id" gorm:"index:idx_user_status"` | |||
| RemoteUserId int `json:"remote_user_id" gorm:"index"` | |||
| RequestId string `json:"request_id" gorm:"type:varchar(128);uniqueIndex"` | |||
| Quota int `json:"quota"` | |||
| EstimatedQuota int `json:"estimated_quota"` | |||
| Status string `json:"status" gorm:"default:'pending';index:idx_user_status;index:idx_status_created"` | |||
| CreatedAt int64 `json:"created_at" gorm:"index:idx_status_created"` | |||
| SyncedAt int64 `json:"synced_at"` | |||
| RetryCount int `json:"retry_count" gorm:"default:0"` | |||
| ErrorMsg string `json:"error_msg"` | |||
| } | |||
| func (PendingSyncRecord) TableName() string { | |||
| return "pending_sync_records" | |||
| } | |||
| // GetPendingSyncQuota 获取用户待同步金额总和 | |||
| func GetPendingSyncQuota(userId int) int { | |||
| var total int64 | |||
| DB.Model(&PendingSyncRecord{}). | |||
| Where("user_id = ? AND status IN ?", userId, []string{PendingSyncStatusPending, PendingSyncStatusFailed}). | |||
| Select("COALESCE(SUM(quota), 0)"). | |||
| Scan(&total) | |||
| return int(total) | |||
| } | |||
| // CreatePendingSyncRecord 创建待同步记录 | |||
| func CreatePendingSyncRecord(userId, remoteUserId int, requestId string, quota, estimatedQuota int) error { | |||
| record := PendingSyncRecord{ | |||
| UserId: userId, | |||
| RemoteUserId: remoteUserId, | |||
| RequestId: requestId, | |||
| Quota: quota, | |||
| EstimatedQuota: estimatedQuota, | |||
| Status: PendingSyncStatusPending, | |||
| CreatedAt: time.Now().Unix(), | |||
| } | |||
| return DB.Create(&record).Error | |||
| } | |||
| // GetPendingRecordsForSync 获取待同步记录 | |||
| func GetPendingRecordsForSync(limit, maxRetry int) []PendingSyncRecord { | |||
| var records []PendingSyncRecord | |||
| DB.Where("status IN ? AND retry_count < ?", | |||
| []string{PendingSyncStatusPending, PendingSyncStatusFailed}, | |||
| maxRetry). | |||
| Limit(limit). | |||
| Order("created_at asc"). | |||
| Find(&records) | |||
| return records | |||
| } | |||
| // MarkRecordSynced 标记记录为已同步 | |||
| func MarkRecordSynced(recordId int) error { | |||
| return DB.Model(&PendingSyncRecord{}).Where("id = ?", recordId).Updates(map[string]interface{}{ | |||
| "status": PendingSyncStatusSynced, | |||
| "synced_at": time.Now().Unix(), | |||
| }).Error | |||
| } | |||
| // MarkRecordFailed 标记记录为失败 | |||
| func MarkRecordFailed(recordId int, errMsg string) error { | |||
| return DB.Model(&PendingSyncRecord{}).Where("id = ?", recordId).Updates(map[string]interface{}{ | |||
| "status": PendingSyncStatusFailed, | |||
| "retry_count": gorm.Expr("retry_count + 1"), | |||
| "error_msg": errMsg, | |||
| }).Error | |||
| } | |||
| // UpdatePendingSyncRecordQuota 更新待同步记录的实际扣费额度,返回影响的行数 | |||
| func UpdatePendingSyncRecordQuota(requestId string, actualQuota int) int64 { | |||
| result := DB.Model(&PendingSyncRecord{}). | |||
| Where("request_id = ? AND status = ?", requestId, PendingSyncStatusPending). | |||
| Update("quota", actualQuota) | |||
| if result.Error != nil { | |||
| return 0 | |||
| } | |||
| return result.RowsAffected | |||
| } | |||
| // CleanupOldSyncedRecords 清理超过指定天数的已同步记录 | |||
| func CleanupOldSyncedRecords(days int) int64 { | |||
| result := DB.Where("status = ? AND synced_at < ?", | |||
| PendingSyncStatusSynced, | |||
| time.Now().Unix()-int64(days*86400)). | |||
| Delete(&PendingSyncRecord{}) | |||
| return result.RowsAffected | |||
| } | |||
| // GetPendingQuotaByUser 获取每个用户的待同步金额 | |||
| func GetPendingQuotaByUser() map[int]int { | |||
| type result struct { | |||
| UserId int | |||
| Total int | |||
| } | |||
| var results []result | |||
| DB.Model(&PendingSyncRecord{}). | |||
| Select("user_id, COALESCE(SUM(quota), 0) as total"). | |||
| Where("status IN ?", []string{PendingSyncStatusPending, PendingSyncStatusFailed}). | |||
| Group("user_id"). | |||
| Scan(&results) | |||
| m := make(map[int]int) | |||
| for _, r := range results { | |||
| m[r.UserId] = r.Total | |||
| } | |||
| return m | |||
| } | |||
| @@ -0,0 +1,55 @@ | |||
| package model | |||
| import "time" | |||
| const ( | |||
| SyncTypeUserCreate = "user_create" | |||
| SyncTypeQuotaChange = "quota_change" | |||
| SyncTypePreConsumeQuery = "pre_consume_query" | |||
| SyncTypeBatchSync = "batch_sync" | |||
| SyncTypeManualSync = "manual_sync" | |||
| SyncTypeStartupSync = "startup_sync" | |||
| SyncTypeSettle = "settle" | |||
| ) | |||
| const ( | |||
| SyncDirectionCnToOv = "cn_to_ov" | |||
| SyncDirectionOvToCn = "ov_to_cn" | |||
| ) | |||
| const ( | |||
| SyncStatusSuccess = "success" | |||
| SyncStatusFailed = "failed" | |||
| ) | |||
| // QuotaSyncLog 余额同步日志 | |||
| type QuotaSyncLog struct { | |||
| Id int `json:"id" gorm:"primaryKey"` | |||
| UserId int `json:"user_id" gorm:"index:idx_user_id"` | |||
| RemoteUserId int `json:"remote_user_id" gorm:"index"` | |||
| SyncType string `json:"sync_type" gorm:"size:32;index:idx_sync_type"` | |||
| BeforeQuota int `json:"before_quota"` | |||
| AfterQuota int `json:"after_quota"` | |||
| ChangeAmount int `json:"change_amount"` | |||
| MasterQuota int `json:"master_quota"` | |||
| PendingQuota int `json:"pending_quota"` | |||
| RequestId string `json:"request_id"` | |||
| Model string `json:"model"` | |||
| Direction string `json:"direction" gorm:"size:16"` | |||
| Status string `json:"status" gorm:"size:16;default:'success';index:idx_status"` | |||
| ErrorMsg string `json:"error_msg"` | |||
| CreatedAt int64 `json:"created_at" gorm:"index:idx_created_at"` | |||
| SyncedAt int64 `json:"synced_at"` | |||
| } | |||
| func (QuotaSyncLog) TableName() string { | |||
| return "quota_sync_logs" | |||
| } | |||
| // CreateSyncLog 创建同步日志 | |||
| func CreateSyncLog(log *QuotaSyncLog) error { | |||
| if log.CreatedAt == 0 { | |||
| log.CreatedAt = time.Now().Unix() | |||
| } | |||
| return DB.Create(log).Error | |||
| } | |||
| @@ -154,6 +154,21 @@ func Redeem(key string, userId int) (quota int, err error) { | |||
| common.SysError("redemption failed: " + err.Error()) | |||
| return 0, ErrRedeemFailed | |||
| } | |||
| // 触发余额更新回调(在事务外) | |||
| if quotaUpdateCallback != nil { | |||
| newQuota, _ := GetUserQuota(userId, true) | |||
| // 使用 defer/recover 保护回调 | |||
| func() { | |||
| defer func() { | |||
| if r := recover(); r != nil { | |||
| common.SysError(fmt.Sprintf("quota update callback panic in Redeem: %v", r)) | |||
| } | |||
| }() | |||
| quotaUpdateCallback(userId, newQuota) | |||
| }() | |||
| } | |||
| RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s,兑换码ID %d", logger.LogQuota(redemption.Quota), redemption.Id)) | |||
| return redemption.Quota, nil | |||
| } | |||
| @@ -0,0 +1,254 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupRedemptionDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| sqlDB, _ := db.DB() | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := DB | |||
| origLogDB := LOG_DB | |||
| DB = db | |||
| LOG_DB = db | |||
| common.UsingSQLite = true | |||
| common.RedisEnabled = false | |||
| require.NoError(t, db.AutoMigrate(&User{}, &Redemption{}, &Log{})) | |||
| t.Cleanup(func() { | |||
| DB = origDB | |||
| LOG_DB = origLogDB | |||
| sqlDB.Close() | |||
| }) | |||
| return db | |||
| } | |||
| func TestRedeem_Success(t *testing.T) { | |||
| db := setupRedemptionDB(t) | |||
| // 创建测试用户 | |||
| user := User{ | |||
| Id: 100, | |||
| Username: "redeem_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 创建兑换码 | |||
| redemption := Redemption{ | |||
| Id: 1, | |||
| UserId: 1, | |||
| Key: "test-key-123", | |||
| Status: common.RedemptionCodeStatusEnabled, | |||
| Quota: 50000, | |||
| CreatedTime: common.GetTimestamp(), | |||
| ExpiredTime: 0, // 永不过期 | |||
| } | |||
| require.NoError(t, db.Create(&redemption).Error) | |||
| // 设置回调 | |||
| callbackCalled := false | |||
| var receivedUserId int | |||
| var receivedQuota int | |||
| SetQuotaUpdateCallback(func(userId int, quota int) { | |||
| callbackCalled = true | |||
| receivedUserId = userId | |||
| receivedQuota = quota | |||
| }) | |||
| defer SetQuotaUpdateCallback(nil) | |||
| // 执行兑换 | |||
| quota, err := Redeem("test-key-123", 100) | |||
| require.NoError(t, err) | |||
| assert.Equal(t, 50000, quota) | |||
| // 验证回调被触发 | |||
| assert.True(t, callbackCalled) | |||
| assert.Equal(t, 100, receivedUserId) | |||
| assert.Equal(t, 150000, receivedQuota) // 原始 100000 + 兑换 50000 | |||
| // 验证用户额度已更新 | |||
| var updatedUser User | |||
| require.NoError(t, db.First(&updatedUser, 100).Error) | |||
| assert.Equal(t, 150000, updatedUser.Quota) | |||
| // 验证兑换码状态已更新 | |||
| var updatedRedemption Redemption | |||
| require.NoError(t, db.First(&updatedRedemption, 1).Error) | |||
| assert.Equal(t, common.RedemptionCodeStatusUsed, updatedRedemption.Status) | |||
| assert.Equal(t, 100, updatedRedemption.UsedUserId) | |||
| } | |||
| func TestRedeem_InvalidKey(t *testing.T) { | |||
| setupRedemptionDB(t) | |||
| // 测试无效的兑换码 | |||
| quota, err := Redeem("invalid-key", 100) | |||
| assert.Error(t, err) | |||
| assert.Equal(t, 0, quota) | |||
| } | |||
| func TestRedeem_AlreadyUsed(t *testing.T) { | |||
| db := setupRedemptionDB(t) | |||
| // 创建测试用户 | |||
| user := User{ | |||
| Id: 100, | |||
| Username: "redeem_used_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 创建已使用的兑换码 | |||
| redemption := Redemption{ | |||
| Id: 1, | |||
| UserId: 1, | |||
| Key: "used-key", | |||
| Status: common.RedemptionCodeStatusUsed, | |||
| Quota: 50000, | |||
| CreatedTime: common.GetTimestamp(), | |||
| RedeemedTime: common.GetTimestamp(), | |||
| UsedUserId: 99, | |||
| } | |||
| require.NoError(t, db.Create(&redemption).Error) | |||
| // 尝试兑换已使用的兑换码 | |||
| quota, err := Redeem("used-key", 100) | |||
| assert.Error(t, err) | |||
| assert.Equal(t, 0, quota) | |||
| // 错误会被包装为 ErrRedeemFailed,但底层错误消息会记录在日志中 | |||
| assert.Equal(t, ErrRedeemFailed, err) | |||
| } | |||
| func TestRedeem_Expired(t *testing.T) { | |||
| db := setupRedemptionDB(t) | |||
| // 创建测试用户 | |||
| user := User{ | |||
| Id: 100, | |||
| Username: "redeem_expired_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 创建已过期的兑换码 | |||
| redemption := Redemption{ | |||
| Id: 1, | |||
| UserId: 1, | |||
| Key: "expired-key", | |||
| Status: common.RedemptionCodeStatusEnabled, | |||
| Quota: 50000, | |||
| CreatedTime: common.GetTimestamp() - 86400, | |||
| ExpiredTime: common.GetTimestamp() - 3600, // 已过期 | |||
| } | |||
| require.NoError(t, db.Create(&redemption).Error) | |||
| // 尝试兑换已过期的兑换码 | |||
| quota, err := Redeem("expired-key", 100) | |||
| assert.Error(t, err) | |||
| assert.Equal(t, 0, quota) | |||
| // 错误会被包装为 ErrRedeemFailed | |||
| assert.Equal(t, ErrRedeemFailed, err) | |||
| } | |||
| func TestRedeem_CallbackPanic(t *testing.T) { | |||
| db := setupRedemptionDB(t) | |||
| // 创建测试用户 | |||
| user := User{ | |||
| Id: 100, | |||
| Username: "redeem_panic_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 创建兑换码 | |||
| redemption := Redemption{ | |||
| Id: 1, | |||
| UserId: 1, | |||
| Key: "panic-key", | |||
| Status: common.RedemptionCodeStatusEnabled, | |||
| Quota: 50000, | |||
| CreatedTime: common.GetTimestamp(), | |||
| } | |||
| require.NoError(t, db.Create(&redemption).Error) | |||
| // 设置一个会 panic 的回调 | |||
| SetQuotaUpdateCallback(func(userId int, quota int) { | |||
| panic("callback panic in redeem") | |||
| }) | |||
| defer SetQuotaUpdateCallback(nil) | |||
| // 执行兑换(回调 panic 不应影响兑换成功) | |||
| quota, err := Redeem("panic-key", 100) | |||
| require.NoError(t, err) | |||
| assert.Equal(t, 50000, quota) | |||
| // 验证用户额度已更新 | |||
| var updatedUser User | |||
| require.NoError(t, db.First(&updatedUser, 100).Error) | |||
| assert.Equal(t, 150000, updatedUser.Quota) | |||
| } | |||
| func TestRedeem_SyncedUser(t *testing.T) { | |||
| db := setupRedemptionDB(t) | |||
| // 创建同步用户(从 Master 节点同步过来的) | |||
| user := User{ | |||
| Id: 100, | |||
| Username: "synced_redeem_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| Source: common.UserSourceSynced, | |||
| RemoteUserId: 1000, | |||
| SyncedQuota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 创建兑换码 | |||
| redemption := Redemption{ | |||
| Id: 1, | |||
| UserId: 1, | |||
| Key: "synced-redeem-key", | |||
| Status: common.RedemptionCodeStatusEnabled, | |||
| Quota: 50000, | |||
| CreatedTime: common.GetTimestamp(), | |||
| } | |||
| require.NoError(t, db.Create(&redemption).Error) | |||
| // 设置回调 | |||
| callbackCalled := false | |||
| SetQuotaUpdateCallback(func(userId int, quota int) { | |||
| callbackCalled = true | |||
| }) | |||
| defer SetQuotaUpdateCallback(nil) | |||
| // 执行兑换 | |||
| quota, err := Redeem("synced-redeem-key", 100) | |||
| require.NoError(t, err) | |||
| assert.Equal(t, 50000, quota) | |||
| // 验证回调被触发(Master 节点需要同步到 Slave) | |||
| assert.True(t, callbackCalled) | |||
| // 验证用户额度已更新 | |||
| var updatedUser User | |||
| require.NoError(t, db.First(&updatedUser, 100).Error) | |||
| assert.Equal(t, 150000, updatedUser.Quota) | |||
| } | |||
| @@ -0,0 +1,538 @@ | |||
| package model | |||
| import ( | |||
| "fmt" | |||
| "testing" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupRegionSyncDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| sqlDB, _ := db.DB() | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := DB | |||
| origLogDB := LOG_DB | |||
| DB = db | |||
| LOG_DB = db | |||
| common.UsingSQLite = true | |||
| common.RedisEnabled = false | |||
| require.NoError(t, db.AutoMigrate(&User{})) | |||
| t.Cleanup(func() { | |||
| DB = origDB | |||
| LOG_DB = origLogDB | |||
| sqlDB.Close() | |||
| }) | |||
| return db | |||
| } | |||
| func TestUser_IsSyncedUser_Method(t *testing.T) { | |||
| tests := []struct { | |||
| name string | |||
| user User | |||
| want bool | |||
| }{ | |||
| {"synced valid", User{Source: "synced", Id: 100}, true}, | |||
| {"synced boundary", User{Source: "synced", Id: common.ForeignUserIDStart - 1}, true}, | |||
| {"synced id 0", User{Source: "synced", Id: 0}, false}, | |||
| {"synced too large", User{Source: "synced", Id: common.ForeignUserIDStart}, false}, | |||
| {"local user", User{Source: "local", Id: 100}, false}, | |||
| } | |||
| for _, tt := range tests { | |||
| t.Run(tt.name, func(t *testing.T) { | |||
| assert.Equal(t, tt.want, tt.user.IsSyncedUser()) | |||
| }) | |||
| } | |||
| } | |||
| func TestUser_IsLocalUser_Method(t *testing.T) { | |||
| tests := []struct { | |||
| name string | |||
| user User | |||
| want bool | |||
| }{ | |||
| {"local source", User{Source: "local", Id: common.ForeignUserIDStart}, true}, | |||
| {"foreign id", User{Source: "", Id: common.ForeignUserIDStart + 1}, true}, | |||
| {"synced source", User{Source: "synced", Id: 100}, false}, | |||
| } | |||
| for _, tt := range tests { | |||
| t.Run(tt.name, func(t *testing.T) { | |||
| assert.Equal(t, tt.want, tt.user.IsLocalUser()) | |||
| }) | |||
| } | |||
| } | |||
| func TestUser_SyncFields_CRUD(t *testing.T) { | |||
| db := setupRegionSyncDB(t) | |||
| user := User{ | |||
| Id: 100, Username: "synced_test", Password: "hashed_password", | |||
| Source: common.UserSourceSynced, RemoteUserId: 100, | |||
| SyncedQuota: 500000, LastSyncAt: 1700000000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| var loaded User | |||
| require.NoError(t, db.First(&loaded, 100).Error) | |||
| assert.Equal(t, common.UserSourceSynced, loaded.Source) | |||
| assert.Equal(t, 100, loaded.RemoteUserId) | |||
| assert.Equal(t, 500000, loaded.SyncedQuota) | |||
| assert.Equal(t, int64(1700000000), loaded.LastSyncAt) | |||
| } | |||
| func TestUser_DefaultSourceIsLocal(t *testing.T) { | |||
| db := setupRegionSyncDB(t) | |||
| user := User{ | |||
| Id: 101, Username: "default_source_test", Password: "hashed_password", | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| var loaded User | |||
| require.NoError(t, db.Where("username = ?", "default_source_test").First(&loaded).Error) | |||
| assert.Equal(t, common.UserSourceLocal, loaded.Source) | |||
| } | |||
| // PendingSyncRecord tests | |||
| func setupPendingSyncDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| sqlDB, _ := db.DB() | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := DB | |||
| origLogDB := LOG_DB | |||
| DB = db | |||
| LOG_DB = db | |||
| common.UsingSQLite = true | |||
| common.RedisEnabled = false | |||
| require.NoError(t, db.AutoMigrate(&PendingSyncRecord{})) | |||
| t.Cleanup(func() { | |||
| DB = origDB | |||
| LOG_DB = origLogDB | |||
| sqlDB.Close() | |||
| }) | |||
| return db | |||
| } | |||
| func TestPendingSyncRecord_Create(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| userId := 100 | |||
| remoteUserId := 200 | |||
| requestId := "req-test-001" | |||
| quota := 1000 | |||
| estimatedQuota := 500 | |||
| err := CreatePendingSyncRecord(userId, remoteUserId, requestId, quota, estimatedQuota) | |||
| require.NoError(t, err) | |||
| var record PendingSyncRecord | |||
| require.NoError(t, db.Where("request_id = ?", requestId).First(&record).Error) | |||
| assert.Equal(t, userId, record.UserId) | |||
| assert.Equal(t, remoteUserId, record.RemoteUserId) | |||
| assert.Equal(t, requestId, record.RequestId) | |||
| assert.Equal(t, quota, record.Quota) | |||
| assert.Equal(t, estimatedQuota, record.EstimatedQuota) | |||
| assert.Equal(t, PendingSyncStatusPending, record.Status) | |||
| assert.NotZero(t, record.CreatedAt) | |||
| assert.Zero(t, record.SyncedAt) | |||
| assert.Zero(t, record.RetryCount) | |||
| } | |||
| func TestPendingSyncRecord_RequestIdUnique(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| requestId := "req-unique-001" | |||
| // 创建第一条记录 | |||
| err := CreatePendingSyncRecord(100, 200, requestId, 1000, 500) | |||
| require.NoError(t, err) | |||
| // 尝试用相同的 requestId 创建第二条记录 | |||
| err = CreatePendingSyncRecord(101, 201, requestId, 2000, 600) | |||
| assert.Error(t, err) // 应该失败,因为 requestId 唯一 | |||
| // 验证只有一条记录 | |||
| var count int64 | |||
| db.Model(&PendingSyncRecord{}).Where("request_id = ?", requestId).Count(&count) | |||
| assert.Equal(t, int64(1), count) | |||
| } | |||
| func TestGetPendingSyncQuota(t *testing.T) { | |||
| setupPendingSyncDB(t) | |||
| userId := 100 | |||
| // 创建多条记录 | |||
| require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-001", 1000, 500)) | |||
| require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-002", 2000, 600)) | |||
| require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-003", 500, 200)) | |||
| // 创建一个 failed 状态的记录,应该被计入 | |||
| require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-004", 300, 100)) | |||
| DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-004").Update("status", PendingSyncStatusFailed) | |||
| // 创建一个 synced 状态的记录,不应该被计入 | |||
| require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-005", 9999, 0)) | |||
| DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-005").Update("status", PendingSyncStatusSynced) | |||
| total := GetPendingSyncQuota(userId) | |||
| assert.Equal(t, 3800, total) // 1000 + 2000 + 500 + 300 = 3800 | |||
| } | |||
| func TestGetPendingSyncQuota_Empty(t *testing.T) { | |||
| setupPendingSyncDB(t) | |||
| // 用户没有待同步记录 | |||
| total := GetPendingSyncQuota(99999) | |||
| assert.Equal(t, 0, total) | |||
| } | |||
| // QuotaSyncLog tests | |||
| func TestQuotaSyncLog_Create(t *testing.T) { | |||
| db := setupRegionSyncDB(t) | |||
| require.NoError(t, db.AutoMigrate(&QuotaSyncLog{})) | |||
| log := QuotaSyncLog{ | |||
| UserId: 100, | |||
| RemoteUserId: 100, | |||
| SyncType: SyncTypeUserCreate, | |||
| BeforeQuota: 0, | |||
| AfterQuota: 500000, | |||
| ChangeAmount: 500000, | |||
| MasterQuota: 500000, | |||
| PendingQuota: 0, | |||
| Direction: SyncDirectionCnToOv, | |||
| Status: SyncStatusSuccess, | |||
| CreatedAt: 1700000000, | |||
| SyncedAt: 1700000000, | |||
| } | |||
| require.NoError(t, db.Create(&log).Error) | |||
| var loaded QuotaSyncLog | |||
| require.NoError(t, db.First(&loaded, log.Id).Error) | |||
| assert.Equal(t, SyncTypeUserCreate, loaded.SyncType) | |||
| assert.Equal(t, 500000, loaded.ChangeAmount) | |||
| assert.Equal(t, SyncDirectionCnToOv, loaded.Direction) | |||
| } | |||
| func TestQuotaSyncLog_QueryByUser(t *testing.T) { | |||
| db := setupRegionSyncDB(t) | |||
| require.NoError(t, db.AutoMigrate(&QuotaSyncLog{})) | |||
| for i := 0; i < 3; i++ { | |||
| log := QuotaSyncLog{ | |||
| UserId: 100, RemoteUserId: 100, | |||
| SyncType: SyncTypeBatchSync, Direction: SyncDirectionOvToCn, | |||
| Status: SyncStatusSuccess, ChangeAmount: 100 * (i + 1), | |||
| CreatedAt: 1700000000 + int64(i), | |||
| } | |||
| require.NoError(t, db.Create(&log).Error) | |||
| } | |||
| var logs []QuotaSyncLog | |||
| db.Where("user_id = ?", 100).Order("created_at desc").Find(&logs) | |||
| assert.Len(t, logs, 3) | |||
| } | |||
| // --------------------------------------------------------------------------- | |||
| // 阶段 1:Model 层缺失测试 | |||
| // --------------------------------------------------------------------------- | |||
| func TestGetPendingRecordsForSync_Basic(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| now := time.Now().Unix() | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-rs-001", | |||
| Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now - 100, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 101, RemoteUserId: 201, RequestId: "req-rs-002", | |||
| Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now - 50, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 102, RemoteUserId: 202, RequestId: "req-rs-003", | |||
| Quota: 300, Status: PendingSyncStatusSynced, CreatedAt: now, | |||
| }).Error) | |||
| records := GetPendingRecordsForSync(10, 3) | |||
| assert.Len(t, records, 2) | |||
| // 按 created_at 升序 | |||
| assert.Equal(t, "req-rs-001", records[0].RequestId) | |||
| assert.Equal(t, "req-rs-002", records[1].RequestId) | |||
| } | |||
| func TestGetPendingRecordsForSync_MaxRetry(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| now := time.Now().Unix() | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-retry-001", | |||
| Quota: 100, Status: PendingSyncStatusFailed, RetryCount: 3, CreatedAt: now, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 101, RemoteUserId: 201, RequestId: "req-retry-002", | |||
| Quota: 200, Status: PendingSyncStatusPending, RetryCount: 2, CreatedAt: now, | |||
| }).Error) | |||
| records := GetPendingRecordsForSync(10, 3) | |||
| assert.Len(t, records, 1) | |||
| assert.Equal(t, "req-retry-002", records[0].RequestId) | |||
| } | |||
| func TestGetPendingRecordsForSync_Limit(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| now := time.Now().Unix() | |||
| for i := 0; i < 5; i++ { | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: fmt.Sprintf("req-lim-%d", i), | |||
| Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now + int64(i), | |||
| }).Error) | |||
| } | |||
| records := GetPendingRecordsForSync(2, 3) | |||
| assert.Len(t, records, 2) | |||
| } | |||
| func TestMarkRecordSynced(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-sync-001", | |||
| Quota: 100, Status: PendingSyncStatusPending, CreatedAt: time.Now().Unix(), | |||
| }).Error) | |||
| var record PendingSyncRecord | |||
| require.NoError(t, db.Where("request_id = ?", "req-sync-001").First(&record).Error) | |||
| require.NoError(t, MarkRecordSynced(record.Id)) | |||
| var updated PendingSyncRecord | |||
| require.NoError(t, db.First(&updated, record.Id).Error) | |||
| assert.Equal(t, PendingSyncStatusSynced, updated.Status) | |||
| assert.NotZero(t, updated.SyncedAt) | |||
| } | |||
| func TestMarkRecordFailed(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-fail-001", | |||
| Quota: 100, Status: PendingSyncStatusPending, RetryCount: 0, CreatedAt: time.Now().Unix(), | |||
| }).Error) | |||
| var record PendingSyncRecord | |||
| require.NoError(t, db.Where("request_id = ?", "req-fail-001").First(&record).Error) | |||
| require.NoError(t, MarkRecordFailed(record.Id, "connection refused")) | |||
| var updated PendingSyncRecord | |||
| require.NoError(t, db.First(&updated, record.Id).Error) | |||
| assert.Equal(t, PendingSyncStatusFailed, updated.Status) | |||
| assert.Equal(t, 1, updated.RetryCount) | |||
| assert.Equal(t, "connection refused", updated.ErrorMsg) | |||
| } | |||
| func TestCleanupOldSyncedRecords(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| now := time.Now().Unix() | |||
| // 40 天前已同步的记录 | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-old-001", | |||
| Quota: 100, Status: PendingSyncStatusSynced, SyncedAt: now - 40*86400, CreatedAt: now - 45*86400, | |||
| }).Error) | |||
| // 5 天前已同步的记录(不应被清理) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-old-002", | |||
| Quota: 200, Status: PendingSyncStatusSynced, SyncedAt: now - 5*86400, CreatedAt: now - 10*86400, | |||
| }).Error) | |||
| // pending 状态的记录(不应被清理) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-old-003", | |||
| Quota: 300, Status: PendingSyncStatusPending, CreatedAt: now - 50*86400, | |||
| }).Error) | |||
| deleted := CleanupOldSyncedRecords(30) | |||
| assert.Equal(t, int64(1), deleted) | |||
| var count int64 | |||
| db.Model(&PendingSyncRecord{}).Count(&count) | |||
| assert.Equal(t, int64(2), count) | |||
| } | |||
| func TestCleanupOldSyncedRecords_Empty(t *testing.T) { | |||
| setupPendingSyncDB(t) | |||
| deleted := CleanupOldSyncedRecords(30) | |||
| assert.Equal(t, int64(0), deleted) | |||
| } | |||
| func TestGetPendingQuotaByUser(t *testing.T) { | |||
| db := setupPendingSyncDB(t) | |||
| now := time.Now().Unix() | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-pq-001", | |||
| Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-pq-002", | |||
| Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 101, RemoteUserId: 201, RequestId: "req-pq-003", | |||
| Quota: 500, Status: PendingSyncStatusPending, CreatedAt: now, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&PendingSyncRecord{ | |||
| UserId: 100, RemoteUserId: 200, RequestId: "req-pq-004", | |||
| Quota: 9999, Status: PendingSyncStatusSynced, CreatedAt: now, | |||
| }).Error) | |||
| m := GetPendingQuotaByUser() | |||
| assert.Equal(t, 300, m[100]) // 100 + 200, synced 不计入 | |||
| assert.Equal(t, 500, m[101]) | |||
| _, exists := m[102] | |||
| assert.False(t, exists) | |||
| } | |||
| // --------------------------------------------------------------------------- | |||
| // QuotaUpdateCallback tests | |||
| // --------------------------------------------------------------------------- | |||
| func TestSetQuotaUpdateCallback(t *testing.T) { | |||
| // 测试设置回调函数 | |||
| called := false | |||
| var callbackUserId int | |||
| var callbackQuota int | |||
| callback := func(userId int, quota int) { | |||
| called = true | |||
| callbackUserId = userId | |||
| callbackQuota = quota | |||
| } | |||
| SetQuotaUpdateCallback(callback) | |||
| // 清理:重置为 nil | |||
| defer SetQuotaUpdateCallback(nil) | |||
| // 验证回调被设置(通过调用测试) | |||
| if quotaUpdateCallback != nil { | |||
| quotaUpdateCallback(123, 500000) | |||
| } | |||
| assert.True(t, called) | |||
| assert.Equal(t, 123, callbackUserId) | |||
| assert.Equal(t, 500000, callbackQuota) | |||
| } | |||
| func TestIncreaseUserQuota_TriggersCallback(t *testing.T) { | |||
| db := setupRegionSyncDB(t) | |||
| // 创建测试用户 | |||
| user := User{ | |||
| Id: 1000, | |||
| Username: "callback_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 设置回调 | |||
| callbackCalled := false | |||
| var receivedUserId int | |||
| var receivedQuota int | |||
| SetQuotaUpdateCallback(func(userId int, quota int) { | |||
| callbackCalled = true | |||
| receivedUserId = userId | |||
| receivedQuota = quota | |||
| }) | |||
| defer SetQuotaUpdateCallback(nil) // 清理 | |||
| // 执行增加额度操作 | |||
| err := IncreaseUserQuota(1000, 50000, true) | |||
| require.NoError(t, err) | |||
| // 验证回调被触发 | |||
| assert.True(t, callbackCalled) | |||
| assert.Equal(t, 1000, receivedUserId) | |||
| assert.Equal(t, 150000, receivedQuota) // 原始 100000 + 增加 50000 | |||
| // 验证数据库已更新 | |||
| var updated User | |||
| require.NoError(t, db.First(&updated, 1000).Error) | |||
| assert.Equal(t, 150000, updated.Quota) | |||
| } | |||
| func TestIncreaseUserQuota_NoCallback(t *testing.T) { | |||
| db := setupRegionSyncDB(t) | |||
| // 创建测试用户 | |||
| user := User{ | |||
| Id: 1001, | |||
| Username: "no_callback_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 确保没有回调 | |||
| SetQuotaUpdateCallback(nil) | |||
| // 执行增加额度操作(不应 panic) | |||
| err := IncreaseUserQuota(1001, 50000, true) | |||
| require.NoError(t, err) | |||
| // 验证数据库已更新 | |||
| var updated User | |||
| require.NoError(t, db.First(&updated, 1001).Error) | |||
| assert.Equal(t, 150000, updated.Quota) | |||
| } | |||
| func TestIncreaseUserQuota_CallbackError(t *testing.T) { | |||
| db := setupRegionSyncDB(t) | |||
| // 创建测试用户 | |||
| user := User{ | |||
| Id: 1002, | |||
| Username: "callback_error_test", | |||
| Password: "hashed_password", | |||
| Quota: 100000, | |||
| } | |||
| require.NoError(t, db.Create(&user).Error) | |||
| // 设置一个会 panic 的回调,验证不影响主流程 | |||
| SetQuotaUpdateCallback(func(userId int, quota int) { | |||
| panic("callback error") | |||
| }) | |||
| defer SetQuotaUpdateCallback(nil) | |||
| // 执行增加额度操作(回调 panic 不应影响数据库更新) | |||
| err := IncreaseUserQuota(1002, 50000, true) | |||
| require.NoError(t, err) | |||
| // 验证数据库已更新(即使回调失败) | |||
| var updated User | |||
| require.NoError(t, db.First(&updated, 1002).Error) | |||
| assert.Equal(t, 150000, updated.Quota) | |||
| } | |||
| @@ -6,6 +6,7 @@ import ( | |||
| "fmt" | |||
| "strconv" | |||
| "strings" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| @@ -17,6 +18,16 @@ import ( | |||
| const UserNameMaxLength = 20 | |||
| // QuotaUpdateCallback 余额更新回调函数类型 | |||
| type QuotaUpdateCallback func(userId int, quota int) | |||
| var quotaUpdateCallback QuotaUpdateCallback | |||
| // SetQuotaUpdateCallback 设置余额更新回调 | |||
| func SetQuotaUpdateCallback(cb QuotaUpdateCallback) { | |||
| quotaUpdateCallback = cb | |||
| } | |||
| // User if you add sensitive fields, don't forget to clean them in setupLogin function. | |||
| // Otherwise, the sensitive information will be saved on local storage in plain text! | |||
| type User struct { | |||
| @@ -49,6 +60,21 @@ type User struct { | |||
| Setting string `json:"setting" gorm:"type:text;column:setting"` | |||
| Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"` | |||
| StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"` | |||
| // 跨地区同步相关 | |||
| Source string `json:"source" gorm:"type:varchar(20);default:'local'"` | |||
| RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"` | |||
| SyncedQuota int `json:"synced_quota" gorm:"type:int;default:0;column:synced_quota"` | |||
| LastSyncAt int64 `json:"last_sync_at" gorm:"type:bigint;default:0;column:last_sync_at"` | |||
| } | |||
| // IsSyncedUser 判断是否为国内同步用户 | |||
| func (u *User) IsSyncedUser() bool { | |||
| return common.IsSyncedUser(u.Source, u.Id) | |||
| } | |||
| // IsLocalUser 判断是否为本地用户 | |||
| func (u *User) IsLocalUser() bool { | |||
| return common.IsLocalUser(u.Source, u.Id) | |||
| } | |||
| func (user *User) ToBaseUser() *UserBase { | |||
| @@ -60,6 +86,7 @@ func (user *User) ToBaseUser() *UserBase { | |||
| Username: user.Username, | |||
| Setting: user.Setting, | |||
| Email: user.Email, | |||
| Source: user.Source, | |||
| } | |||
| return cache | |||
| } | |||
| @@ -767,17 +794,8 @@ func ValidateAccessToken(token string) (user *User) { | |||
| } | |||
| // GetUserQuota gets quota from Redis first, falls back to DB if needed | |||
| // 同步用户返回 SyncedQuota,本地用户返回 Quota | |||
| func GetUserQuota(id int, fromDB bool) (quota int, err error) { | |||
| defer func() { | |||
| // Update Redis cache asynchronously on successful DB read | |||
| if shouldUpdateRedis(fromDB, err) { | |||
| gopool.Go(func() { | |||
| if err := updateUserQuotaCache(id, quota); err != nil { | |||
| common.SysLog("failed to update user quota cache: " + err.Error()) | |||
| } | |||
| }) | |||
| } | |||
| }() | |||
| if !fromDB && common.RedisEnabled { | |||
| quota, err := getUserQuotaCache(id) | |||
| if err == nil { | |||
| @@ -785,12 +803,31 @@ func GetUserQuota(id int, fromDB bool) (quota int, err error) { | |||
| } | |||
| // Don't return error - fall through to DB | |||
| } | |||
| fromDB = true | |||
| err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find("a).Error | |||
| // 查询用户的 source 和额度字段 | |||
| var user User | |||
| err = DB.Model(&User{}).Select("source", "quota", "synced_quota").Where("id = ?", id).First(&user).Error | |||
| if err != nil { | |||
| return 0, err | |||
| } | |||
| // 同步用户返回 SyncedQuota,本地用户返回 Quota | |||
| if user.IsSyncedUser() { | |||
| quota = user.SyncedQuota | |||
| } else { | |||
| quota = user.Quota | |||
| } | |||
| // Update Redis cache asynchronously if enabled | |||
| if common.RedisEnabled { | |||
| gopool.Go(func() { | |||
| if err := updateUserQuotaCache(id, quota); err != nil { | |||
| common.SysLog("failed to update user quota cache: " + err.Error()) | |||
| } | |||
| }) | |||
| } | |||
| return quota, nil | |||
| } | |||
| @@ -877,7 +914,26 @@ func IncreaseUserQuota(id int, quota int, db bool) (err error) { | |||
| addNewRecord(BatchUpdateTypeUserQuota, id, quota) | |||
| return nil | |||
| } | |||
| return increaseUserQuota(id, quota) | |||
| err = increaseUserQuota(id, quota) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| // 触发余额更新回调(如果有注册) | |||
| if quotaUpdateCallback != nil { | |||
| newQuota, _ := GetUserQuota(id, true) | |||
| // 使用 defer/recover 保护回调,防止 panic 影响主流程 | |||
| func() { | |||
| defer func() { | |||
| if r := recover(); r != nil { | |||
| common.SysError(fmt.Sprintf("quota update callback panic: %v", r)) | |||
| } | |||
| }() | |||
| quotaUpdateCallback(id, newQuota) | |||
| }() | |||
| } | |||
| return nil | |||
| } | |||
| func increaseUserQuota(id int, quota int) (err error) { | |||
| @@ -924,6 +980,47 @@ func DeltaUpdateUserQuota(id int, delta int) (err error) { | |||
| } | |||
| } | |||
| // GetSyncedUsers 获取所有同步用户 | |||
| func GetSyncedUsers() []User { | |||
| var users []User | |||
| DB.Where("source = ?", common.UserSourceSynced).Find(&users) | |||
| return users | |||
| } | |||
| // UpdateSyncedQuota 更新同步用户的 synced_quota | |||
| func UpdateSyncedQuota(userId int, quota int) error { | |||
| return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]interface{}{ | |||
| "synced_quota": quota, | |||
| "last_sync_at": time.Now().Unix(), | |||
| }).Error | |||
| } | |||
| // AtomicDecreaseSyncedQuota 原子扣减 synced_quota,检查余额并扣减在一条 SQL 中完成。 | |||
| // 返回 (当前余额, 是否成功, 错误)。 | |||
| func AtomicDecreaseSyncedQuota(userId int, amount int, threshold int) (int, bool, error) { | |||
| result := DB.Model(&User{}). | |||
| Where("id = ? AND synced_quota >= ?", userId, amount+threshold). | |||
| Update("synced_quota", gorm.Expr("synced_quota - ?", amount)) | |||
| if result.Error != nil { | |||
| return 0, false, result.Error | |||
| } | |||
| if result.RowsAffected == 0 { | |||
| var quota int | |||
| DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan("a) | |||
| return quota, false, nil | |||
| } | |||
| var newQuota int | |||
| DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan(&newQuota) | |||
| return newQuota, true, nil | |||
| } | |||
| // IncreaseSyncedQuota 原子增加 synced_quota(用于退款) | |||
| func IncreaseSyncedQuota(userId int, amount int) error { | |||
| return DB.Model(&User{}). | |||
| Where("id = ?", userId). | |||
| Update("synced_quota", gorm.Expr("synced_quota + ?", amount)).Error | |||
| } | |||
| //func GetRootUserEmail() (email string) { | |||
| // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email) | |||
| // return email | |||
| @@ -15,6 +15,7 @@ import ( | |||
| // UserBase struct remains the same as it represents the cached data structure | |||
| type UserBase struct { | |||
| Source string `json:"source"` | |||
| Id int `json:"id"` | |||
| Group string `json:"group"` | |||
| Email string `json:"email"` | |||
| @@ -31,6 +32,7 @@ func (user *UserBase) WriteContext(c *gin.Context) { | |||
| common.SetContextKey(c, constant.ContextKeyUserEmail, user.Email) | |||
| common.SetContextKey(c, constant.ContextKeyUserName, user.Username) | |||
| common.SetContextKey(c, constant.ContextKeyUserSetting, user.GetSetting()) | |||
| common.SetContextKey(c, constant.ContextKeyUserSource, user.Source) | |||
| } | |||
| func (user *UserBase) GetSetting() dto.UserSetting { | |||
| @@ -101,6 +103,7 @@ func GetUserCache(userId int) (userCache *UserBase, err error) { | |||
| // Create cache object from user data | |||
| userCache = &UserBase{ | |||
| Id: user.Id, | |||
| Source: user.Source, | |||
| Group: user.Group, | |||
| Quota: user.Quota, | |||
| Status: user.Status, | |||
| @@ -87,9 +87,8 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens | |||
| if meta.MaxTokens != 0 { | |||
| preConsumedTokens += meta.MaxTokens | |||
| } | |||
| // 如果没有找到渠道定价,需要获取全局 modelRatio 和 completionRatio | |||
| // 但如果 ChannelMeta 为 nil(渠道选择前),跳过检查,延迟到渠道选择后 | |||
| if channelMetaAvailable && !channelPricingFound { | |||
| // 如果没有找到渠道定价,获取全局 modelRatio 和 completionRatio | |||
| if !channelPricingFound { | |||
| var success bool | |||
| var matchName string | |||
| modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName) | |||
| @@ -388,5 +388,15 @@ func SetApiRouter(router *gin.Engine) { | |||
| deploymentsRoute.POST("/:id/extend", controller.ExtendDeployment) | |||
| deploymentsRoute.DELETE("/:id", controller.DeleteDeployment) | |||
| } | |||
| // Region sync API routes (for inter-node communication) | |||
| syncRoute := apiRouter.Group("/internal/sync") | |||
| syncRoute.Use(middleware.SyncAuth()) | |||
| { | |||
| syncRoute.POST("/user/create", controller.ReceiveSyncedUserCreate) | |||
| syncRoute.POST("/quota/update", controller.ReceiveQuotaUpdate) | |||
| syncRoute.POST("/quota/query", controller.QueryUserQuota) | |||
| syncRoute.POST("/quota/batch-deduct", controller.BatchDeductQuota) | |||
| } | |||
| } | |||
| } | |||
| @@ -7,6 +7,7 @@ import ( | |||
| "sync" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| @@ -132,6 +133,10 @@ func (s *BillingSession) needsRefundLocked() bool { | |||
| if sub, ok := s.funding.(*SubscriptionFunding); ok && sub.preConsumed > 0 { | |||
| return true | |||
| } | |||
| // 同步用户可能在 tokenConsumed=0 时仍预扣了额度 | |||
| if synced, ok := s.funding.(*SyncedUserFunding); ok && synced.consumed > 0 { | |||
| return true | |||
| } | |||
| return false | |||
| } | |||
| @@ -257,6 +262,18 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons | |||
| return nil, types.NewError(fmt.Errorf("relayInfo is nil"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) | |||
| } | |||
| // 检查是否为同步用户(从上下文中获取,避免热路径数据库查询) | |||
| userSource := common.GetContextKeyString(c, constant.ContextKeyUserSource) | |||
| common.SysLog(fmt.Sprintf("[RegionSync] NewBillingSession: userId=%d, userSource=%s, preConsumedQuota=%d, requestId=%s", relayInfo.UserId, userSource, preConsumedQuota, relayInfo.RequestId)) | |||
| if userSource == common.UserSourceSynced { | |||
| user, userErr := model.GetUserById(relayInfo.UserId, false) | |||
| if userErr == nil { | |||
| common.SysLog(fmt.Sprintf("[RegionSync] synced user detected: userId=%d, remoteUserId=%d, syncedQuota=%d, requestId=%s", user.Id, user.RemoteUserId, user.SyncedQuota, relayInfo.RequestId)) | |||
| return newSyncedUserBillingSession(c, relayInfo, user, preConsumedQuota) | |||
| } | |||
| common.SysError(fmt.Sprintf("[RegionSync] GetUserById failed for synced user: userId=%d, err=%v", relayInfo.UserId, userErr)) | |||
| } | |||
| pref := common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference) | |||
| // 钱包路径需要先检查用户额度 | |||
| @@ -345,3 +362,29 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons | |||
| return session, nil | |||
| } | |||
| } | |||
| // newSyncedUserBillingSession 为同步用户创建 BillingSession | |||
| // 同步用户的余额在 master 节点,本地使用 SyncedUserFunding 管理预扣费和结算 | |||
| func newSyncedUserBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, user *model.User, preConsumedQuota int) (*BillingSession, *types.NewAPIError) { | |||
| syncedQuota := user.SyncedQuota | |||
| if syncedQuota <= 0 { | |||
| syncedQuota = user.Quota | |||
| } | |||
| session := &BillingSession{ | |||
| relayInfo: relayInfo, | |||
| funding: NewSyncedUserFunding( | |||
| user.Id, | |||
| user.RemoteUserId, | |||
| relayInfo.RequestId, | |||
| syncedQuota, | |||
| ), | |||
| } | |||
| relayInfo.UserQuota = syncedQuota | |||
| if apiErr := session.preConsume(c, preConsumedQuota); apiErr != nil { | |||
| return nil, apiErr | |||
| } | |||
| return session, nil | |||
| } | |||
| @@ -0,0 +1,135 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupBillingSyncDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| sqlDB, _ := db.DB() | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := model.DB | |||
| origLogDB := model.LOG_DB | |||
| model.DB = db | |||
| model.LOG_DB = db | |||
| common.UsingSQLite = true | |||
| common.RedisEnabled = false | |||
| require.NoError(t, db.AutoMigrate(&model.PendingSyncRecord{})) | |||
| t.Cleanup(func() { | |||
| model.DB = origDB | |||
| model.LOG_DB = origLogDB | |||
| sqlDB.Close() | |||
| }) | |||
| return db | |||
| } | |||
| // TestNeedsRefundLocked_SyncedUser 测试同步用户在 consumed > 0 时需要退款 | |||
| func TestNeedsRefundLocked_SyncedUser(t *testing.T) { | |||
| setupBillingSyncDB(t) | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| origThreshold := settings.MinBalanceThreshold | |||
| defer func() { settings.MinBalanceThreshold = origThreshold }() | |||
| settings.MinBalanceThreshold = 100000 | |||
| f := NewSyncedUserFunding(100, 200, "req-nr-001", 500000) | |||
| require.NoError(t, f.PreConsume(5000)) | |||
| session := &BillingSession{ | |||
| funding: f, | |||
| } | |||
| assert.True(t, session.needsRefundLocked()) | |||
| } | |||
| // TestNeedsRefundLocked_SyncedUser_NoConsumed 测试同步用户在 consumed == 0 时不需要退款 | |||
| func TestNeedsRefundLocked_SyncedUser_NoConsumed(t *testing.T) { | |||
| setupBillingSyncDB(t) | |||
| f := NewSyncedUserFunding(100, 200, "req-nr-002", 500000) | |||
| session := &BillingSession{ | |||
| funding: f, | |||
| } | |||
| assert.False(t, session.needsRefundLocked()) | |||
| } | |||
| // TestNeedsRefundLocked_SyncedUser_FundingSettled 测试 fundingSettled=true 后不需要退款 | |||
| func TestNeedsRefundLocked_SyncedUser_FundingSettled(t *testing.T) { | |||
| setupBillingSyncDB(t) | |||
| settings := system_setting.GetRegionSyncSettings() | |||
| origThreshold := settings.MinBalanceThreshold | |||
| defer func() { settings.MinBalanceThreshold = origThreshold }() | |||
| settings.MinBalanceThreshold = 100000 | |||
| f := NewSyncedUserFunding(100, 200, "req-nr-003", 500000) | |||
| require.NoError(t, f.PreConsume(5000)) | |||
| session := &BillingSession{ | |||
| funding: f, | |||
| fundingSettled: true, // 资金已提交 | |||
| } | |||
| // fundingSettled=true 时不应再退 | |||
| assert.False(t, session.needsRefundLocked()) | |||
| } | |||
| // TestShouldTrust_SyncedUser_SourceNotWalletOrSub 验证同步用户 Source 不是 wallet/subscription, | |||
| // 因此在 shouldTrust 的 switch 中走 default 分支返回 false | |||
| func TestShouldTrust_SyncedUser_SourceNotWalletOrSub(t *testing.T) { | |||
| f := NewSyncedUserFunding(100, 200, "req-trust", 500000) | |||
| assert.Equal(t, "synced_wallet", f.Source()) | |||
| assert.NotEqual(t, BillingSourceWallet, f.Source()) | |||
| assert.NotEqual(t, BillingSourceSubscription, f.Source()) | |||
| } | |||
| // TestSyncedUserFunding_BillingSource 测试同步用户计费来源标识 | |||
| func TestSyncedUserFunding_BillingSource(t *testing.T) { | |||
| f := NewSyncedUserFunding(100, 200, "req-src", 500000) | |||
| assert.Equal(t, BillingSourceSyncedWallet, f.Source()) | |||
| assert.Equal(t, "synced_wallet", f.Source()) | |||
| } | |||
| // TestNeedsRefundLocked_WalletAndSynced 对比钱包和同步用户的行为 | |||
| func TestNeedsRefundLocked_WalletAndSynced(t *testing.T) { | |||
| // 钱包用户:tokenConsumed > 0 才需要退 | |||
| walletSession := &BillingSession{ | |||
| funding: &WalletFunding{userId: 1, consumed: 10 | |||