diff --git a/controller/user.go b/controller/user.go index 59aaf584..9b8d931e 100644 --- a/controller/user.go +++ b/controller/user.go @@ -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 } diff --git a/model/user.go b/model/user.go index b25de5e7..eb4ea086 100644 --- a/model/user.go +++ b/model/user.go @@ -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(¤t).Omit("quota", "used_quota", "request_count", "auth_version").Updates(newUser).Error; err != nil { + if err = tx.Model(¤t).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 diff --git a/model/user_update_test.go b/model/user_update_test.go index c68d8f33..8e69ba1e 100644 --- a/model/user_update_test.go +++ b/model/user_update_test.go @@ -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) { diff --git a/router/api-router.go b/router/api-router.go index 907cf1ed..2b0bfd6c 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -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)