user.go 26 KB

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