user.go 25 KB

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