Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

1310 wiersze
37 KiB

  1. package model
  2. import (
  3. "encoding/json"
  4. "errors"
  5. "fmt"
  6. "strconv"
  7. "strings"
  8. "time"
  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. // QuotaUpdateCallback 余额更新回调函数类型
  17. type QuotaUpdateCallback func(userId int, quota int)
  18. var quotaUpdateCallback QuotaUpdateCallback
  19. // SetQuotaUpdateCallback 设置余额更新回调
  20. func SetQuotaUpdateCallback(cb QuotaUpdateCallback) {
  21. quotaUpdateCallback = cb
  22. }
  23. // User if you add sensitive fields, don't forget to clean them in setupLogin function.
  24. // Otherwise, the sensitive information will be saved on local storage in plain text!
  25. type User struct {
  26. Id int `json:"id"`
  27. Username string `json:"username" gorm:"unique;index" validate:"max=20"`
  28. Password string `json:"password" gorm:"not null;" validate:"min=8,max=20"`
  29. OriginalPassword string `json:"original_password" gorm:"-:all"` // this field is only for Password change verification, don't save it to database!
  30. DisplayName string `json:"display_name" gorm:"index" validate:"max=20"`
  31. Role int `json:"role" gorm:"type:int;default:1"` // admin, common
  32. Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled
  33. Email string `json:"email" gorm:"index" validate:"max=50"`
  34. GitHubId string `json:"github_id" gorm:"column:github_id;index"`
  35. DiscordId string `json:"discord_id" gorm:"column:discord_id;index"`
  36. OidcId string `json:"oidc_id" gorm:"column:oidc_id;index"`
  37. WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"`
  38. TelegramId string `json:"telegram_id" gorm:"column:telegram_id;index"`
  39. VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
  40. AccessToken *string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
  41. Quota int `json:"quota" gorm:"type:int;default:0"`
  42. UsedQuota int `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
  43. RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number
  44. Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
  45. AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
  46. AffCount int `json:"aff_count" gorm:"type:int;default:0;column:aff_count"`
  47. AffQuota int `json:"aff_quota" gorm:"type:int;default:0;column:aff_quota"` // 邀请剩余额度
  48. AffHistoryQuota int `json:"aff_history_quota" gorm:"type:int;default:0;column:aff_history"` // 邀请历史额度
  49. InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
  50. CreatedAt int64 `json:"created_at" gorm:"bigint;default:0"`
  51. DeletedAt gorm.DeletedAt `gorm:"index"`
  52. LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"`
  53. Setting string `json:"setting" gorm:"type:text;column:setting"`
  54. Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"`
  55. StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"`
  56. // 跨地区同步相关
  57. Source string `json:"source" gorm:"type:varchar(20);default:'local'"`
  58. RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"`
  59. SyncedQuota int `json:"synced_quota" gorm:"type:int;default:0;column:synced_quota"`
  60. LastSyncAt int64 `json:"last_sync_at" gorm:"type:bigint;default:0;column:last_sync_at"`
  61. // Relay 抓包开关
  62. CaptureRelay bool `json:"capture_relay" gorm:"default:false"`
  63. }
  64. // IsSyncedUser 判断是否为国内同步用户
  65. func (u *User) IsSyncedUser() bool {
  66. return common.IsSyncedUser(u.Source)
  67. }
  68. // IsLocalUser 判断是否为本地用户
  69. func (u *User) IsLocalUser() bool {
  70. return common.IsLocalUser(u.Source)
  71. }
  72. func (user *User) ToBaseUser() *UserBase {
  73. cache := &UserBase{
  74. Id: user.Id,
  75. Group: user.Group,
  76. Quota: user.Quota,
  77. Status: user.Status,
  78. Username: user.Username,
  79. Setting: user.Setting,
  80. Email: user.Email,
  81. Source: user.Source,
  82. }
  83. return cache
  84. }
  85. func (user *User) GetAccessToken() string {
  86. if user.AccessToken == nil {
  87. return ""
  88. }
  89. return *user.AccessToken
  90. }
  91. func (user *User) SetAccessToken(token string) {
  92. user.AccessToken = &token
  93. }
  94. func (user *User) GetSetting() dto.UserSetting {
  95. setting := dto.UserSetting{}
  96. if user.Setting != "" {
  97. err := json.Unmarshal([]byte(user.Setting), &setting)
  98. if err != nil {
  99. common.SysLog("failed to unmarshal setting: " + err.Error())
  100. }
  101. }
  102. return setting
  103. }
  104. func (user *User) SetSetting(setting dto.UserSetting) {
  105. settingBytes, err := json.Marshal(setting)
  106. if err != nil {
  107. common.SysLog("failed to marshal setting: " + err.Error())
  108. return
  109. }
  110. user.Setting = string(settingBytes)
  111. }
  112. // 根据用户角色生成默认的边栏配置
  113. func generateDefaultSidebarConfigForRole(userRole int) string {
  114. defaultConfig := map[string]interface{}{}
  115. // 聊天区域 - 所有用户都可以访问
  116. defaultConfig["chat"] = map[string]interface{}{
  117. "enabled": true,
  118. "playground": true,
  119. "chat": true,
  120. }
  121. // 控制台区域 - 所有用户都可以访问
  122. defaultConfig["console"] = map[string]interface{}{
  123. "enabled": true,
  124. "detail": true,
  125. "token": true,
  126. "log": true,
  127. "midjourney": true,
  128. "task": true,
  129. }
  130. // 个人中心区域 - 所有用户都可以访问
  131. defaultConfig["personal"] = map[string]interface{}{
  132. "enabled": true,
  133. "topup": true,
  134. "personal": true,
  135. }
  136. // 管理员区域 - 根据角色决定
  137. if userRole == common.RoleAdminUser {
  138. // 管理员可以访问管理员区域,但不能访问系统设置
  139. defaultConfig["admin"] = map[string]interface{}{
  140. "enabled": true,
  141. "channel": true,
  142. "models": true,
  143. "redemption": true,
  144. "user": true,
  145. "setting": false, // 管理员不能访问系统设置
  146. }
  147. } else if userRole == common.RoleRootUser {
  148. // 超级管理员可以访问所有功能
  149. defaultConfig["admin"] = map[string]interface{}{
  150. "enabled": true,
  151. "channel": true,
  152. "models": true,
  153. "redemption": true,
  154. "user": true,
  155. "setting": true,
  156. }
  157. }
  158. // 普通用户不包含admin区域
  159. // 转换为JSON字符串
  160. configBytes, err := json.Marshal(defaultConfig)
  161. if err != nil {
  162. common.SysLog("生成默认边栏配置失败: " + err.Error())
  163. return ""
  164. }
  165. return string(configBytes)
  166. }
  167. // CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil
  168. func CheckUserExistOrDeleted(username string, email string) (bool, error) {
  169. var user User
  170. // err := DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
  171. // check email if empty
  172. var err error
  173. if email == "" {
  174. err = DB.Unscoped().First(&user, "username = ?", username).Error
  175. } else {
  176. err = DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
  177. }
  178. if err != nil {
  179. if errors.Is(err, gorm.ErrRecordNotFound) {
  180. // not exist, return false, nil
  181. return false, nil
  182. }
  183. // other error, return false, err
  184. return false, err
  185. }
  186. // exist, return true, nil
  187. return true, nil
  188. }
  189. func GetMaxUserId() int {
  190. var user User
  191. DB.Unscoped().Last(&user)
  192. return user.Id
  193. }
  194. // ApplySyncedQuota 将同步用户的 quota 替换为 synced_quota,
  195. // 使 API 返回的 quota 对所有用户类型都表示实际可用余额。
  196. func (u *User) ApplySyncedQuota() {
  197. if u.IsSyncedUser() {
  198. u.Quota = u.SyncedQuota
  199. }
  200. }
  201. func applySyncedUserQuota(users []*User) {
  202. for _, u := range users {
  203. u.ApplySyncedQuota()
  204. }
  205. }
  206. func GetAllUsers(pageInfo *common.PageInfo) (users []*User, total int64, err error) {
  207. // Start transaction
  208. tx := DB.Begin()
  209. if tx.Error != nil {
  210. return nil, 0, tx.Error
  211. }
  212. defer func() {
  213. if r := recover(); r != nil {
  214. tx.Rollback()
  215. }
  216. }()
  217. // Get total count within transaction
  218. err = tx.Unscoped().Model(&User{}).Count(&total).Error
  219. if err != nil {
  220. tx.Rollback()
  221. return nil, 0, err
  222. }
  223. // Get paginated users within same transaction
  224. err = tx.Unscoped().Order("id desc").Limit(pageInfo.GetPageSize()).Offset(pageInfo.GetStartIdx()).Omit("password").Find(&users).Error
  225. if err != nil {
  226. tx.Rollback()
  227. return nil, 0, err
  228. }
  229. // Commit transaction
  230. if err = tx.Commit().Error; err != nil {
  231. return nil, 0, err
  232. }
  233. applySyncedUserQuota(users)
  234. return users, total, nil
  235. }
  236. func SearchUsers(keyword string, group string, startIdx int, num int) ([]*User, int64, error) {
  237. var users []*User
  238. var total int64
  239. var err error
  240. // 开始事务
  241. tx := DB.Begin()
  242. if tx.Error != nil {
  243. return nil, 0, tx.Error
  244. }
  245. defer func() {
  246. if r := recover(); r != nil {
  247. tx.Rollback()
  248. }
  249. }()
  250. // 构建基础查询
  251. query := tx.Unscoped().Model(&User{})
  252. // 构建搜索条件
  253. likeCondition := "username LIKE ? OR email LIKE ? OR display_name LIKE ?"
  254. // 尝试将关键字转换为整数ID
  255. keywordInt, err := strconv.Atoi(keyword)
  256. if err == nil {
  257. // 如果是数字,同时搜索ID和其他字段
  258. likeCondition = "id = ? OR " + likeCondition
  259. if group != "" {
  260. query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
  261. keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
  262. } else {
  263. query = query.Where(likeCondition,
  264. keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
  265. }
  266. } else {
  267. // 非数字关键字,只搜索字符串字段
  268. if group != "" {
  269. query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
  270. "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
  271. } else {
  272. query = query.Where(likeCondition,
  273. "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
  274. }
  275. }
  276. // 获取总数
  277. err = query.Count(&total).Error
  278. if err != nil {
  279. tx.Rollback()
  280. return nil, 0, err
  281. }
  282. // 获取分页数据
  283. err = query.Omit("password").Order("id desc").Limit(num).Offset(startIdx).Find(&users).Error
  284. if err != nil {
  285. tx.Rollback()
  286. return nil, 0, err
  287. }
  288. // 提交事务
  289. if err = tx.Commit().Error; err != nil {
  290. return nil, 0, err
  291. }
  292. applySyncedUserQuota(users)
  293. return users, total, nil
  294. }
  295. func GetUserById(id int, selectAll bool) (*User, error) {
  296. if id == 0 {
  297. return nil, errors.New("id 为空!")
  298. }
  299. user := User{Id: id}
  300. var err error = nil
  301. if selectAll {
  302. err = DB.First(&user, "id = ?", id).Error
  303. } else {
  304. err = DB.Omit("password").First(&user, "id = ?", id).Error
  305. }
  306. return &user, err
  307. }
  308. func getUserByStringColumn(column string, value string) (*User, error) {
  309. if value == "" {
  310. return nil, nil
  311. }
  312. var user User
  313. result := DB.Where(column+" = ?", value).Limit(1).Find(&user)
  314. if result.Error != nil {
  315. return nil, result.Error
  316. }
  317. if result.RowsAffected == 0 {
  318. return nil, nil
  319. }
  320. return &user, nil
  321. }
  322. func GetUserByUsername(username string) (*User, error) {
  323. return getUserByStringColumn("username", username)
  324. }
  325. func GetUserByEmail(email string) (*User, error) {
  326. return getUserByStringColumn("email", email)
  327. }
  328. func GetUserByGitHubId(githubId string) (*User, error) {
  329. return getUserByStringColumn("github_id", githubId)
  330. }
  331. func GetUserByDiscordId(discordId string) (*User, error) {
  332. return getUserByStringColumn("discord_id", discordId)
  333. }
  334. func GetUserByOidcId(oidcId string) (*User, error) {
  335. return getUserByStringColumn("oidc_id", oidcId)
  336. }
  337. func GetUserByWeChatId(wechatId string) (*User, error) {
  338. return getUserByStringColumn("wechat_id", wechatId)
  339. }
  340. func GetUserByTelegramId(telegramId string) (*User, error) {
  341. return getUserByStringColumn("telegram_id", telegramId)
  342. }
  343. func GetUserByLinuxDOId(linuxDOId string) (*User, error) {
  344. return getUserByStringColumn("linux_do_id", linuxDOId)
  345. }
  346. func GetUserIdByAffCode(affCode string) (int, error) {
  347. if affCode == "" {
  348. return 0, errors.New("affCode 为空!")
  349. }
  350. var user User
  351. result := DB.Select("id").Where("aff_code = ?", affCode).Limit(1).Find(&user)
  352. if result.Error != nil {
  353. return 0, result.Error
  354. }
  355. if result.RowsAffected == 0 {
  356. return 0, gorm.ErrRecordNotFound
  357. }
  358. return user.Id, nil
  359. }
  360. func DeleteUserById(id int) (err error) {
  361. if id == 0 {
  362. return errors.New("id 为空!")
  363. }
  364. user := User{Id: id}
  365. return user.Delete()
  366. }
  367. func HardDeleteUserById(id int) error {
  368. if id == 0 {
  369. return errors.New("id 为空!")
  370. }
  371. err := DB.Unscoped().Delete(&User{}, "id = ?", id).Error
  372. return err
  373. }
  374. func inviteUser(inviterId int) (err error) {
  375. user, err := GetUserById(inviterId, true)
  376. if err != nil {
  377. return err
  378. }
  379. user.AffCount++
  380. user.AffQuota += common.QuotaForInviter
  381. user.AffHistoryQuota += common.QuotaForInviter
  382. return DB.Save(user).Error
  383. }
  384. func (user *User) TransferAffQuotaToQuota(quota int) error {
  385. // 检查quota是否小于最小额度
  386. if float64(quota) < common.QuotaPerUnit {
  387. return fmt.Errorf("转移额度最小为%s!", logger.LogQuota(int(common.QuotaPerUnit)))
  388. }
  389. // 开始数据库事务
  390. tx := DB.Begin()
  391. if tx.Error != nil {
  392. return tx.Error
  393. }
  394. defer tx.Rollback() // 确保在函数退出时事务能回滚
  395. // 加锁查询用户以确保数据一致性
  396. err := tx.Set("gorm:query_option", "FOR UPDATE").First(&user, user.Id).Error
  397. if err != nil {
  398. return err
  399. }
  400. // 再次检查用户的AffQuota是否足够
  401. if user.AffQuota < quota {
  402. return errors.New("邀请额度不足!")
  403. }
  404. // 更新用户额度
  405. user.AffQuota -= quota
  406. user.Quota += quota
  407. // 保存用户状态
  408. if err := tx.Save(user).Error; err != nil {
  409. return err
  410. }
  411. // 提交事务
  412. return tx.Commit().Error
  413. }
  414. func (user *User) Insert(inviterId int) error {
  415. var err error
  416. if user.Password != "" {
  417. user.Password, err = common.Password2Hash(user.Password)
  418. if err != nil {
  419. return err
  420. }
  421. }
  422. matchedQuota := MatchEmailQuotaRule(user.Email)
  423. if matchedQuota >= 0 {
  424. user.Quota = int(matchedQuota)
  425. } else {
  426. user.Quota = common.QuotaForNewUser
  427. }
  428. //user.SetAccessToken(common.GetUUID())
  429. user.AffCode = common.GetRandomString(4)
  430. user.CreatedAt = time.Now().Unix()
  431. // 初始化用户设置,包括默认的边栏配置
  432. if user.Setting == "" {
  433. defaultSetting := dto.UserSetting{}
  434. // 这里暂时不设置SidebarModules,因为需要在用户创建后根据角色设置
  435. user.SetSetting(defaultSetting)
  436. }
  437. result := DB.Create(user)
  438. if result.Error != nil {
  439. return result.Error
  440. }
  441. // 用户创建成功后,根据角色初始化边栏配置
  442. // 需要重新获取用户以确保有正确的ID和Role
  443. var createdUser User
  444. if err := DB.Where("username = ?", user.Username).First(&createdUser).Error; err == nil {
  445. // 生成基于角色的默认边栏配置
  446. defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
  447. if defaultSidebarConfig != "" {
  448. currentSetting := createdUser.GetSetting()
  449. currentSetting.SidebarModules = defaultSidebarConfig
  450. createdUser.SetSetting(currentSetting)
  451. createdUser.Update(false)
  452. common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
  453. }
  454. }
  455. if user.Quota > 0 {
  456. if MatchEmailQuotaRule(user.Email) >= 0 {
  457. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s(邮箱后缀规则匹配)", logger.LogQuota(user.Quota)))
  458. } else {
  459. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(user.Quota)))
  460. }
  461. }
  462. if inviterId != 0 {
  463. if common.QuotaForInvitee > 0 {
  464. _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
  465. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
  466. }
  467. if common.QuotaForInviter > 0 {
  468. //_ = IncreaseUserQuota(inviterId, common.QuotaForInviter)
  469. RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
  470. _ = inviteUser(inviterId)
  471. }
  472. }
  473. return nil
  474. }
  475. // InsertWithTx inserts a new user within an existing transaction.
  476. // This is used for OAuth registration where user creation and binding need to be atomic.
  477. // Post-creation tasks (sidebar config, logs, inviter rewards) are handled after the transaction commits.
  478. func (user *User) InsertWithTx(tx *gorm.DB, inviterId int) error {
  479. var err error
  480. if user.Password != "" {
  481. user.Password, err = common.Password2Hash(user.Password)
  482. if err != nil {
  483. return err
  484. }
  485. }
  486. matchedQuota := MatchEmailQuotaRule(user.Email)
  487. if matchedQuota >= 0 {
  488. user.Quota = int(matchedQuota)
  489. } else {
  490. user.Quota = common.QuotaForNewUser
  491. }
  492. user.AffCode = common.GetRandomString(4)
  493. user.CreatedAt = time.Now().Unix()
  494. // 初始化用户设置
  495. if user.Setting == "" {
  496. defaultSetting := dto.UserSetting{}
  497. user.SetSetting(defaultSetting)
  498. }
  499. result := tx.Create(user)
  500. if result.Error != nil {
  501. return result.Error
  502. }
  503. return nil
  504. }
  505. // InsertImportedUserWithHashedPassword 专用于迁移导入,直接写入已有密码哈希。
  506. func InsertImportedUserWithHashedPassword(user *User) error {
  507. return InsertImportedUserWithHashedPasswordTx(DB, user)
  508. }
  509. func InsertImportedUserWithHashedPasswordTx(tx *gorm.DB, user *User) error {
  510. if user == nil {
  511. return errors.New("user is nil")
  512. }
  513. if user.Id != 0 {
  514. return errors.New("InsertImportedUserWithHashedPassword: Id must not be set, let DB auto-assign")
  515. }
  516. if user.AffCode == "" {
  517. user.AffCode = common.GetRandomString(4)
  518. }
  519. if user.CreatedAt == 0 {
  520. user.CreatedAt = common.GetTimestamp()
  521. }
  522. if user.Setting == "" {
  523. defaultSetting := dto.UserSetting{}
  524. user.SetSetting(defaultSetting)
  525. }
  526. return tx.Create(user).Error
  527. }
  528. // FinalizeOAuthUserCreation performs post-transaction tasks for OAuth user creation.
  529. // This should be called after the transaction commits successfully.
  530. func (user *User) FinalizeOAuthUserCreation(inviterId int) {
  531. // 用户创建成功后,根据角色初始化边栏配置
  532. var createdUser User
  533. if err := DB.Where("id = ?", user.Id).First(&createdUser).Error; err == nil {
  534. defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
  535. if defaultSidebarConfig != "" {
  536. currentSetting := createdUser.GetSetting()
  537. currentSetting.SidebarModules = defaultSidebarConfig
  538. createdUser.SetSetting(currentSetting)
  539. createdUser.Update(false)
  540. common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
  541. }
  542. }
  543. if user.Quota > 0 {
  544. if matchedQuota := MatchEmailQuotaRule(user.Email); matchedQuota >= 0 {
  545. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s(邮箱后缀规则匹配)", logger.LogQuota(user.Quota)))
  546. } else {
  547. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(user.Quota)))
  548. }
  549. }
  550. if inviterId != 0 {
  551. if common.QuotaForInvitee > 0 {
  552. _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
  553. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
  554. }
  555. if common.QuotaForInviter > 0 {
  556. RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
  557. _ = inviteUser(inviterId)
  558. }
  559. }
  560. }
  561. func (user *User) Update(updatePassword bool) error {
  562. var err error
  563. if updatePassword {
  564. user.Password, err = common.Password2Hash(user.Password)
  565. if err != nil {
  566. return err
  567. }
  568. }
  569. newUser := *user
  570. DB.First(&user, user.Id)
  571. if err = DB.Model(user).Updates(newUser).Error; err != nil {
  572. return err
  573. }
  574. // Update cache
  575. return updateUserCache(*user)
  576. }
  577. func (user *User) Edit(updatePassword bool) error {
  578. var err error
  579. if updatePassword {
  580. user.Password, err = common.Password2Hash(user.Password)
  581. if err != nil {
  582. return err
  583. }
  584. }
  585. newUser := *user
  586. updates := map[string]interface{}{
  587. "username": newUser.Username,
  588. "display_name": newUser.DisplayName,
  589. "group": newUser.Group,
  590. "quota": newUser.Quota,
  591. "remark": newUser.Remark,
  592. "capture_relay": newUser.CaptureRelay,
  593. }
  594. if updatePassword {
  595. updates["password"] = newUser.Password
  596. }
  597. DB.First(&user, user.Id)
  598. if err = DB.Model(user).Updates(updates).Error; err != nil {
  599. return err
  600. }
  601. // Update cache
  602. return updateUserCache(*user)
  603. }
  604. func (user *User) ClearBinding(bindingType string) error {
  605. if user.Id == 0 {
  606. return errors.New("user id is empty")
  607. }
  608. bindingColumnMap := map[string]string{
  609. "email": "email",
  610. "github": "github_id",
  611. "discord": "discord_id",
  612. "oidc": "oidc_id",
  613. "wechat": "wechat_id",
  614. "telegram": "telegram_id",
  615. "linuxdo": "linux_do_id",
  616. }
  617. column, ok := bindingColumnMap[bindingType]
  618. if !ok {
  619. return errors.New("invalid binding type")
  620. }
  621. if err := DB.Model(&User{}).Where("id = ?", user.Id).Update(column, "").Error; err != nil {
  622. return err
  623. }
  624. if err := DB.Where("id = ?", user.Id).First(user).Error; err != nil {
  625. return err
  626. }
  627. return updateUserCache(*user)
  628. }
  629. func (user *User) Delete() error {
  630. if user.Id == 0 {
  631. return errors.New("id 为空!")
  632. }
  633. if err := DB.Delete(user).Error; err != nil {
  634. return err
  635. }
  636. // 清除缓存
  637. return invalidateUserCache(user.Id)
  638. }
  639. func (user *User) HardDelete() error {
  640. if user.Id == 0 {
  641. return errors.New("id 为空!")
  642. }
  643. err := DB.Unscoped().Delete(user).Error
  644. return err
  645. }
  646. // ValidateAndFill check password & user status
  647. func (user *User) ValidateAndFill() (err error) {
  648. // When querying with struct, GORM will only query with non-zero fields,
  649. // that means if your field's value is 0, '', false or other zero values,
  650. // it won't be used to build query conditions
  651. password := user.Password
  652. username := strings.TrimSpace(user.Username)
  653. if username == "" || password == "" {
  654. return errors.New("用户名或密码为空")
  655. }
  656. // find buy username or email
  657. DB.Where("username = ? OR email = ?", username, username).First(user)
  658. okay := common.ValidatePasswordAndHash(password, user.Password)
  659. if !okay || user.Status != common.UserStatusEnabled {
  660. return errors.New("用户名或密码错误,或用户已被封禁")
  661. }
  662. return nil
  663. }
  664. func (user *User) FillUserById() error {
  665. if user.Id == 0 {
  666. return errors.New("id 为空!")
  667. }
  668. DB.Where(User{Id: user.Id}).First(user)
  669. return nil
  670. }
  671. func (user *User) FillUserByEmail() error {
  672. if user.Email == "" {
  673. return errors.New("email 为空!")
  674. }
  675. DB.Where(User{Email: user.Email}).First(user)
  676. return nil
  677. }
  678. func (user *User) FillUserByGitHubId() error {
  679. if user.GitHubId == "" {
  680. return errors.New("GitHub id 为空!")
  681. }
  682. DB.Where(User{GitHubId: user.GitHubId}).First(user)
  683. return nil
  684. }
  685. // UpdateGitHubId updates the user's GitHub ID (used for migration from login to numeric ID)
  686. func (user *User) UpdateGitHubId(newGitHubId string) error {
  687. if user.Id == 0 {
  688. return errors.New("user id is empty")
  689. }
  690. return DB.Model(user).Update("github_id", newGitHubId).Error
  691. }
  692. func (user *User) FillUserByDiscordId() error {
  693. if user.DiscordId == "" {
  694. return errors.New("discord id 为空!")
  695. }
  696. DB.Where(User{DiscordId: user.DiscordId}).First(user)
  697. return nil
  698. }
  699. func (user *User) FillUserByOidcId() error {
  700. if user.OidcId == "" {
  701. return errors.New("oidc id 为空!")
  702. }
  703. DB.Where(User{OidcId: user.OidcId}).First(user)
  704. return nil
  705. }
  706. func (user *User) FillUserByWeChatId() error {
  707. if user.WeChatId == "" {
  708. return errors.New("WeChat id 为空!")
  709. }
  710. DB.Where(User{WeChatId: user.WeChatId}).First(user)
  711. return nil
  712. }
  713. func (user *User) FillUserByTelegramId() error {
  714. if user.TelegramId == "" {
  715. return errors.New("Telegram id 为空!")
  716. }
  717. err := DB.Where(User{TelegramId: user.TelegramId}).First(user).Error
  718. if errors.Is(err, gorm.ErrRecordNotFound) {
  719. return errors.New("该 Telegram 账户未绑定")
  720. }
  721. return nil
  722. }
  723. func IsEmailAlreadyTaken(email string) bool {
  724. return DB.Unscoped().Where("email = ?", email).Find(&User{}).RowsAffected == 1
  725. }
  726. func IsWeChatIdAlreadyTaken(wechatId string) bool {
  727. return DB.Unscoped().Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1
  728. }
  729. func IsGitHubIdAlreadyTaken(githubId string) bool {
  730. return DB.Unscoped().Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1
  731. }
  732. func IsDiscordIdAlreadyTaken(discordId string) bool {
  733. return DB.Unscoped().Where("discord_id = ?", discordId).Find(&User{}).RowsAffected == 1
  734. }
  735. func IsOidcIdAlreadyTaken(oidcId string) bool {
  736. return DB.Where("oidc_id = ?", oidcId).Find(&User{}).RowsAffected == 1
  737. }
  738. func IsTelegramIdAlreadyTaken(telegramId string) bool {
  739. return DB.Unscoped().Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1
  740. }
  741. func ResetUserPasswordByEmail(email string, password string) error {
  742. if email == "" || password == "" {
  743. return errors.New("邮箱地址或密码为空!")
  744. }
  745. hashedPassword, err := common.Password2Hash(password)
  746. if err != nil {
  747. return err
  748. }
  749. err = DB.Model(&User{}).Where("email = ?", email).Update("password", hashedPassword).Error
  750. return err
  751. }
  752. func IsAdmin(userId int) bool {
  753. if userId == 0 {
  754. return false
  755. }
  756. var user User
  757. err := DB.Where("id = ?", userId).Select("role").Find(&user).Error
  758. if err != nil {
  759. common.SysLog("no such user " + err.Error())
  760. return false
  761. }
  762. return user.Role >= common.RoleAdminUser
  763. }
  764. //// IsUserEnabled checks user status from Redis first, falls back to DB if needed
  765. //func IsUserEnabled(id int, fromDB bool) (status bool, err error) {
  766. // defer func() {
  767. // // Update Redis cache asynchronously on successful DB read
  768. // if shouldUpdateRedis(fromDB, err) {
  769. // gopool.Go(func() {
  770. // if err := updateUserStatusCache(id, status); err != nil {
  771. // common.SysError("failed to update user status cache: " + err.Error())
  772. // }
  773. // })
  774. // }
  775. // }()
  776. // if !fromDB && common.RedisEnabled {
  777. // // Try Redis first
  778. // status, err := getUserStatusCache(id)
  779. // if err == nil {
  780. // return status == common.UserStatusEnabled, nil
  781. // }
  782. // // Don't return error - fall through to DB
  783. // }
  784. // fromDB = true
  785. // var user User
  786. // err = DB.Where("id = ?", id).Select("status").Find(&user).Error
  787. // if err != nil {
  788. // return false, err
  789. // }
  790. //
  791. // return user.Status == common.UserStatusEnabled, nil
  792. //}
  793. func ValidateAccessToken(token string) (user *User) {
  794. if token == "" {
  795. return nil
  796. }
  797. token = strings.Replace(token, "Bearer ", "", 1)
  798. user = &User{}
  799. if DB.Where("access_token = ?", token).First(user).RowsAffected == 1 {
  800. return user
  801. }
  802. return nil
  803. }
  804. // GetUserQuota gets quota from Redis first, falls back to DB if needed
  805. // 同步用户返回 SyncedQuota,本地用户返回 Quota
  806. func GetUserQuota(id int, fromDB bool) (quota int, err error) {
  807. if !fromDB && common.RedisEnabled {
  808. quota, err := getUserQuotaCache(id)
  809. if err == nil {
  810. return quota, nil
  811. }
  812. // Don't return error - fall through to DB
  813. }
  814. fromDB = true
  815. // 查询用户的 source 和额度字段
  816. var user User
  817. err = DB.Model(&User{}).Select("source", "quota", "synced_quota").Where("id = ?", id).First(&user).Error
  818. if err != nil {
  819. return 0, err
  820. }
  821. // 同步用户返回 SyncedQuota,本地用户返回 Quota
  822. if user.IsSyncedUser() {
  823. quota = user.SyncedQuota
  824. } else {
  825. quota = user.Quota
  826. }
  827. // Update Redis cache asynchronously if enabled
  828. if common.RedisEnabled {
  829. gopool.Go(func() {
  830. if err := updateUserQuotaCache(id, quota); err != nil {
  831. common.SysLog("failed to update user quota cache: " + err.Error())
  832. }
  833. })
  834. }
  835. return quota, nil
  836. }
  837. func GetUserUsedQuota(id int) (quota int, err error) {
  838. err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error
  839. return quota, err
  840. }
  841. func GetUserEmail(id int) (email string, err error) {
  842. err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
  843. return email, err
  844. }
  845. // GetUserGroup gets group from Redis first, falls back to DB if needed
  846. func GetUserGroup(id int, fromDB bool) (group string, err error) {
  847. defer func() {
  848. // Update Redis cache asynchronously on successful DB read
  849. if shouldUpdateRedis(fromDB, err) {
  850. gopool.Go(func() {
  851. if err := updateUserGroupCache(id, group); err != nil {
  852. common.SysLog("failed to update user group cache: " + err.Error())
  853. }
  854. })
  855. }
  856. }()
  857. if !fromDB && common.RedisEnabled {
  858. group, err := getUserGroupCache(id)
  859. if err == nil {
  860. return group, nil
  861. }
  862. // Don't return error - fall through to DB
  863. }
  864. fromDB = true
  865. err = DB.Model(&User{}).Where("id = ?", id).Select(commonGroupCol).Find(&group).Error
  866. if err != nil {
  867. return "", err
  868. }
  869. return group, nil
  870. }
  871. // GetUserSetting gets setting from Redis first, falls back to DB if needed
  872. func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error) {
  873. var setting string
  874. defer func() {
  875. // Update Redis cache asynchronously on successful DB read
  876. if shouldUpdateRedis(fromDB, err) {
  877. gopool.Go(func() {
  878. if err := updateUserSettingCache(id, setting); err != nil {
  879. common.SysLog("failed to update user setting cache: " + err.Error())
  880. }
  881. })
  882. }
  883. }()
  884. if !fromDB && common.RedisEnabled {
  885. setting, err := getUserSettingCache(id)
  886. if err == nil {
  887. return setting, nil
  888. }
  889. // Don't return error - fall through to DB
  890. }
  891. fromDB = true
  892. err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&setting).Error
  893. if err != nil {
  894. return settingMap, err
  895. }
  896. userBase := &UserBase{
  897. Setting: setting,
  898. }
  899. return userBase.GetSetting(), nil
  900. }
  901. func IncreaseUserQuota(id int, quota int, db bool) (err error) {
  902. if quota < 0 {
  903. return errors.New("quota 不能为负数!")
  904. }
  905. gopool.Go(func() {
  906. err := cacheIncrUserQuota(id, int64(quota))
  907. if err != nil {
  908. common.SysLog("failed to increase user quota: " + err.Error())
  909. }
  910. })
  911. if !db && common.BatchUpdateEnabled {
  912. addNewRecord(BatchUpdateTypeUserQuota, id, quota)
  913. return nil
  914. }
  915. err = increaseUserQuota(id, quota)
  916. if err != nil {
  917. return err
  918. }
  919. // 触发余额更新回调(如果有注册)
  920. if quotaUpdateCallback != nil {
  921. newQuota, _ := GetUserQuota(id, true)
  922. // 使用 defer/recover 保护回调,防止 panic 影响主流程
  923. func() {
  924. defer func() {
  925. if r := recover(); r != nil {
  926. common.SysError(fmt.Sprintf("quota update callback panic: %v", r))
  927. }
  928. }()
  929. quotaUpdateCallback(id, newQuota)
  930. }()
  931. }
  932. return nil
  933. }
  934. func increaseUserQuota(id int, quota int) (err error) {
  935. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
  936. if err != nil {
  937. return err
  938. }
  939. return err
  940. }
  941. func DecreaseUserQuota(id int, quota int) (err error) {
  942. if quota < 0 {
  943. return errors.New("quota 不能为负数!")
  944. }
  945. gopool.Go(func() {
  946. err := cacheDecrUserQuota(id, int64(quota))
  947. if err != nil {
  948. common.SysLog("failed to decrease user quota: " + err.Error())
  949. }
  950. })
  951. if common.BatchUpdateEnabled {
  952. addNewRecord(BatchUpdateTypeUserQuota, id, -quota)
  953. return nil
  954. }
  955. return decreaseUserQuota(id, quota)
  956. }
  957. func decreaseUserQuota(id int, quota int) (err error) {
  958. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
  959. if err != nil {
  960. return err
  961. }
  962. return err
  963. }
  964. func DeltaUpdateUserQuota(id int, delta int) (err error) {
  965. if delta == 0 {
  966. return nil
  967. }
  968. if delta > 0 {
  969. return IncreaseUserQuota(id, delta, false)
  970. } else {
  971. return DecreaseUserQuota(id, -delta)
  972. }
  973. }
  974. func ListLocalUsersForMigration(page, pageSize int) ([]*User, int64, error) {
  975. if page < 0 {
  976. page = 0
  977. }
  978. if pageSize <= 0 {
  979. pageSize = 20
  980. }
  981. var users []*User
  982. var total int64
  983. query := DB.Model(&User{}).Where("source = ?", common.UserSourceLocal)
  984. if err := query.Count(&total).Error; err != nil {
  985. return nil, 0, err
  986. }
  987. err := query.Order("id asc").Offset(page * pageSize).Limit(pageSize).Find(&users).Error
  988. return users, total, err
  989. }
  990. func HasOVSyncedCopyByRemoteId(cnUserId int) (bool, error) {
  991. var count int64
  992. err := DB.Model(&User{}).
  993. Where("remote_user_id = ? AND source = ?", cnUserId, common.UserSourceSynced).
  994. Count(&count).Error
  995. return count > 0, err
  996. }
  997. func ConvertUserToSynced(userId, remoteUserId, syncedQuota int) error {
  998. var user User
  999. if err := DB.Where("id = ?", userId).First(&user).Error; err != nil {
  1000. return err
  1001. }
  1002. now := common.GetTimestamp()
  1003. if user.Source == common.UserSourceSynced {
  1004. if user.RemoteUserId != remoteUserId {
  1005. return fmt.Errorf("user %d is already synced to remote_user_id=%d, cannot change to %d",
  1006. userId, user.RemoteUserId, remoteUserId)
  1007. }
  1008. if user.SyncedQuota == syncedQuota {
  1009. return nil
  1010. }
  1011. return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]any{
  1012. "synced_quota": syncedQuota,
  1013. "last_sync_at": now,
  1014. }).Error
  1015. }
  1016. return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]any{
  1017. "source": common.UserSourceSynced,
  1018. "remote_user_id": remoteUserId,
  1019. "synced_quota": syncedQuota,
  1020. "last_sync_at": now,
  1021. }).Error
  1022. }
  1023. // GetSyncedUsers 获取所有同步用户
  1024. func GetSyncedUsers() []User {
  1025. var users []User
  1026. DB.Where("source = ?", common.UserSourceSynced).Find(&users)
  1027. return users
  1028. }
  1029. // UpdateSyncedQuota 更新同步用户的 synced_quota
  1030. func UpdateSyncedQuota(userId int, quota int) error {
  1031. return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]interface{}{
  1032. "synced_quota": quota,
  1033. "last_sync_at": time.Now().Unix(),
  1034. }).Error
  1035. }
  1036. // AtomicDecreaseSyncedQuota 原子扣减 synced_quota,检查余额并扣减在一条 SQL 中完成。
  1037. // 语义:扣减后 synced_quota 必须 >= threshold,即 synced_quota >= amount + threshold。
  1038. // 返回 (当前余额, 是否成功, 错误)。
  1039. func AtomicDecreaseSyncedQuota(userId int, amount int, threshold int) (int, bool, error) {
  1040. result := DB.Model(&User{}).
  1041. Where("id = ? AND synced_quota >= ?", userId, amount+threshold).
  1042. Update("synced_quota", gorm.Expr("synced_quota - ?", amount))
  1043. if result.Error != nil {
  1044. return 0, false, result.Error
  1045. }
  1046. if result.RowsAffected == 0 {
  1047. var quota int
  1048. DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan(&quota)
  1049. return quota, false, nil
  1050. }
  1051. var newQuota int
  1052. DB.Model(&User{}).Where("id = ?", userId).Select("synced_quota").Scan(&newQuota)
  1053. return newQuota, true, nil
  1054. }
  1055. // IncreaseSyncedQuota 原子增加 synced_quota(用于退款)
  1056. func IncreaseSyncedQuota(userId int, amount int) error {
  1057. return DB.Model(&User{}).
  1058. Where("id = ?", userId).
  1059. Update("synced_quota", gorm.Expr("synced_quota + ?", amount)).Error
  1060. }
  1061. // AdjustSyncedQuota 原子调整 synced_quota(增量更新,避免并发覆盖)
  1062. // delta > 0 表示额外扣减,delta < 0 表示退还。
  1063. func AdjustSyncedQuota(userId int, delta int) error {
  1064. return DB.Model(&User{}).Where("id = ?", userId).
  1065. Update("synced_quota", gorm.Expr("synced_quota - ?", delta)).Error
  1066. }
  1067. //func GetRootUserEmail() (email string) {
  1068. // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email)
  1069. // return email
  1070. //}
  1071. func GetRootUser() (user *User) {
  1072. DB.Where("role = ?", common.RoleRootUser).First(&user)
  1073. return user
  1074. }
  1075. func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
  1076. if common.BatchUpdateEnabled {
  1077. addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
  1078. addNewRecord(BatchUpdateTypeRequestCount, id, 1)
  1079. return
  1080. }
  1081. updateUserUsedQuotaAndRequestCount(id, quota, 1)
  1082. }
  1083. func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
  1084. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  1085. map[string]interface{}{
  1086. "used_quota": gorm.Expr("used_quota + ?", quota),
  1087. "request_count": gorm.Expr("request_count + ?", count),
  1088. },
  1089. ).Error
  1090. if err != nil {
  1091. common.SysLog("failed to update user used quota and request count: " + err.Error())
  1092. return
  1093. }
  1094. //// 更新缓存
  1095. //if err := invalidateUserCache(id); err != nil {
  1096. // common.SysError("failed to invalidate user cache: " + err.Error())
  1097. //}
  1098. }
  1099. func updateUserUsedQuota(id int, quota int) {
  1100. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  1101. map[string]interface{}{
  1102. "used_quota": gorm.Expr("used_quota + ?", quota),
  1103. },
  1104. ).Error
  1105. if err != nil {
  1106. common.SysLog("failed to update user used quota: " + err.Error())
  1107. }
  1108. }
  1109. func updateUserRequestCount(id int, count int) {
  1110. err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error
  1111. if err != nil {
  1112. common.SysLog("failed to update user request count: " + err.Error())
  1113. }
  1114. }
  1115. // GetUsernameById gets username from Redis first, falls back to DB if needed
  1116. func GetUsernameById(id int, fromDB bool) (username string, err error) {
  1117. defer func() {
  1118. // Update Redis cache asynchronously on successful DB read
  1119. if shouldUpdateRedis(fromDB, err) {
  1120. gopool.Go(func() {
  1121. if err := updateUserNameCache(id, username); err != nil {
  1122. common.SysLog("failed to update user name cache: " + err.Error())
  1123. }
  1124. })
  1125. }
  1126. }()
  1127. if !fromDB && common.RedisEnabled {
  1128. username, err := getUserNameCache(id)
  1129. if err == nil {
  1130. return username, nil
  1131. }
  1132. // Don't return error - fall through to DB
  1133. }
  1134. fromDB = true
  1135. err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error
  1136. if err != nil {
  1137. return "", err
  1138. }
  1139. return username, nil
  1140. }
  1141. func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
  1142. var user User
  1143. err := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).First(&user).Error
  1144. return !errors.Is(err, gorm.ErrRecordNotFound)
  1145. }
  1146. func (user *User) FillUserByLinuxDOId() error {
  1147. if user.LinuxDOId == "" {
  1148. return errors.New("linux do id is empty")
  1149. }
  1150. err := DB.Where("linux_do_id = ?", user.LinuxDOId).First(user).Error
  1151. return err
  1152. }
  1153. func RootUserExists() bool {
  1154. var user User
  1155. err := DB.Where("role = ?", common.RoleRootUser).First(&user).Error
  1156. if err != nil {
  1157. return false
  1158. }
  1159. return true
  1160. }