Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

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