Skip to content

phsl / new-api

Go to a project
Commit 0cd9dc85 by RedwindA Committed by GitHub
Options

Merge commit from fork

parent ea4f0210
Showing with 109 additions and 22 deletions
controller/user.go
...@@ -399,11 +399,6 @@ func GetUser(c *gin.Context) { ...@@ -399,11 +399,6 @@ func GetUser(c *gin.Context) {
func GenerateAccessToken(c *gin.Context) { func GenerateAccessToken(c *gin.Context) {
id := c.GetInt("id") id := c.GetInt("id")
user, err := model.GetUserById(id, true)
if err != nil {
common.ApiError(c, err)
return
}
// get rand int 28-32 // get rand int 28-32
randI := common.GetRandomInt(4) randI := common.GetRandomInt(4)
key, err := common.GenerateRandomKey(29 + randI) key, err := common.GenerateRandomKey(29 + randI)
...@@ -412,14 +407,12 @@ func GenerateAccessToken(c *gin.Context) { ...@@ -412,14 +407,12 @@ func GenerateAccessToken(c *gin.Context) {
common.SysLog("failed to generate key: " + err.Error()) common.SysLog("failed to generate key: " + err.Error())
return return
} }
user.SetAccessToken(key) if model.DB.Where("access_token = ?", key).First(&model.User{}).RowsAffected != 0 {
if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 {
common.ApiErrorI18n(c, i18n.MsgUuidDuplicate) common.ApiErrorI18n(c, i18n.MsgUuidDuplicate)
return return
} }
if err := user.Update(false); err != nil { if err := model.UpdateUserAccessToken(id, key); err != nil {
common.ApiError(c, err) common.ApiError(c, err)
return return
} }
...@@ -427,7 +420,7 @@ func GenerateAccessToken(c *gin.Context) { ...@@ -427,7 +420,7 @@ func GenerateAccessToken(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": true, "success": true,
"message": "", "message": "",
"data": user.AccessToken, "data": key,
}) })
return return
} }
......
model/user.go
...@@ -139,6 +139,22 @@ func (user *User) SetAccessToken(token string) { ...@@ -139,6 +139,22 @@ func (user *User) SetAccessToken(token string) {
user.AccessToken = &token 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 { func (user *User) GetSetting() dto.UserSetting {
setting := dto.UserSetting{} setting := dto.UserSetting{}
if user.Setting != "" { if user.Setting != "" {
...@@ -489,15 +505,19 @@ func HardDeleteUserById(id int) error { ...@@ -489,15 +505,19 @@ func HardDeleteUserById(id int) error {
return user.HardDelete() return user.HardDelete()
} }
func inviteUser(inviterId int) (err error) { func inviteUser(inviterId int) error {
user, err := GetUserById(inviterId, true) result := DB.Model(&User{}).Where("id = ?", inviterId).Updates(map[string]interface{}{
if err != nil { "aff_count": gorm.Expr("aff_count + ?", 1),
return err "aff_quota": gorm.Expr("aff_quota + ?", common.QuotaForInviter),
"aff_history": gorm.Expr("aff_history + ?", common.QuotaForInviter),
})
if result.Error != nil {
return result.Error
}
if result.RowsAffected == 0 {
return gorm.ErrRecordNotFound
} }
user.AffCount++ return nil
user.AffQuota += common.QuotaForInviter
user.AffHistoryQuota += common.QuotaForInviter
return DB.Save(user).Error
} }
func (user *User) TransferAffQuotaToQuota(quota int) error { func (user *User) TransferAffQuotaToQuota(quota int) error {
...@@ -514,7 +534,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error { ...@@ -514,7 +534,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error {
defer tx.Rollback() // 确保在函数退出时事务能回滚 defer tx.Rollback() // 确保在函数退出时事务能回滚
// 加锁查询用户以确保数据一致性 // 加锁查询用户以确保数据一致性
err := lockForUpdate(tx).First(&user, user.Id).Error err := lockForUpdate(tx).First(user, user.Id).Error
if err != nil { if err != nil {
return err return err
} }
...@@ -748,7 +768,16 @@ func (user *User) UpdateWithTx(tx *gorm.DB, updatePassword bool) error { ...@@ -748,7 +768,16 @@ func (user *User) UpdateWithTx(tx *gorm.DB, updatePassword bool) error {
return err 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 err
} }
return tx.First(user, user.Id).Error return tx.First(user, user.Id).Error
......
model/user_update_test.go
...@@ -27,7 +27,7 @@ func setupUserUpdateTestState(t *testing.T) { ...@@ -27,7 +27,7 @@ func setupUserUpdateTestState(t *testing.T) {
}) })
} }
func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing.T) {
setupUserUpdateTestState(t) setupUserUpdateTestState(t)
user := User{ user := User{
...@@ -39,7 +39,11 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { ...@@ -39,7 +39,11 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) {
Quota: 1000, Quota: 1000,
UsedQuota: 20, UsedQuota: 20,
RequestCount: 3, RequestCount: 3,
AffCount: 2,
AffQuota: 800,
AffHistoryQuota: 1200,
} }
user.SetAccessToken("old-token")
require.NoError(t, DB.Create(&user).Error) require.NoError(t, DB.Create(&user).Error)
staleUser, err := GetUserById(user.Id, true) staleUser, err := GetUserById(user.Id, true)
...@@ -49,6 +53,10 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { ...@@ -49,6 +53,10 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) {
"quota": gorm.Expr("quota - ?", 400), "quota": gorm.Expr("quota - ?", 400),
"used_quota": gorm.Expr("used_quota + ?", 400), "used_quota": gorm.Expr("used_quota + ?", 400),
"request_count": gorm.Expr("request_count + ?", 1), "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) }).Error)
staleUser.DisplayName = "after" staleUser.DisplayName = "after"
...@@ -60,6 +68,63 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { ...@@ -60,6 +68,63 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) {
assert.Equal(t, 600, got.Quota) assert.Equal(t, 600, got.Quota)
assert.Equal(t, 420, got.UsedQuota) assert.Equal(t, 420, got.UsedQuota)
assert.Equal(t, 4, got.RequestCount) 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) { func TestUpdateUserSettingOnlyUpdatesSetting(t *testing.T) {
......
router/api-router.go
...@@ -90,7 +90,7 @@ func SetApiRouter(router *gin.Engine) { ...@@ -90,7 +90,7 @@ func SetApiRouter(router *gin.Engine) {
selfRoute.GET("/models", controller.GetUserModels) selfRoute.GET("/models", controller.GetUserModels)
selfRoute.PUT("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.UpdateSelf) selfRoute.PUT("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.UpdateSelf)
selfRoute.DELETE("/self", controller.DeleteSelf) 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.GET("/passkey", controller.PasskeyStatus)
selfRoute.POST("/passkey/register/begin", middleware.DisableCache(), controller.PasskeyRegisterBegin) selfRoute.POST("/passkey/register/begin", middleware.DisableCache(), controller.PasskeyRegisterBegin)
selfRoute.POST("/passkey/register/finish", middleware.DisableCache(), controller.PasskeyRegisterFinish) selfRoute.POST("/passkey/register/finish", middleware.DisableCache(), controller.PasskeyRegisterFinish)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment