user.go 29 KB

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