user.go 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029
  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. // InsertWithTx inserts a new user within an existing transaction.
  382. // This is used for OAuth registration where user creation and binding need to be atomic.
  383. // Post-creation tasks (sidebar config, logs, inviter rewards) are handled after the transaction commits.
  384. func (user *User) InsertWithTx(tx *gorm.DB, inviterId int) error {
  385. var err error
  386. if user.Password != "" {
  387. user.Password, err = common.Password2Hash(user.Password)
  388. if err != nil {
  389. return err
  390. }
  391. }
  392. user.Quota = common.QuotaForNewUser
  393. user.AffCode = common.GetRandomString(4)
  394. // 初始化用户设置
  395. if user.Setting == "" {
  396. defaultSetting := dto.UserSetting{}
  397. user.SetSetting(defaultSetting)
  398. }
  399. result := tx.Create(user)
  400. if result.Error != nil {
  401. return result.Error
  402. }
  403. return nil
  404. }
  405. // FinalizeOAuthUserCreation performs post-transaction tasks for OAuth user creation.
  406. // This should be called after the transaction commits successfully.
  407. func (user *User) FinalizeOAuthUserCreation(inviterId int) {
  408. // 用户创建成功后,根据角色初始化边栏配置
  409. var createdUser User
  410. if err := DB.Where("id = ?", user.Id).First(&createdUser).Error; err == nil {
  411. defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
  412. if defaultSidebarConfig != "" {
  413. currentSetting := createdUser.GetSetting()
  414. currentSetting.SidebarModules = defaultSidebarConfig
  415. createdUser.SetSetting(currentSetting)
  416. createdUser.Update(false)
  417. common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
  418. }
  419. }
  420. if common.QuotaForNewUser > 0 {
  421. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(common.QuotaForNewUser)))
  422. }
  423. if inviterId != 0 {
  424. if common.QuotaForInvitee > 0 {
  425. _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
  426. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
  427. }
  428. if common.QuotaForInviter > 0 {
  429. RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
  430. _ = inviteUser(inviterId)
  431. }
  432. }
  433. }
  434. func (user *User) Update(updatePassword bool) error {
  435. var err error
  436. if updatePassword {
  437. user.Password, err = common.Password2Hash(user.Password)
  438. if err != nil {
  439. return err
  440. }
  441. }
  442. newUser := *user
  443. DB.First(&user, user.Id)
  444. if err = DB.Model(user).Updates(newUser).Error; err != nil {
  445. return err
  446. }
  447. // Update cache
  448. return updateUserCache(*user)
  449. }
  450. func (user *User) Edit(updatePassword bool) error {
  451. var err error
  452. if updatePassword {
  453. user.Password, err = common.Password2Hash(user.Password)
  454. if err != nil {
  455. return err
  456. }
  457. }
  458. newUser := *user
  459. updates := map[string]interface{}{
  460. "username": newUser.Username,
  461. "display_name": newUser.DisplayName,
  462. "group": newUser.Group,
  463. "quota": newUser.Quota,
  464. "remark": newUser.Remark,
  465. }
  466. if updatePassword {
  467. updates["password"] = newUser.Password
  468. }
  469. DB.First(&user, user.Id)
  470. if err = DB.Model(user).Updates(updates).Error; err != nil {
  471. return err
  472. }
  473. // Update cache
  474. return updateUserCache(*user)
  475. }
  476. func (user *User) ClearBinding(bindingType string) error {
  477. if user.Id == 0 {
  478. return errors.New("user id is empty")
  479. }
  480. bindingColumnMap := map[string]string{
  481. "email": "email",
  482. "github": "github_id",
  483. "discord": "discord_id",
  484. "oidc": "oidc_id",
  485. "wechat": "wechat_id",
  486. "telegram": "telegram_id",
  487. "linuxdo": "linux_do_id",
  488. }
  489. column, ok := bindingColumnMap[bindingType]
  490. if !ok {
  491. return errors.New("invalid binding type")
  492. }
  493. if err := DB.Model(&User{}).Where("id = ?", user.Id).Update(column, "").Error; err != nil {
  494. return err
  495. }
  496. if err := DB.Where("id = ?", user.Id).First(user).Error; err != nil {
  497. return err
  498. }
  499. return updateUserCache(*user)
  500. }
  501. func (user *User) Delete() error {
  502. if user.Id == 0 {
  503. return errors.New("id 为空!")
  504. }
  505. if err := DB.Delete(user).Error; err != nil {
  506. return err
  507. }
  508. // 清除缓存
  509. return invalidateUserCache(user.Id)
  510. }
  511. func (user *User) HardDelete() error {
  512. if user.Id == 0 {
  513. return errors.New("id 为空!")
  514. }
  515. err := DB.Unscoped().Delete(user).Error
  516. return err
  517. }
  518. // ValidateAndFill check password & user status
  519. func (user *User) ValidateAndFill() (err error) {
  520. // When querying with struct, GORM will only query with non-zero fields,
  521. // that means if your field's value is 0, '', false or other zero values,
  522. // it won't be used to build query conditions
  523. password := user.Password
  524. username := strings.TrimSpace(user.Username)
  525. if username == "" || password == "" {
  526. return errors.New("用户名或密码为空")
  527. }
  528. // find buy username or email
  529. DB.Where("username = ? OR email = ?", username, username).First(user)
  530. okay := common.ValidatePasswordAndHash(password, user.Password)
  531. if !okay || user.Status != common.UserStatusEnabled {
  532. return errors.New("用户名或密码错误,或用户已被封禁")
  533. }
  534. return nil
  535. }
  536. func (user *User) FillUserById() error {
  537. if user.Id == 0 {
  538. return errors.New("id 为空!")
  539. }
  540. DB.Where(User{Id: user.Id}).First(user)
  541. return nil
  542. }
  543. func (user *User) FillUserByEmail() error {
  544. if user.Email == "" {
  545. return errors.New("email 为空!")
  546. }
  547. DB.Where(User{Email: user.Email}).First(user)
  548. return nil
  549. }
  550. func (user *User) FillUserByGitHubId() error {
  551. if user.GitHubId == "" {
  552. return errors.New("GitHub id 为空!")
  553. }
  554. DB.Where(User{GitHubId: user.GitHubId}).First(user)
  555. return nil
  556. }
  557. // UpdateGitHubId updates the user's GitHub ID (used for migration from login to numeric ID)
  558. func (user *User) UpdateGitHubId(newGitHubId string) error {
  559. if user.Id == 0 {
  560. return errors.New("user id is empty")
  561. }
  562. return DB.Model(user).Update("github_id", newGitHubId).Error
  563. }
  564. func (user *User) FillUserByDiscordId() error {
  565. if user.DiscordId == "" {
  566. return errors.New("discord id 为空!")
  567. }
  568. DB.Where(User{DiscordId: user.DiscordId}).First(user)
  569. return nil
  570. }
  571. func (user *User) FillUserByOidcId() error {
  572. if user.OidcId == "" {
  573. return errors.New("oidc id 为空!")
  574. }
  575. DB.Where(User{OidcId: user.OidcId}).First(user)
  576. return nil
  577. }
  578. func (user *User) FillUserByWeChatId() error {
  579. if user.WeChatId == "" {
  580. return errors.New("WeChat id 为空!")
  581. }
  582. DB.Where(User{WeChatId: user.WeChatId}).First(user)
  583. return nil
  584. }
  585. func (user *User) FillUserByTelegramId() error {
  586. if user.TelegramId == "" {
  587. return errors.New("Telegram id 为空!")
  588. }
  589. err := DB.Where(User{TelegramId: user.TelegramId}).First(user).Error
  590. if errors.Is(err, gorm.ErrRecordNotFound) {
  591. return errors.New("该 Telegram 账户未绑定")
  592. }
  593. return nil
  594. }
  595. func IsEmailAlreadyTaken(email string) bool {
  596. return DB.Unscoped().Where("email = ?", email).Find(&User{}).RowsAffected == 1
  597. }
  598. func IsWeChatIdAlreadyTaken(wechatId string) bool {
  599. return DB.Unscoped().Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1
  600. }
  601. func IsGitHubIdAlreadyTaken(githubId string) bool {
  602. return DB.Unscoped().Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1
  603. }
  604. func IsDiscordIdAlreadyTaken(discordId string) bool {
  605. return DB.Unscoped().Where("discord_id = ?", discordId).Find(&User{}).RowsAffected == 1
  606. }
  607. func IsOidcIdAlreadyTaken(oidcId string) bool {
  608. return DB.Where("oidc_id = ?", oidcId).Find(&User{}).RowsAffected == 1
  609. }
  610. func IsTelegramIdAlreadyTaken(telegramId string) bool {
  611. return DB.Unscoped().Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1
  612. }
  613. func ResetUserPasswordByEmail(email string, password string) error {
  614. if email == "" || password == "" {
  615. return errors.New("邮箱地址或密码为空!")
  616. }
  617. hashedPassword, err := common.Password2Hash(password)
  618. if err != nil {
  619. return err
  620. }
  621. err = DB.Model(&User{}).Where("email = ?", email).Update("password", hashedPassword).Error
  622. return err
  623. }
  624. func IsAdmin(userId int) bool {
  625. if userId == 0 {
  626. return false
  627. }
  628. var user User
  629. err := DB.Where("id = ?", userId).Select("role").Find(&user).Error
  630. if err != nil {
  631. common.SysLog("no such user " + err.Error())
  632. return false
  633. }
  634. return user.Role >= common.RoleAdminUser
  635. }
  636. //// IsUserEnabled checks user status from Redis first, falls back to DB if needed
  637. //func IsUserEnabled(id int, fromDB bool) (status bool, err error) {
  638. // defer func() {
  639. // // Update Redis cache asynchronously on successful DB read
  640. // if shouldUpdateRedis(fromDB, err) {
  641. // gopool.Go(func() {
  642. // if err := updateUserStatusCache(id, status); err != nil {
  643. // common.SysError("failed to update user status cache: " + err.Error())
  644. // }
  645. // })
  646. // }
  647. // }()
  648. // if !fromDB && common.RedisEnabled {
  649. // // Try Redis first
  650. // status, err := getUserStatusCache(id)
  651. // if err == nil {
  652. // return status == common.UserStatusEnabled, nil
  653. // }
  654. // // Don't return error - fall through to DB
  655. // }
  656. // fromDB = true
  657. // var user User
  658. // err = DB.Where("id = ?", id).Select("status").Find(&user).Error
  659. // if err != nil {
  660. // return false, err
  661. // }
  662. //
  663. // return user.Status == common.UserStatusEnabled, nil
  664. //}
  665. func ValidateAccessToken(token string) (user *User) {
  666. if token == "" {
  667. return nil
  668. }
  669. token = strings.Replace(token, "Bearer ", "", 1)
  670. user = &User{}
  671. if DB.Where("access_token = ?", token).First(user).RowsAffected == 1 {
  672. return user
  673. }
  674. return nil
  675. }
  676. // GetUserQuota gets quota from Redis first, falls back to DB if needed
  677. func GetUserQuota(id int, fromDB bool) (quota int, err error) {
  678. defer func() {
  679. // Update Redis cache asynchronously on successful DB read
  680. if shouldUpdateRedis(fromDB, err) {
  681. gopool.Go(func() {
  682. if err := updateUserQuotaCache(id, quota); err != nil {
  683. common.SysLog("failed to update user quota cache: " + err.Error())
  684. }
  685. })
  686. }
  687. }()
  688. if !fromDB && common.RedisEnabled {
  689. quota, err := getUserQuotaCache(id)
  690. if err == nil {
  691. return quota, nil
  692. }
  693. // Don't return error - fall through to DB
  694. }
  695. fromDB = true
  696. err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error
  697. if err != nil {
  698. return 0, err
  699. }
  700. return quota, nil
  701. }
  702. func GetUserUsedQuota(id int) (quota int, err error) {
  703. err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error
  704. return quota, err
  705. }
  706. func GetUserEmail(id int) (email string, err error) {
  707. err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
  708. return email, err
  709. }
  710. // GetUserGroup gets group from Redis first, falls back to DB if needed
  711. func GetUserGroup(id int, fromDB bool) (group string, err error) {
  712. defer func() {
  713. // Update Redis cache asynchronously on successful DB read
  714. if shouldUpdateRedis(fromDB, err) {
  715. gopool.Go(func() {
  716. if err := updateUserGroupCache(id, group); err != nil {
  717. common.SysLog("failed to update user group cache: " + err.Error())
  718. }
  719. })
  720. }
  721. }()
  722. if !fromDB && common.RedisEnabled {
  723. group, err := getUserGroupCache(id)
  724. if err == nil {
  725. return group, nil
  726. }
  727. // Don't return error - fall through to DB
  728. }
  729. fromDB = true
  730. err = DB.Model(&User{}).Where("id = ?", id).Select(commonGroupCol).Find(&group).Error
  731. if err != nil {
  732. return "", err
  733. }
  734. return group, nil
  735. }
  736. // GetUserSetting gets setting from Redis first, falls back to DB if needed
  737. func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error) {
  738. var setting string
  739. defer func() {
  740. // Update Redis cache asynchronously on successful DB read
  741. if shouldUpdateRedis(fromDB, err) {
  742. gopool.Go(func() {
  743. if err := updateUserSettingCache(id, setting); err != nil {
  744. common.SysLog("failed to update user setting cache: " + err.Error())
  745. }
  746. })
  747. }
  748. }()
  749. if !fromDB && common.RedisEnabled {
  750. setting, err := getUserSettingCache(id)
  751. if err == nil {
  752. return setting, nil
  753. }
  754. // Don't return error - fall through to DB
  755. }
  756. fromDB = true
  757. err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&setting).Error
  758. if err != nil {
  759. return settingMap, err
  760. }
  761. userBase := &UserBase{
  762. Setting: setting,
  763. }
  764. return userBase.GetSetting(), nil
  765. }
  766. func IncreaseUserQuota(id int, quota int, db bool) (err error) {
  767. if quota < 0 {
  768. return errors.New("quota 不能为负数!")
  769. }
  770. gopool.Go(func() {
  771. err := cacheIncrUserQuota(id, int64(quota))
  772. if err != nil {
  773. common.SysLog("failed to increase user quota: " + err.Error())
  774. }
  775. })
  776. if !db && common.BatchUpdateEnabled {
  777. addNewRecord(BatchUpdateTypeUserQuota, id, quota)
  778. return nil
  779. }
  780. return increaseUserQuota(id, quota)
  781. }
  782. func increaseUserQuota(id int, quota int) (err error) {
  783. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
  784. if err != nil {
  785. return err
  786. }
  787. return err
  788. }
  789. func DecreaseUserQuota(id int, quota int) (err error) {
  790. if quota < 0 {
  791. return errors.New("quota 不能为负数!")
  792. }
  793. gopool.Go(func() {
  794. err := cacheDecrUserQuota(id, int64(quota))
  795. if err != nil {
  796. common.SysLog("failed to decrease user quota: " + err.Error())
  797. }
  798. })
  799. if common.BatchUpdateEnabled {
  800. addNewRecord(BatchUpdateTypeUserQuota, id, -quota)
  801. return nil
  802. }
  803. return decreaseUserQuota(id, quota)
  804. }
  805. func decreaseUserQuota(id int, quota int) (err error) {
  806. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
  807. if err != nil {
  808. return err
  809. }
  810. return err
  811. }
  812. func DeltaUpdateUserQuota(id int, delta int) (err error) {
  813. if delta == 0 {
  814. return nil
  815. }
  816. if delta > 0 {
  817. return IncreaseUserQuota(id, delta, false)
  818. } else {
  819. return DecreaseUserQuota(id, -delta)
  820. }
  821. }
  822. //func GetRootUserEmail() (email string) {
  823. // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email)
  824. // return email
  825. //}
  826. func GetRootUser() (user *User) {
  827. DB.Where("role = ?", common.RoleRootUser).First(&user)
  828. return user
  829. }
  830. func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
  831. if common.BatchUpdateEnabled {
  832. addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
  833. addNewRecord(BatchUpdateTypeRequestCount, id, 1)
  834. return
  835. }
  836. updateUserUsedQuotaAndRequestCount(id, quota, 1)
  837. }
  838. func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
  839. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  840. map[string]interface{}{
  841. "used_quota": gorm.Expr("used_quota + ?", quota),
  842. "request_count": gorm.Expr("request_count + ?", count),
  843. },
  844. ).Error
  845. if err != nil {
  846. common.SysLog("failed to update user used quota and request count: " + err.Error())
  847. return
  848. }
  849. //// 更新缓存
  850. //if err := invalidateUserCache(id); err != nil {
  851. // common.SysError("failed to invalidate user cache: " + err.Error())
  852. //}
  853. }
  854. func updateUserUsedQuota(id int, quota int) {
  855. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  856. map[string]interface{}{
  857. "used_quota": gorm.Expr("used_quota + ?", quota),
  858. },
  859. ).Error
  860. if err != nil {
  861. common.SysLog("failed to update user used quota: " + err.Error())
  862. }
  863. }
  864. func updateUserRequestCount(id int, count int) {
  865. err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error
  866. if err != nil {
  867. common.SysLog("failed to update user request count: " + err.Error())
  868. }
  869. }
  870. // GetUsernameById gets username from Redis first, falls back to DB if needed
  871. func GetUsernameById(id int, fromDB bool) (username string, err error) {
  872. defer func() {
  873. // Update Redis cache asynchronously on successful DB read
  874. if shouldUpdateRedis(fromDB, err) {
  875. gopool.Go(func() {
  876. if err := updateUserNameCache(id, username); err != nil {
  877. common.SysLog("failed to update user name cache: " + err.Error())
  878. }
  879. })
  880. }
  881. }()
  882. if !fromDB && common.RedisEnabled {
  883. username, err := getUserNameCache(id)
  884. if err == nil {
  885. return username, nil
  886. }
  887. // Don't return error - fall through to DB
  888. }
  889. fromDB = true
  890. err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error
  891. if err != nil {
  892. return "", err
  893. }
  894. return username, nil
  895. }
  896. func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
  897. var user User
  898. err := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).First(&user).Error
  899. return !errors.Is(err, gorm.ErrRecordNotFound)
  900. }
  901. func (user *User) FillUserByLinuxDOId() error {
  902. if user.LinuxDOId == "" {
  903. return errors.New("linux do id is empty")
  904. }
  905. err := DB.Where("linux_do_id = ?", user.LinuxDOId).First(user).Error
  906. return err
  907. }
  908. func RootUserExists() bool {
  909. var user User
  910. err := DB.Where("role = ?", common.RoleRootUser).First(&user).Error
  911. if err != nil {
  912. return false
  913. }
  914. return true
  915. }