Browse Source

feat: 跨区域配额同步(Master-Slave Region Sync)

实现 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
fengsilin 2 months ago
parent
commit
1cecd5349d
42 changed files with 5348 additions and 31 deletions
  1. +33
    -0
      common/region_sync.go
  2. +93
    -0
      common/region_sync_test.go
  3. +1
    -0
      constant/context_key.go
  4. +273
    -0
      controller/region_sync.go
  5. +466
    -0
      controller/region_sync_test.go
  6. +2
    -0
      controller/relay.go
  7. +10
    -8
      controller/user.go
  8. +22
    -1
      main.go
  9. +1
    -1
      makefile
  10. +96
    -0
      middleware/sync_auth.go
  11. +149
    -0
      middleware/sync_auth_test.go
  12. +4
    -0
      model/main.go
  13. +126
    -0
      model/pending_sync_record.go
  14. +55
    -0
      model/quota_sync_log.go
  15. +15
    -0
      model/redemption.go
  16. +254
    -0
      model/redemption_test.go
  17. +538
    -0
      model/region_sync_test.go
  18. +109
    -12
      model/user.go
  19. +3
    -0
      model/user_cache.go
  20. +2
    -3
      relay/helper/price.go
  21. +10
    -0
      router/api-router.go
  22. +43
    -0
      service/billing_session.go
  23. +135
    -0
      service/billing_session_sync_test.go
  24. +132
    -0
      service/region_sync/sync_client.go
  25. +113
    -0
      service/region_sync/sync_client_test.go
  26. +216
    -0
      service/region_sync/sync_manager.go
  27. +324
    -0
      service/region_sync/sync_manager_test.go
  28. +116
    -0
      service/region_sync/sync_push.go
  29. +171
    -0
      service/region_sync/sync_push_test.go
  30. +73
    -0
      service/region_sync/sync_types.go
  31. +37
    -0
      service/region_sync/sync_types_test.go
  32. +137
    -0
      service/synced_user_funding.go
  33. +212
    -0
      service/synced_user_funding_test.go
  34. +14
    -0
      setting/config/config.go
  35. +38
    -0
      setting/system_setting/region_sync.go
  36. +27
    -0
      setting/system_setting/region_sync_test.go
  37. +153
    -0
      test_region_sync.sh
  38. +156
    -0
      test_region_sync_concurrent.sh
  39. +533
    -0
      test_region_sync_e2e.sh
  40. BIN
      web/public/logo.png
  41. +31
    -6
      web/src/components/settings/OperationSetting.jsx
  42. +425
    -0
      web/src/pages/Setting/Operation/SettingsRegionSync.jsx

+ 33
- 0
common/region_sync.go View File

@@ -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
}

+ 93
- 0
common/region_sync_test.go View File

@@ -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)
}
})
}
}

+ 1
- 0
constant/context_key.go View File

@@ -51,6 +51,7 @@ const (
ContextKeyUserGroup ContextKey = "user_group"
ContextKeyUsingGroup ContextKey = "group"
ContextKeyUserName ContextKey = "username"
ContextKeyUserSource ContextKey = "user_source"

ContextKeyLocalCountTokens ContextKey = "local_count_tokens"



+ 273
- 0
controller/region_sync.go View File

@@ -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())
}
}

+ 466
- 0
controller/region_sync_test.go View File

@@ -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)
}

+ 2
- 0
controller/relay.go View File

@@ -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


+ 10
- 8
controller/user.go View File

@@ -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
- 1
main.go View File

@@ -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
}

+ 1
- 1
makefile View File

@@ -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)"


+ 96
- 0
middleware/sync_auth.go View File

@@ -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 ""
}

+ 149
- 0
middleware/sync_auth_test.go View File

@@ -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 "))
}

+ 4
- 0
model/main.go View File

@@ -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))


+ 126
- 0
model/pending_sync_record.go View File

@@ -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
}

+ 55
- 0
model/quota_sync_log.go View File

@@ -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
}

+ 15
- 0
model/redemption.go View File

@@ -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
}


+ 254
- 0
model/redemption_test.go View File

@@ -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)
}

+ 538
- 0
model/region_sync_test.go View File

@@ -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)
}

+ 109
- 12
model/user.go View File

@@ -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(&quota).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(&quota)
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


+ 3
- 0
model/user_cache.go View File

@@ -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,


+ 2
- 3
relay/helper/price.go View File

@@ -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)


+ 10
- 0
router/api-router.go View File

@@ -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)
}
}
}

+ 43
- 0
service/billing_session.go View File

@@ -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
}

+ 135
- 0
service/billing_session_sync_test.go View File

@@ -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