Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge commit from fork

Browse Source
This commit is contained in:
RedwindA
2026-08-06 17:53:55 +08:00
committed by GitHub
parent ea4f021012
commit 0cd9dc85e3
4 changed files with 117 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
controller/user.go
+3 -10
View File
@@ -399,11 +399,6 @@ func GetUser(c *gin.Context) {
func GenerateAccessToken(c *gin.Context) {
id := c.GetInt("id")
user, err := model.GetUserById(id, true)
if err != nil {
common.ApiError(c, err)
return
}
// get rand int 28-32
randI := common.GetRandomInt(4)
key, err := common.GenerateRandomKey(29 + randI)
@@ -412,14 +407,12 @@ func GenerateAccessToken(c *gin.Context) {
common.SysLog("failed to generate key: " + err.Error())
return
}
user.SetAccessToken(key)
if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 {
if model.DB.Where("access_token = ?", key).First(&model.User{}).RowsAffected != 0 {
common.ApiErrorI18n(c, i18n.MsgUuidDuplicate)
return
}
if err := user.Update(false); err != nil {
if err := model.UpdateUserAccessToken(id, key); err != nil {
common.ApiError(c, err)
return
}
@@ -427,7 +420,7 @@ func GenerateAccessToken(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": user.AccessToken,
"data": key,
})
return
}
model/user.go
+39 -10
View File
@@ -139,6 +139,22 @@ func (user *User) SetAccessToken(token string) {
user.AccessToken = &token
}
// UpdateUserAccessToken rotates a dashboard personal access token without
// writing a stale user snapshot back over concurrently updated fields.
func UpdateUserAccessToken(id int, token string) error {
if id == 0 {
return errors.New("id 为空！")
}
result := DB.Model(&User{}).Where("id = ?", id).Update("access_token", token)
if result.Error != nil {
return result.Error
}
if result.RowsAffected == 0 {
return gorm.ErrRecordNotFound
}
return nil
}
func (user *User) GetSetting() dto.UserSetting {
setting := dto.UserSetting{}
if user.Setting != "" {
@@ -489,15 +505,19 @@ func HardDeleteUserById(id int) error {
return user.HardDelete()
}
func inviteUser(inviterId int) (err error) {
user, err := GetUserById(inviterId, true)
if err != nil {
return err
func inviteUser(inviterId int) error {
result := DB.Model(&User{}).Where("id = ?", inviterId).Updates(map[string]interface{}{
"aff_count": gorm.Expr("aff_count + ?", 1),
"aff_quota": gorm.Expr("aff_quota + ?", common.QuotaForInviter),
"aff_history": gorm.Expr("aff_history + ?", common.QuotaForInviter),
})
if result.Error != nil {
return result.Error
}
user.AffCount++
user.AffQuota += common.QuotaForInviter
user.AffHistoryQuota += common.QuotaForInviter
return DB.Save(user).Error
if result.RowsAffected == 0 {
return gorm.ErrRecordNotFound
}
return nil
}
func (user *User) TransferAffQuotaToQuota(quota int) error {
@@ -514,7 +534,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error {
defer tx.Rollback() // 确保在函数退出时事务能回滚
// 加锁查询用户以确保数据一致性
err := lockForUpdate(tx).First(&user, user.Id).Error
err := lockForUpdate(tx).First(user, user.Id).Error
if err != nil {
return err
}
@@ -748,7 +768,16 @@ func (user *User) UpdateWithTx(tx *gorm.DB, updatePassword bool) error {
return err
}
}
if err = tx.Model(&current).Omit("quota", "used_quota", "request_count", "auth_version").Updates(newUser).Error; err != nil {
if err = tx.Model(&current).Omit(
"access_token",
"quota",
"used_quota",
"request_count",
"aff_count",
"aff_quota",
"aff_history",
"auth_version",
).Updates(newUser).Error; err != nil {
return err
}
return tx.First(user, user.Id).Error
model/user_update_test.go
+74 -9
View File
@@ -27,19 +27,23 @@ func setupUserUpdateTestState(t *testing.T) {
})
}
func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) {
func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing.T) {
setupUserUpdateTestState(t)
user := User{
Id: 1,
Username: "quota-race-user",
Password: "password",
DisplayName: "before",
Status: common.UserStatusEnabled,
Quota: 1000,
UsedQuota: 20,
RequestCount: 3,
Id: 1,
Username: "quota-race-user",
Password: "password",
DisplayName: "before",
Status: common.UserStatusEnabled,
Quota: 1000,
UsedQuota: 20,
RequestCount: 3,
AffCount: 2,
AffQuota: 800,
AffHistoryQuota: 1200,
}
user.SetAccessToken("old-token")
require.NoError(t, DB.Create(&user).Error)
staleUser, err := GetUserById(user.Id, true)
@@ -49,6 +53,10 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) {
"quota": gorm.Expr("quota - ?", 400),
"used_quota": gorm.Expr("used_quota + ?", 400),
"request_count": gorm.Expr("request_count + ?", 1),
"aff_count": gorm.Expr("aff_count + ?", 1),
"aff_quota": gorm.Expr("aff_quota - ?", 500),
"aff_history": gorm.Expr("aff_history + ?", 500),
"access_token": "rotated-token",
}).Error)
staleUser.DisplayName = "after"
@@ -60,6 +68,63 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) {
assert.Equal(t, 600, got.Quota)
assert.Equal(t, 420, got.UsedQuota)
assert.Equal(t, 4, got.RequestCount)
assert.Equal(t, 3, got.AffCount)
assert.Equal(t, 300, got.AffQuota)
assert.Equal(t, 1700, got.AffHistoryQuota)
assert.Equal(t, "rotated-token", got.GetAccessToken())
}
func TestUpdateUserAccessTokenOnlyUpdatesAccessToken(t *testing.T) {
setupUserUpdateTestState(t)
user := User{
Id: 2,
Username: "token-rotation-user",
Password: "password",
DisplayName: "before",
Status: common.UserStatusEnabled,
Quota: 1000,
AffQuota: 800,
AffHistoryQuota: 1200,
}
require.NoError(t, DB.Create(&user).Error)
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{
"quota": gorm.Expr("quota + ?", 500),
"aff_quota": gorm.Expr("aff_quota - ?", 500),
"display_name": "concurrent-update",
}).Error)
require.NoError(t, UpdateUserAccessToken(user.Id, "rotated-token"))
var got User
require.NoError(t, DB.First(&got, user.Id).Error)
assert.Equal(t, "rotated-token", got.GetAccessToken())
assert.Equal(t, "concurrent-update", got.DisplayName)
assert.Equal(t, 1500, got.Quota)
assert.Equal(t, 300, got.AffQuota)
assert.Equal(t, 1200, got.AffHistoryQuota)
}
func TestUpdateUserAccessTokenRejectsSoftDeletedUser(t *testing.T) {
setupUserUpdateTestState(t)
user := User{
Id: 3,
Username: "deleted-token-rotation-user",
Password: "password",
Status: common.UserStatusEnabled,
}
user.SetAccessToken("old-token")
require.NoError(t, DB.Create(&user).Error)
require.NoError(t, DB.Delete(&user).Error)
err := UpdateUserAccessToken(user.Id, "orphaned-token")
require.ErrorIs(t, err, gorm.ErrRecordNotFound)
var got User
require.NoError(t, DB.Unscoped().First(&got, user.Id).Error)
assert.Equal(t, "old-token", got.GetAccessToken())
}
func TestUpdateUserSettingOnlyUpdatesSetting(t *testing.T) {
router/api-router.go
+1 -1
View File
@@ -90,7 +90,7 @@ func SetApiRouter(router *gin.Engine) {
selfRoute.GET("/models", controller.GetUserModels)
selfRoute.PUT("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.UpdateSelf)
selfRoute.DELETE("/self", controller.DeleteSelf)
selfRoute.GET("/token", middleware.DisableCache(), controller.GenerateAccessToken)
selfRoute.GET("/token", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GenerateAccessToken)
selfRoute.GET("/passkey", controller.PasskeyStatus)
selfRoute.POST("/passkey/register/begin", middleware.DisableCache(), controller.PasskeyRegisterBegin)
selfRoute.POST("/passkey/register/finish", middleware.DisableCache(), controller.PasskeyRegisterFinish)