user.go 28 KB

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