Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.
 
 
 

1171 lignes
33 KiB

  1. package model
  2. import (
  3. "encoding/json"
  4. "errors"
  5. "fmt"
  6. "strconv"
  7. "strings"
  8. "time"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/dto"
  11. "github.com/QuantumNous/new-api/logger"
  12. "github.com/bytedance/gopkg/util/gopool"
  13. "gorm.io/gorm"
  14. )
  15. const UserNameMaxLength = 20
  16. // QuotaUpdateCallback 余额更新回调函数类型
  17. type QuotaUpdateCallback func(userId int, quota int)
  18. var quotaUpdateCallback QuotaUpdateCallback
  19. // SetQuotaUpdateCallback 设置余额更新回调
  20. func SetQuotaUpdateCallback(cb QuotaUpdateCallback) {
  21. quotaUpdateCallback = cb
  22. }
  23. // User if you add sensitive fields, don't forget to clean them in setupLogin function.
  24. // Otherwise, the sensitive information will be saved on local storage in plain text!
  25. type User struct {
  26. Id int `json:"id"`
  27. Username string `json:"username" gorm:"unique;index" validate:"max=20"`
  28. Password string `json:"password" gorm:"not null;" validate:"min=8,max=20"`
  29. OriginalPassword string `json:"original_password" gorm:"-:all"` // this field is only for Password change verification, don't save it to database!
  30. DisplayName string `json:"display_name" gorm:"index" validate:"max=20"`
  31. Role int `json:"role" gorm:"type:int;default:1"` // admin, common
  32. Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled
  33. Email string `json:"email" gorm:"index" validate:"max=50"`
  34. GitHubId string `json:"github_id" gorm:"column:github_id;index"`
  35. DiscordId string `json:"discord_id" gorm:"column:discord_id;index"`
  36. OidcId string `json:"oidc_id" gorm:"column:oidc_id;index"`
  37. WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"`
  38. TelegramId string `json:"telegram_id" gorm:"column:telegram_id;index"`
  39. VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
  40. AccessToken *string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
  41. Quota int `json:"quota" gorm:"type:int;default:0"`
  42. UsedQuota int `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
  43. RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number
  44. Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
  45. AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
  46. AffCount int `json:"aff_count" gorm:"type:int;default:0;column:aff_count"`
  47. AffQuota int `json:"aff_quota" gorm:"type:int;default:0;column:aff_quota"` // 邀请剩余额度
  48. AffHistoryQuota int `json:"aff_history_quota" gorm:"type:int;default:0;column:aff_history"` // 邀请历史额度
  49. InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
  50. DeletedAt gorm.DeletedAt `gorm:"index"`
  51. LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"`
  52. Setting string `json:"setting" gorm:"type:text;column:setting"`
  53. Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"`
  54. StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"`
  55. // 跨地区同步相关
  56. Source string `json:"source" gorm:"type:varchar(20);default:'local'"`
  57. RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"`
  58. SyncedQuota int `json:"synced_quota" gorm:"type:int;default:0;column:synced_quota"`
  59. LastSyncAt int64 `json:"last_sync_at" gorm:"type:bigint;default:0;column:last_sync_at"`
  60. }
  61. // IsSyncedUser 判断是否为国内同步用户
  62. func (u *User) IsSyncedUser() bool {
  63. return common.IsSyncedUser(u.Source)
  64. }
  65. // IsLocalUser 判断是否为本地用户
  66. func (u *User) IsLocalUser() bool {
  67. return common.IsLocalUser(u.Source)
  68. }
  69. func (user *User) ToBaseUser() *UserBase {
  70. cache := &UserBase{
  71. Id: user.Id,
  72. Group: user.Group,
  73. Quota: user.Quota,
  74. Status: user.Status,
  75. Username: user.Username,
  76. Setting: user.Setting,
  77. Email: user.Email,
  78. Source: user.Source,
  79. }
  80. return cache
  81. }
  82. func (user *User) GetAccessToken() string {
  83. if user.AccessToken == nil {
  84. return ""
  85. }
  86. return *user.AccessToken
  87. }
  88. func (user *User) SetAccessToken(token string) {
  89. user.AccessToken = &token
  90. }
  91. func (user *User) GetSetting() dto.UserSetting {
  92. setting := dto.UserSetting{}
  93. if user.Setting != "" {
  94. err := json.Unmarshal([]byte(user.Setting), &setting)
  95. if err != nil {
  96. common.SysLog("failed to unmarshal setting: " + err.Error())
  97. }
  98. }
  99. return setting
  100. }
  101. func (user *User) SetSetting(setting dto.UserSetting) {
  102. settingBytes, err := json.Marshal(setting)
  103. if err != nil {
  104. common.SysLog("failed to marshal setting: " + err.Error())
  105. return
  106. }
  107. user.Setting = string(settingBytes)
  108. }
  109. // 根据用户角色生成默认的边栏配置
  110. func generateDefaultSidebarConfigForRole(userRole int) string {
  111. defaultConfig := map[string]interface{}{}
  112. // 聊天区域 - 所有用户都可以访问
  113. defaultConfig["chat"] = map[string]interface{}{
  114. "enabled": true,
  115. "playground": true,
  116. "chat": true,
  117. }
  118. // 控制台区域 - 所有用户都可以访问
  119. defaultConfig["console"] = map[string]interface{}{
  120. "enabled": true,
  121. "detail": true,
  122. "token": true,
  123. "log": true,
  124. "midjourney": true,
  125. "task": true,
  126. }
  127. // 个人中心区域 - 所有用户都可以访问
  128. defaultConfig["personal"] = map[string]interface{}{
  129. "enabled": true,
  130. "topup": true,
  131. "personal": true,
  132. }
  133. // 管理员区域 - 根据角色决定
  134. if userRole == common.RoleAdminUser {
  135. // 管理员可以访问管理员区域,但不能访问系统设置
  136. defaultConfig["admin"] = map[string]interface{}{
  137. "enabled": true,
  138. "channel": true,
  139. "models": true,
  140. "redemption": true,
  141. "user": true,
  142. "setting": false, // 管理员不能访问系统设置
  143. }
  144. } else if userRole == common.RoleRootUser {
  145. // 超级管理员可以访问所有功能
  146. defaultConfig["admin"] = map[string]interface{}{
  147. "enabled": true,
  148. "channel": true,
  149. "models": true,
  150. "redemption": true,
  151. "user": true,
  152. "setting": true,
  153. }
  154. }
  155. // 普通用户不包含admin区域
  156. // 转换为JSON字符串
  157. configBytes, err := json.Marshal(defaultConfig)
  158. if err != nil {
  159. common.SysLog("生成默认边栏配置失败: " + err.Error())
  160. return ""
  161. }
  162. return string(configBytes)
  163. }
  164. // CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil
  165. func CheckUserExistOrDeleted(username string, email string) (bool, error) {
  166. var user User
  167. // err := DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
  168. // check email if empty
  169. var err error
  170. if email == "" {
  171. err = DB.Unscoped().First(&user, "username = ?", username).Error
  172. } else {
  173. err = DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
  174. }
  175. if err != nil {
  176. if errors.Is(err, gorm.ErrRecordNotFound) {
  177. // not exist, return false, nil
  178. return false, nil
  179. }
  180. // other error, return false, err
  181. return false, err
  182. }
  183. // exist, return true, nil
  184. return true, nil
  185. }
  186. func GetMaxUserId() int {
  187. var user User
  188. DB.Unscoped().Last(&user)
  189. return user.Id
  190. }
  191. // ApplySyncedQuota 将同步用户的 quota 替换为 synced_quota,
  192. // 使 API 返回的 quota 对所有用户类型都表示实际可用余额。
  193. func (u *User) ApplySyncedQuota() {
  194. if u.IsSyncedUser() {
  195. u.Quota = u.SyncedQuota
  196. }
  197. }
  198. func applySyncedUserQuota(users []*User) {
  199. for _, u := range users {
  200. u.ApplySyncedQuota()
  201. }
  202. }
  203. func GetAllUsers(pageInfo *common.PageInfo) (users []*User, total int64, err error) {
  204. // Start transaction
  205. tx := DB.Begin()
  206. if tx.Error != nil {
  207. return nil, 0, tx.Error
  208. }
  209. defer func() {
  210. if r := recover(); r != nil {
  211. tx.Rollback()
  212. }
  213. }()
  214. // Get total count within transaction
  215. err = tx.Unscoped().Model(&User{}).Count(&total).Error
  216. if err != nil {
  217. tx.Rollback()
  218. return nil, 0, err
  219. }
  220. // Get paginated users within same transaction
  221. err = tx.Unscoped().Order("id desc").Limit(pageInfo.GetPageSize()).Offset(pageInfo.GetStartIdx()).Omit("password").Find(&users).Error
  222. if err != nil {
  223. tx.Rollback()
  224. return nil, 0, err
  225. }
  226. // Commit transaction
  227. if err = tx.Commit().Error; err != nil {
  228. return nil, 0, err
  229. }
  230. applySyncedUserQuota(users)
  231. return users, total, nil
  232. }
  233. func SearchUsers(keyword string, group string, startIdx int, num int) ([]*User, int64, error) {
  234. var users []*User
  235. var total int64
  236. var err error
  237. // 开始事务
  238. tx := DB.Begin()
  239. if tx.Error != nil {
  240. return nil, 0, tx.Error
  241. }
  242. defer func() {
  243. if r := recover(); r != nil {
  244. tx.Rollback()
  245. }
  246. }()
  247. // 构建基础查询
  248. query := tx.Unscoped().Model(&User{})
  249. // 构建搜索条件
  250. likeCondition := "username LIKE ? OR email LIKE ? OR display_name LIKE ?"
  251. // 尝试将关键字转换为整数ID
  252. keywordInt, err := strconv.Atoi(keyword)
  253. if err == nil {
  254. // 如果是数字,同时搜索ID和其他字段
  255. likeCondition = "id = ? OR " + likeCondition
  256. if group != "" {
  257. query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
  258. keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
  259. } else {
  260. query = query.Where(likeCondition,
  261. keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
  262. }
  263. } else {
  264. // 非数字关键字,只搜索字符串字段
  265. if group != "" {
  266. query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
  267. "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
  268. } else {
  269. query = query.Where(likeCondition,
  270. "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
  271. }
  272. }
  273. // 获取总数
  274. err = query.Count(&total).Error
  275. if err != nil {
  276. tx.Rollback()
  277. return nil, 0, err
  278. }
  279. // 获取分页数据
  280. err = query.Omit("password").Order("id desc").Limit(num).Offset(startIdx).Find(&users).Error
  281. if err != nil {
  282. tx.Rollback()
  283. return nil, 0, err
  284. }
  285. // 提交事务
  286. if err = tx.Commit().Error; err != nil {
  287. return nil, 0, err
  288. }
  289. applySyncedUserQuota(users)
  290. return users, total, nil
  291. }
  292. func GetUserById(id int, selectAll bool) (*User, error) {
  293. if id == 0 {
  294. return nil, errors.New("id 为空!")
  295. }
  296. user := User{Id: id}
  297. var err error = nil
  298. if selectAll {
  299. err = DB.First(&user, "id = ?", id).Error
  300. } else {
  301. err = DB.Omit("password").First(&user, "id = ?", id).Error
  302. }
  303. return &user, err
  304. }
  305. func GetUserIdByAffCode(affCode string) (int, error) {
  306. if affCode == "" {
  307. return 0, errors.New("affCode 为空!")
  308. }
  309. var user User
  310. err := DB.Select("id").First(&user, "aff_code = ?", affCode).Error
  311. return user.Id, err
  312. }
  313. func DeleteUserById(id int) (err error) {
  314. if id == 0 {
  315. return errors.New("id 为空!")
  316. }
  317. user := User{Id: id}
  318. return user.Delete()
  319. }
  320. func HardDeleteUserById(id int) error {
  321. if id == 0 {
  322. return errors.New("id 为空!")
  323. }
  324. err := DB.Unscoped().Delete(&User{}, "id = ?", id).Error
  325. return err
  326. }
  327. func inviteUser(inviterId int) (err error) {
  328. user, err := GetUserById(inviterId, true)
  329. if err != nil {
  330. return err
  331. }
  332. user.AffCount++
  333. user.AffQuota += common.QuotaForInviter
  334. user.AffHistoryQuota += common.QuotaForInviter
  335. return DB.Save(user).Error
  336. }
  337. func (user *User) TransferAffQuotaToQuota(quota int) error {
  338. // 检查quota是否小于最小额度
  339. if float64(quota) < common.QuotaPerUnit {
  340. return fmt.Errorf("转移额度最小为%s!", logger.LogQuota(int(common.QuotaPerUnit)))
  341. }
  342. // 开始数据库事务
  343. tx := DB.Begin()
  344. if tx.Error != nil {
  345. return tx.Error
  346. }
  347. defer tx.Rollback() // 确保在函数退出时事务能回滚
  348. // 加锁查询用户以确保数据一致性
  349. err := tx.Set("gorm:query_option", "FOR UPDATE").First(&user, user.Id).Error
  350. if err != nil {
  351. return err
  352. }
  353. // 再次检查用户的AffQuota是否足够
  354. if user.AffQuota < quota {
  355. return errors.New("邀请额度不足!")
  356. }
  357. // 更新用户额度
  358. user.AffQuota -= quota
  359. user.Quota += quota
  360. // 保存用户状态
  361. if err := tx.Save(user).Error; err != nil {
  362. return err
  363. }
  364. // 提交事务
  365. return tx.Commit().Error
  366. }
  367. func (user *User) Insert(inviterId int) error {
  368. var err error
  369. if user.Password != "" {
  370. user.Password, err = common.Password2Hash(user.Password)
  371. if err != nil {
  372. return err
  373. }
  374. }
  375. matchedQuota := MatchEmailQuotaRule(user.Email)
  376. if matchedQuota >= 0 {
  377. user.Quota = int(matchedQuota)
  378. } else {
  379. user.Quota = common.QuotaForNewUser
  380. }
  381. //user.SetAccessToken(common.GetUUID())
  382. user.AffCode = common.GetRandomString(4)
  383. // 初始化用户设置,包括默认的边栏配置
  384. if user.Setting == "" {
  385. defaultSetting := dto.UserSetting{}
  386. // 这里暂时不设置SidebarModules,因为需要在用户创建后根据角色设置
  387. user.SetSetting(defaultSetting)
  388. }
  389. result := DB.Create(user)
  390. if result.Error != nil {
  391. return result.Error
  392. }
  393. // 用户创建成功后,根据角色初始化边栏配置
  394. // 需要重新获取用户以确保有正确的ID和Role
  395. var createdUser User
  396. if err := DB.Where("username = ?", user.Username).First(&createdUser).Error; err == nil {
  397. // 生成基于角色的默认边栏配置
  398. defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
  399. if defaultSidebarConfig != "" {
  400. currentSetting := createdUser.GetSetting()
  401. currentSetting.SidebarModules = defaultSidebarConfig
  402. createdUser.SetSetting(currentSetting)
  403. createdUser.Update(false)
  404. common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
  405. }
  406. }
  407. if user.Quota > 0 {
  408. if MatchEmailQuotaRule(user.Email) >= 0 {
  409. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s(邮箱后缀规则匹配)", logger.LogQuota(user.Quota)))
  410. } else {
  411. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(user.Quota)))
  412. }
  413. }
  414. if inviterId != 0 {
  415. if common.QuotaForInvitee > 0 {
  416. _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
  417. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
  418. }
  419. if common.QuotaForInviter > 0 {
  420. //_ = IncreaseUserQuota(inviterId, common.QuotaForInviter)
  421. RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
  422. _ = inviteUser(inviterId)
  423. }
  424. }
  425. return nil
  426. }
  427. // InsertWithTx inserts a new user within an existing transaction.
  428. // This is used for OAuth registration where user creation and binding need to be atomic.
  429. // Post-creation tasks (sidebar config, logs, inviter rewards) are handled after the transaction commits.
  430. func (user *User) InsertWithTx(tx *gorm.DB, inviterId int) error {
  431. var err error
  432. if user.Password != "" {
  433. user.Password, err = common.Password2Hash(user.Password)
  434. if err != nil {
  435. return err
  436. }
  437. }
  438. matchedQuota := MatchEmailQuotaRule(user.Email)
  439. if matchedQuota >= 0 {
  440. user.Quota = int(matchedQuota)
  441. } else {
  442. user.Quota = common.QuotaForNewUser
  443. }
  444. user.AffCode = common.GetRandomString(4)
  445. // 初始化用户设置
  446. if user.Setting == "" {
  447. defaultSetting := dto.UserSetting{}
  448. user.SetSetting(defaultSetting)
  449. }
  450. result := tx.Create(user)
  451. if result.Error != nil {
  452. return result.Error
  453. }
  454. return nil
  455. }
  456. // FinalizeOAuthUserCreation performs post-transaction tasks for OAuth user creation.
  457. // This should be called after the transaction commits successfully.
  458. func (user *User) FinalizeOAuthUserCreation(inviterId int) {
  459. // 用户创建成功后,根据角色初始化边栏配置
  460. var createdUser User
  461. if err := DB.Where("id = ?", user.Id).First(&createdUser).Error; err == nil {
  462. defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
  463. if defaultSidebarConfig != "" {
  464. currentSetting := createdUser.GetSetting()
  465. currentSetting.SidebarModules = defaultSidebarConfig
  466. createdUser.SetSetting(currentSetting)
  467. createdUser.Update(false)
  468. common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
  469. }
  470. }
  471. if user.Quota > 0 {
  472. if matchedQuota := MatchEmailQuotaRule(user.Email); matchedQuota >= 0 {
  473. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s(邮箱后缀规则匹配)", logger.LogQuota(user.Quota)))
  474. } else {
  475. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(user.Quota)))
  476. }
  477. }
  478. if inviterId != 0 {
  479. if common.QuotaForInvitee > 0 {
  480. _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
  481. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
  482. }
  483. if common.QuotaForInviter > 0 {
  484. RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
  485. _ = inviteUser(inviterId)
  486. }
  487. }
  488. }
  489. func (user *User) Update(updatePassword bool) error {
  490. var err error
  491. if updatePassword {
  492. user.Password, err = common.Password2Hash(user.Password)
  493. if err != nil {
  494. return err
  495. }
  496. }
  497. newUser := *user
  498. DB.First(&user, user.Id)
  499. if err = DB.Model(user).Updates(newUser).Error; err != nil {
  500. return err
  501. }
  502. // Update cache
  503. return updateUserCache(*user)
  504. }
  505. func (user *User) Edit(updatePassword bool) error {
  506. var err error
  507. if updatePassword {
  508. user.Password, err = common.Password2Hash(user.Password)
  509. if err != nil {
  510. return err
  511. }
  512. }
  513. newUser := *user
  514. updates := map[string]interface{}{
  515. "username": newUser.Username,
  516. "display_name": newUser.DisplayName,
  517. "group": newUser.Group,
  518. "quota": newUser.Quota,
  519. "remark": newUser.Remark,
  520. }
  521. if updatePassword {
  522. updates["password"] = newUser.Password
  523. }
  524. DB.First(&user, user.Id)
  525. if err = DB.Model(user).Updates(updates).Error; err != nil {
  526. return err
  527. }
  528. // Update cache
  529. return updateUserCache(*user)
  530. }
  531. func (user *User) ClearBinding(bindingType string) error {
  532. if user.Id == 0 {
  533. return errors.New("user id is empty")
  534. }
  535. bindingColumnMap := map[string]string{
  536. "email": "email",
  537. "github": "github_id",
  538. "discord": "discord_id",
  539. "oidc": "oidc_id",
  540. "wechat": "wechat_id",
  541. "telegram": "telegram_id",
  542. "linuxdo": "linux_do_id",
  543. }
  544. column, ok := bindingColumnMap[bindingType]
  545. if !ok {
  546. return errors.New("invalid binding type")
  547. }
  548. if err := DB.Model(&User{}).Where("id = ?", user.Id).Update(column, "").Error; err != nil {
  549. return err
  550. }
  551. if err := DB.Where("id = ?", user.Id).First(user).Error; err != nil {
  552. return err
  553. }
  554. return updateUserCache(*user)
  555. }
  556. func (user *User) Delete() error {
  557. if user.Id == 0 {
  558. return errors.New("id 为空!")
  559. }
  560. if err := DB.Delete(user).Error; err != nil {
  561. return err
  562. }
  563. // 清除缓存
  564. return invalidateUserCache(user.Id)
  565. }
  566. func (user *User) HardDelete() error {
  567. if user.Id == 0 {
  568. return errors.New("id 为空!")
  569. }
  570. err := DB.Unscoped().Delete(user).Error
  571. return err
  572. }
  573. // ValidateAndFill check password & user status
  574. func (user *User) ValidateAndFill() (err error) {
  575. // When querying with struct, GORM will only query with non-zero fields,
  576. // that means if your field's value is 0, '', false or other zero values,
  577. // it won't be used to build query conditions
  578. password := user.Password
  579. username := strings.TrimSpace(user.Username)
  580. if username == "" || password == "" {
  581. return errors.New("用户名或密码为空")
  582. }
  583. // find buy username or email
  584. DB.Where("username = ? OR email = ?", username, username).First(user)
  585. okay := common.ValidatePasswordAndHash(password, user.Password)
  586. if !okay || user.Status != common.UserStatusEnabled {
  587. return errors.New("用户名或密码错误,或用户已被封禁")
  588. }
  589. return nil
  590. }
  591. func (user *User) FillUserById() error {
  592. if user.Id == 0 {
  593. return errors.New("id 为空!")
  594. }
  595. DB.Where(User{Id: user.Id}).First(user)
  596. return nil
  597. }
  598. func (user *User) FillUserByEmail() error {
  599. if user.Email == "" {
  600. return errors.New("email 为空!")
  601. }
  602. DB.Where(User{Email: user.Email}).First(user)
  603. return nil
  604. }
  605. func (user *User) FillUserByGitHubId() error {
  606. if user.GitHubId == "" {
  607. return errors.New("GitHub id 为空!")
  608. }
  609. DB.Where(User{GitHubId: user.GitHubId}).First(user)
  610. return nil
  611. }
  612. // UpdateGitHubId updates the user's GitHub ID (used for migration from login to numeric ID)
  613. func (user *User) UpdateGitHubId(newGitHubId string) error {
  614. if user.Id == 0 {
  615. return errors.New("user id is empty")
  616. }
  617. return DB.Model(user).Update("github_id", newGitHubId).Error
  618. }
  619. func (user *User) FillUserByDiscordId() error {
  620. if user.DiscordId == "" {
  621. return errors.New("discord id 为空!")
  622. }
  623. DB.Where(User{DiscordId: user.DiscordId}).First(user)
  624. return nil
  625. }
  626. func (user *User) FillUserByOidcId() error {
  627. if user.OidcId == "" {
  628. return errors.New("oidc id 为空!")
  629. }
  630. DB.Where(User{OidcId: user.OidcId}).First(user)
  631. return nil
  632. }
  633. func (user *User) FillUserByWeChatId() error {
  634. if user.WeChatId == "" {
  635. return errors.New("WeChat id 为空!")
  636. }
  637. DB.Where(User{WeChatId: user.WeChatId}).First(user)
  638. return nil
  639. }
  640. func (user *User) FillUserByTelegramId() error {
  641. if user.TelegramId == "" {
  642. return errors.New("Telegram id 为空!")
  643. }
  644. err := DB.Where(User{TelegramId: user.TelegramId}).First(user).Error
  645. if errors.Is(err, gorm.ErrRecordNotFound) {
  646. return errors.New("该 Telegram 账户未绑定")
  647. }
  648. return nil
  649. }
  650. func IsEmailAlreadyTaken(email string) bool {
  651. return DB.Unscoped().Where("email = ?", email).Find(&User{}).RowsAffected == 1
  652. }
  653. func IsWeChatIdAlreadyTaken(wechatId string) bool {
  654. return DB.Unscoped().Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1
  655. }
  656. func IsGitHubIdAlreadyTaken(githubId string) bool {
  657. return DB.Unscoped().Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1
  658. }
  659. func IsDiscordIdAlreadyTaken(discordId string) bool {
  660. return DB.Unscoped().Where("discord_id = ?", discordId).Find(&User{}).RowsAffected == 1
  661. }
  662. func IsOidcIdAlreadyTaken(oidcId string) bool {
  663. return DB.Where("oidc_id = ?", oidcId).Find(&User{}).RowsAffected == 1
  664. }
  665. func IsTelegramIdAlreadyTaken(telegramId string) bool {
  666. return DB.Unscoped().Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1
  667. }
  668. func ResetUserPasswordByEmail(email string, password string) error {
  669. if email == "" || password == "" {
  670. return errors.New("邮箱地址或密码为空!")
  671. }
  672. hashedPassword, err := common.Password2Hash(password)
  673. if err != nil {
  674. return err
  675. }
  676. err = DB.Model(&User{}).Where("email = ?", email).Update("password", hashedPassword).Error
  677. return err
  678. }
  679. func IsAdmin(userId int) bool {
  680. if userId == 0 {
  681. return false
  682. }
  683. var user User
  684. err := DB.Where("id = ?", userId).Select("role").Find(&user).Error
  685. if err != nil {
  686. common.SysLog("no such user " + err.Error())
  687. return false
  688. }
  689. return user.Role >= common.RoleAdminUser
  690. }
  691. //// IsUserEnabled checks user status from Redis first, falls back to DB if needed
  692. //func IsUserEnabled(id int, fromDB bool) (status bool, err error) {
  693. // defer func() {
  694. // // Update Redis cache asynchronously on successful DB read
  695. // if shouldUpdateRedis(fromDB, err) {
  696. // gopool.Go(func() {
  697. // if err := updateUserStatusCache(id, status); err != nil {
  698. // common.SysError("failed to update user status cache: " + err.Error())
  699. // }
  700. // })
  701. // }
  702. // }()
  703. // if !fromDB && common.RedisEnabled {
  704. // // Try Redis first
  705. // status, err := getUserStatusCache(id)
  706. // if err == nil {
  707. // return status == common.UserStatusEnabled, nil
  708. // }
  709. // // Don't return error - fall through to DB
  710. // }
  711. // fromDB = true
  712. // var user User
  713. // err = DB.Where("id = ?", id).Select("status").Find(&user).Error
  714. // if err != nil {
  715. // return false, err
  716. // }
  717. //
  718. // return user.Status == common.UserStatusEnabled, nil
  719. //}
  720. func ValidateAccessToken(token string) (user *User) {
  721. if token == "" {
  722. return nil
  723. }
  724. token = strings.Replace(token, "Bearer ", "", 1)
  725. user = &User{}
  726. if DB.Where("access_token = ?", token).First(user).RowsAffected == 1 {
  727. return user
  728. }
  729. return nil
  730. }
  731. // GetUserQuota gets quota from Redis first, falls back to DB if needed
  732. // 同步用户返回 SyncedQuota,本地用户返回 Quota
  733. func GetUserQuota(id int, fromDB bool) (quota int, err error) {
  734. if !fromDB && common.RedisEnabled {
  735. quota, err := getUserQuotaCache(id)
  736. if err == nil {
  737. return quota, nil
  738. }
  739. // Don't return error - fall through to DB
  740. }
  741. fromDB = true
  742. // 查询用户的 source 和额度字段
  743. var user User
  744. err = DB.Model(&User{}).Select("source", "quota", "synced_quota").Where("id = ?", id).First(&user).Error
  745. if err != nil {
  746. return 0, err
  747. }
  748. // 同步用户返回 SyncedQuota,本地用户返回 Quota
  749. if user.IsSyncedUser() {
  750. quota = user.SyncedQuota
  751. } else {
  752. quota = user.Quota
  753. }
  754. // Update Redis cache asynchronously if enabled
  755. if common.RedisEnabled {
  756. gopool.Go(func() {
  757. if err := updateUserQuotaCache(id, quota); err != nil {
  758. common.SysLog("failed to update user quota cache: " + err.Error())
  759. }
  760. })
  761. }
  762. return quota, nil
  763. }
  764. func GetUserUsedQuota(id int) (quota int, err error) {
  765. err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error
  766. return quota, err
  767. }
  768. func GetUserEmail(id int) (email string, err error) {
  769. err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
  770. return email, err
  771. }
  772. // GetUserGroup gets group from Redis first, falls back to DB if needed
  773. func GetUserGroup(id int, fromDB bool) (group string, err error) {
  774. defer func() {
  775. // Update Redis cache asynchronously on successful DB read
  776. if shouldUpdateRedis(fromDB, err) {
  777. gopool.Go(func() {
  778. if err := updateUserGroupCache(id, group); err != nil {
  779. common.SysLog("failed to update user group cache: " + err.Error())
  780. }
  781. })
  782. }
  783. }()
  784. if !fromDB && common.RedisEnabled {
  785. group, err := getUserGroupCache(id)
  786. if err == nil {
  787. return group, nil
  788. }
  789. // Don't return error - fall through to DB
  790. }
  791. fromDB = true
  792. err = DB.Model(&User{}).Where("id = ?", id).Select(commonGroupCol).Find(&group).Error
  793. if err != nil {
  794. return "", err
  795. }
  796. return group, nil
  797. }
  798. // GetUserSetting gets setting from Redis first, falls back to DB if needed
  799. func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error) {
  800. var setting string
  801. defer func() {
  802. // Update Redis cache asynchronously on successful DB read
  803. if shouldUpdateRedis(fromDB, err) {
  804. gopool.Go(func() {
  805. if err := updateUserSettingCache(id, setting); err != nil {
  806. common.SysLog("failed to update user setting cache: " + err.Error())
  807. }
  808. })
  809. }
  810. }()
  811. if !fromDB && common.RedisEnabled {
  812. setting, err := getUserSettingCache(id)
  813. if err == nil {
  814. return setting, nil
  815. }
  816. // Don't return error - fall through to DB
  817. }
  818. fromDB = true
  819. err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&setting).Error
  820. if err != nil {
  821. return settingMap, err
  822. }
  823. userBase := &UserBase{
  824. Setting: setting,
  825. }
  826. return userBase.GetSetting(), nil
  827. }
  828. func IncreaseUserQuota(id int, quota int, db bool) (err error) {
  829. if quota < 0 {
  830. return errors.New("quota 不能为负数!")
  831. }
  832. gopool.Go(func() {
  833. err := cacheIncrUserQuota(id, int64(quota))
  834. if err != nil {
  835. common.SysLog("failed to increase user quota: " + err.Error())
  836. }
  837. })
  838. if !db && common.BatchUpdateEnabled {
  839. addNewRecord(BatchUpdateTypeUserQuota, id, quota)
  840. return nil
  841. }
  842. err = increaseUserQuota(id, quota)
  843. if err != nil {
  844. return err
  845. }
  846. // 触发余额更新回调(如果有注册)
  847. if quotaUpdateCallback != nil {
  848. newQuota, _ := GetUserQuota(id, true)
  849. // 使用 defer/recover 保护回调,防止 panic 影响主流程
  850. func() {
  851. defer func() {
  852. if r := recover(); r != nil {
  853. common.SysError(fmt.Sprintf("quota update callback panic: %v", r))
  854. }
  855. }()
  856. quotaUpdateCallback(id, newQuota)
  857. }()
  858. }
  859. return nil
  860. }
  861. func increaseUserQuota(id int, quota int) (err error) {
  862. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
  863. if err != nil {
  864. return err
  865. }
  866. return err
  867. }
  868. func DecreaseUserQuota(id int, quota int) (err error) {
  869. if quota < 0 {
  870. return errors.New("quota 不能为负数!")
  871. }
  872. gopool.Go(func() {
  873. err := cacheDecrUserQuota(id, int64(quota))
  874. if err != nil {
  875. common.SysLog("failed to decrease user quota: " + err.Error())
  876. }
  877. })
  878. if common.BatchUpdateEnabled {
  879. addNewRecord(BatchUpdateTypeUserQuota, id, -quota)
  880. return nil
  881. }
  882. return decreaseUserQuota(id, quota)
  883. }
  884. func decreaseUserQuota(id int, quota int) (err error) {
  885. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
  886. if err != nil {
  887. return err
  888. }
  889. return err
  890. }
  891. func DeltaUpdateUserQuota(id int, delta int) (err error) {
  892. if delta == 0 {
  893. return nil
  894. }
  895. if delta > 0 {
  896. return IncreaseUserQuota(id, delta, false)
  897. } else {
  898. return DecreaseUserQuota(id, -delta)
  899. }
  900. }
  901. // GetSyncedUsers 获取所有同步用户
  902. func GetSyncedUsers() []User {
  903. var users []User
  904. DB.Where("source = ?", common.UserSourceSynced).Find(&users)
  905. return users
  906. }
  907. // UpdateSyncedQuota 更新同步用户的 synced_quota
  908. func UpdateSyncedQuota(userId int, quota int) error {
  909. return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]interface{}{
  910. "synced_quota": quota,
  911. "last_sync_at": time.Now().Unix(),
  912. }).Error
  913. }
  914. // AtomicDecreaseSyncedQuota 原子扣减 synced_quota,检查余额并扣减在一条 SQL 中完成。
  915. // 语义:扣减后 synced_quota 必须 >= threshold,即 synced_quota >= amount + threshold。
  916. // 返回 (当前余额, 是否成功, 错误)。
  917. func AtomicDecreaseSyncedQuota(userId int, amount int, threshold int) (int, bool, error) {
  918. result := DB.Model(&User{}).
  919. Where("id = ? AND synced_quota >= ?", userId, amount+threshold).
  920. Update("synced_quota", gorm.Expr("synced_quota - ?", amount))
  921. if result.Error != nil {
  922. return 0, false, result.Error
  923. }
  924. if result.RowsAffected == 0 {
  925. var quota int
  926. DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan(&quota)
  927. return quota, false, nil
  928. }
  929. var newQuota int
  930. DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan(&newQuota)
  931. return newQuota, true, nil
  932. }
  933. // IncreaseSyncedQuota 原子增加 synced_quota(用于退款)
  934. func IncreaseSyncedQuota(userId int, amount int) error {
  935. return DB.Model(&User{}).
  936. Where("id = ?", userId).
  937. Update("synced_quota", gorm.Expr("synced_quota + ?", amount)).Error
  938. }
  939. // AdjustSyncedQuota 原子调整 synced_quota(增量更新,避免并发覆盖)
  940. // delta > 0 表示额外扣减,delta < 0 表示退还。
  941. func AdjustSyncedQuota(userId int, delta int) error {
  942. return DB.Model(&User{}).Where("id = ?", userId).
  943. Update("synced_quota", gorm.Expr("synced_quota - ?", delta)).Error
  944. }
  945. //func GetRootUserEmail() (email string) {
  946. // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email)
  947. // return email
  948. //}
  949. func GetRootUser() (user *User) {
  950. DB.Where("role = ?", common.RoleRootUser).First(&user)
  951. return user
  952. }
  953. func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
  954. if common.BatchUpdateEnabled {
  955. addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
  956. addNewRecord(BatchUpdateTypeRequestCount, id, 1)
  957. return
  958. }
  959. updateUserUsedQuotaAndRequestCount(id, quota, 1)
  960. }
  961. func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
  962. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  963. map[string]interface{}{
  964. "used_quota": gorm.Expr("used_quota + ?", quota),
  965. "request_count": gorm.Expr("request_count + ?", count),
  966. },
  967. ).Error
  968. if err != nil {
  969. common.SysLog("failed to update user used quota and request count: " + err.Error())
  970. return
  971. }
  972. //// 更新缓存
  973. //if err := invalidateUserCache(id); err != nil {
  974. // common.SysError("failed to invalidate user cache: " + err.Error())
  975. //}
  976. }
  977. func updateUserUsedQuota(id int, quota int) {
  978. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  979. map[string]interface{}{
  980. "used_quota": gorm.Expr("used_quota + ?", quota),
  981. },
  982. ).Error
  983. if err != nil {
  984. common.SysLog("failed to update user used quota: " + err.Error())
  985. }
  986. }
  987. func updateUserRequestCount(id int, count int) {
  988. err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error
  989. if err != nil {
  990. common.SysLog("failed to update user request count: " + err.Error())
  991. }
  992. }
  993. // GetUsernameById gets username from Redis first, falls back to DB if needed
  994. func GetUsernameById(id int, fromDB bool) (username string, err error) {
  995. defer func() {
  996. // Update Redis cache asynchronously on successful DB read
  997. if shouldUpdateRedis(fromDB, err) {
  998. gopool.Go(func() {
  999. if err := updateUserNameCache(id, username); err != nil {
  1000. common.SysLog("failed to update user name cache: " + err.Error())
  1001. }
  1002. })
  1003. }
  1004. }()
  1005. if !fromDB && common.RedisEnabled {
  1006. username, err := getUserNameCache(id)
  1007. if err == nil {
  1008. return username, nil
  1009. }
  1010. // Don't return error - fall through to DB
  1011. }
  1012. fromDB = true
  1013. err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error
  1014. if err != nil {
  1015. return "", err
  1016. }
  1017. return username, nil
  1018. }
  1019. func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
  1020. var user User
  1021. err := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).First(&user).Error
  1022. return !errors.Is(err, gorm.ErrRecordNotFound)
  1023. }
  1024. func (user *User) FillUserByLinuxDOId() error {
  1025. if user.LinuxDOId == "" {
  1026. return errors.New("linux do id is empty")
  1027. }
  1028. err := DB.Where("linux_do_id = ?", user.LinuxDOId).First(user).Error
  1029. return err
  1030. }
  1031. func RootUserExists() bool {
  1032. var user User
  1033. err := DB.Where("role = ?", common.RoleRootUser).First(&user).Error
  1034. if err != nil {
  1035. return false
  1036. }
  1037. return true
  1038. }