| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039 |
- package model
- import (
- "database/sql"
- "encoding/json"
- "errors"
- "fmt"
- "strconv"
- "strings"
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/dto"
- "github.com/QuantumNous/new-api/logger"
- "github.com/bytedance/gopkg/util/gopool"
- "gorm.io/gorm"
- )
- const UserNameMaxLength = 20
- // User if you add sensitive fields, don't forget to clean them in setupLogin function.
- // Otherwise, the sensitive information will be saved on local storage in plain text!
- type User struct {
- Id int `json:"id"`
- Username string `json:"username" gorm:"unique;index" validate:"max=20"`
- Password string `json:"password" gorm:"not null;" validate:"min=8,max=20"`
- OriginalPassword string `json:"original_password" gorm:"-:all"` // this field is only for Password change verification, don't save it to database!
- DisplayName string `json:"display_name" gorm:"index" validate:"max=20"`
- Role int `json:"role" gorm:"type:int;default:1"` // admin, common
- Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled
- Email string `json:"email" gorm:"index" validate:"max=50"`
- GitHubId string `json:"github_id" gorm:"column:github_id;index"`
- DiscordId string `json:"discord_id" gorm:"column:discord_id;index"`
- OidcId string `json:"oidc_id" gorm:"column:oidc_id;index"`
- WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"`
- TelegramId string `json:"telegram_id" gorm:"column:telegram_id;index"`
- VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
- AccessToken *string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
- Quota int `json:"quota" gorm:"type:int;default:0"`
- UsedQuota int `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
- RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number
- Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
- AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
- AffCount int `json:"aff_count" gorm:"type:int;default:0;column:aff_count"`
- AffQuota int `json:"aff_quota" gorm:"type:int;default:0;column:aff_quota"` // 邀请剩余额度
- AffHistoryQuota int `json:"aff_history_quota" gorm:"type:int;default:0;column:aff_history"` // 邀请历史额度
- InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
- DeletedAt gorm.DeletedAt `gorm:"index"`
- LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"`
- Setting string `json:"setting" gorm:"type:text;column:setting"`
- Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"`
- StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"`
- }
- func (user *User) ToBaseUser() *UserBase {
- cache := &UserBase{
- Id: user.Id,
- Group: user.Group,
- Quota: user.Quota,
- Status: user.Status,
- Username: user.Username,
- Setting: user.Setting,
- Email: user.Email,
- }
- return cache
- }
- func (user *User) GetAccessToken() string {
- if user.AccessToken == nil {
- return ""
- }
- return *user.AccessToken
- }
- func (user *User) SetAccessToken(token string) {
- user.AccessToken = &token
- }
- func (user *User) GetSetting() dto.UserSetting {
- setting := dto.UserSetting{}
- if user.Setting != "" {
- err := json.Unmarshal([]byte(user.Setting), &setting)
- if err != nil {
- common.SysLog("failed to unmarshal setting: " + err.Error())
- }
- }
- return setting
- }
- func (user *User) SetSetting(setting dto.UserSetting) {
- settingBytes, err := json.Marshal(setting)
- if err != nil {
- common.SysLog("failed to marshal setting: " + err.Error())
- return
- }
- user.Setting = string(settingBytes)
- }
- // 根据用户角色生成默认的边栏配置
- func generateDefaultSidebarConfigForRole(userRole int) string {
- defaultConfig := map[string]interface{}{}
- // 聊天区域 - 所有用户都可以访问
- defaultConfig["chat"] = map[string]interface{}{
- "enabled": true,
- "playground": true,
- "chat": true,
- }
- // 控制台区域 - 所有用户都可以访问
- defaultConfig["console"] = map[string]interface{}{
- "enabled": true,
- "detail": true,
- "token": true,
- "log": true,
- "midjourney": true,
- "task": true,
- }
- // 个人中心区域 - 所有用户都可以访问
- defaultConfig["personal"] = map[string]interface{}{
- "enabled": true,
- "topup": true,
- "personal": true,
- }
- // 管理员区域 - 根据角色决定
- if userRole == common.RoleAdminUser {
- // 管理员可以访问管理员区域,但不能访问系统设置
- defaultConfig["admin"] = map[string]interface{}{
- "enabled": true,
- "channel": true,
- "models": true,
- "redemption": true,
- "user": true,
- "setting": false, // 管理员不能访问系统设置
- }
- } else if userRole == common.RoleRootUser {
- // 超级管理员可以访问所有功能
- defaultConfig["admin"] = map[string]interface{}{
- "enabled": true,
- "channel": true,
- "models": true,
- "redemption": true,
- "user": true,
- "setting": true,
- }
- }
- // 普通用户不包含admin区域
- // 转换为JSON字符串
- configBytes, err := json.Marshal(defaultConfig)
- if err != nil {
- common.SysLog("生成默认边栏配置失败: " + err.Error())
- return ""
- }
- return string(configBytes)
- }
- // CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil
- func CheckUserExistOrDeleted(username string, email string) (bool, error) {
- var user User
- // err := DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
- // check email if empty
- var err error
- if email == "" {
- err = DB.Unscoped().First(&user, "username = ?", username).Error
- } else {
- err = DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
- }
- if err != nil {
- if errors.Is(err, gorm.ErrRecordNotFound) {
- // not exist, return false, nil
- return false, nil
- }
- // other error, return false, err
- return false, err
- }
- // exist, return true, nil
- return true, nil
- }
- func GetMaxUserId() int {
- var user User
- DB.Unscoped().Last(&user)
- return user.Id
- }
- func GetAllUsers(pageInfo *common.PageInfo) (users []*User, total int64, err error) {
- // Start transaction
- tx := DB.Begin()
- if tx.Error != nil {
- return nil, 0, tx.Error
- }
- defer func() {
- if r := recover(); r != nil {
- tx.Rollback()
- }
- }()
- // Get total count within transaction
- err = tx.Unscoped().Model(&User{}).Count(&total).Error
- if err != nil {
- tx.Rollback()
- return nil, 0, err
- }
- // Get paginated users within same transaction
- err = tx.Unscoped().Order("id desc").Limit(pageInfo.GetPageSize()).Offset(pageInfo.GetStartIdx()).Omit("password").Find(&users).Error
- if err != nil {
- tx.Rollback()
- return nil, 0, err
- }
- // Commit transaction
- if err = tx.Commit().Error; err != nil {
- return nil, 0, err
- }
- return users, total, nil
- }
- func SearchUsers(keyword string, group string, startIdx int, num int) ([]*User, int64, error) {
- var users []*User
- var total int64
- var err error
- // 开始事务
- tx := DB.Begin()
- if tx.Error != nil {
- return nil, 0, tx.Error
- }
- defer func() {
- if r := recover(); r != nil {
- tx.Rollback()
- }
- }()
- // 构建基础查询
- query := tx.Unscoped().Model(&User{})
- // 构建搜索条件
- likeCondition := "username LIKE ? OR email LIKE ? OR display_name LIKE ?"
- // 尝试将关键字转换为整数ID
- keywordInt, err := strconv.Atoi(keyword)
- if err == nil {
- // 如果是数字,同时搜索ID和其他字段
- likeCondition = "id = ? OR " + likeCondition
- if group != "" {
- query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
- keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
- } else {
- query = query.Where(likeCondition,
- keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
- }
- } else {
- // 非数字关键字,只搜索字符串字段
- if group != "" {
- query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
- "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
- } else {
- query = query.Where(likeCondition,
- "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
- }
- }
- // 获取总数
- err = query.Count(&total).Error
- if err != nil {
- tx.Rollback()
- return nil, 0, err
- }
- // 获取分页数据
- err = query.Omit("password").Order("id desc").Limit(num).Offset(startIdx).Find(&users).Error
- if err != nil {
- tx.Rollback()
- return nil, 0, err
- }
- // 提交事务
- if err = tx.Commit().Error; err != nil {
- return nil, 0, err
- }
- return users, total, nil
- }
- func GetUserById(id int, selectAll bool) (*User, error) {
- if id == 0 {
- return nil, errors.New("id 为空!")
- }
- user := User{Id: id}
- var err error = nil
- if selectAll {
- err = DB.First(&user, "id = ?", id).Error
- } else {
- err = DB.Omit("password").First(&user, "id = ?", id).Error
- }
- return &user, err
- }
- func GetUserIdByAffCode(affCode string) (int, error) {
- if affCode == "" {
- return 0, errors.New("affCode 为空!")
- }
- var user User
- err := DB.Select("id").First(&user, "aff_code = ?", affCode).Error
- return user.Id, err
- }
- func DeleteUserById(id int) (err error) {
- if id == 0 {
- return errors.New("id 为空!")
- }
- user := User{Id: id}
- return user.Delete()
- }
- func HardDeleteUserById(id int) error {
- if id == 0 {
- return errors.New("id 为空!")
- }
- err := DB.Unscoped().Delete(&User{}, "id = ?", id).Error
- return err
- }
- func inviteUser(inviterId int) (err error) {
- user, err := GetUserById(inviterId, true)
- if err != nil {
- return err
- }
- user.AffCount++
- user.AffQuota += common.QuotaForInviter
- user.AffHistoryQuota += common.QuotaForInviter
- return DB.Save(user).Error
- }
- func (user *User) TransferAffQuotaToQuota(quota int) error {
- // 检查quota是否小于最小额度
- if float64(quota) < common.QuotaPerUnit {
- return fmt.Errorf("转移额度最小为%s!", logger.LogQuota(int(common.QuotaPerUnit)))
- }
- // 开始数据库事务
- tx := DB.Begin()
- if tx.Error != nil {
- return tx.Error
- }
- defer tx.Rollback() // 确保在函数退出时事务能回滚
- // 加锁查询用户以确保数据一致性
- err := tx.Set("gorm:query_option", "FOR UPDATE").First(&user, user.Id).Error
- if err != nil {
- return err
- }
- // 再次检查用户的AffQuota是否足够
- if user.AffQuota < quota {
- return errors.New("邀请额度不足!")
- }
- // 更新用户额度
- user.AffQuota -= quota
- user.Quota += quota
- // 保存用户状态
- if err := tx.Save(user).Error; err != nil {
- return err
- }
- // 提交事务
- return tx.Commit().Error
- }
- func (user *User) Insert(inviterId int) error {
- var err error
- if user.Password != "" {
- user.Password, err = common.Password2Hash(user.Password)
- if err != nil {
- return err
- }
- }
- user.Quota = common.QuotaForNewUser
- //user.SetAccessToken(common.GetUUID())
- user.AffCode = common.GetRandomString(4)
- // 初始化用户设置,包括默认的边栏配置
- if user.Setting == "" {
- defaultSetting := dto.UserSetting{}
- // 这里暂时不设置SidebarModules,因为需要在用户创建后根据角色设置
- user.SetSetting(defaultSetting)
- }
- result := DB.Create(user)
- if result.Error != nil {
- return result.Error
- }
- // 用户创建成功后,根据角色初始化边栏配置
- // 需要重新获取用户以确保有正确的ID和Role
- var createdUser User
- if err := DB.Where("username = ?", user.Username).First(&createdUser).Error; err == nil {
- // 生成基于角色的默认边栏配置
- defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
- if defaultSidebarConfig != "" {
- currentSetting := createdUser.GetSetting()
- currentSetting.SidebarModules = defaultSidebarConfig
- createdUser.SetSetting(currentSetting)
- createdUser.Update(false)
- common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
- }
- }
- if common.QuotaForNewUser > 0 {
- RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(common.QuotaForNewUser)))
- }
- if inviterId != 0 {
- if common.QuotaForInvitee > 0 {
- _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
- RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
- }
- if common.QuotaForInviter > 0 {
- //_ = IncreaseUserQuota(inviterId, common.QuotaForInviter)
- RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
- _ = inviteUser(inviterId)
- }
- }
- return nil
- }
- // InsertWithTx inserts a new user within an existing transaction.
- // This is used for OAuth registration where user creation and binding need to be atomic.
- // Post-creation tasks (sidebar config, logs, inviter rewards) are handled after the transaction commits.
- func (user *User) InsertWithTx(tx *gorm.DB, inviterId int) error {
- var err error
- if user.Password != "" {
- user.Password, err = common.Password2Hash(user.Password)
- if err != nil {
- return err
- }
- }
- user.Quota = common.QuotaForNewUser
- user.AffCode = common.GetRandomString(4)
- // 初始化用户设置
- if user.Setting == "" {
- defaultSetting := dto.UserSetting{}
- user.SetSetting(defaultSetting)
- }
- result := tx.Create(user)
- if result.Error != nil {
- return result.Error
- }
- return nil
- }
- // FinalizeOAuthUserCreation performs post-transaction tasks for OAuth user creation.
- // This should be called after the transaction commits successfully.
- func (user *User) FinalizeOAuthUserCreation(inviterId int) {
- // 用户创建成功后,根据角色初始化边栏配置
- var createdUser User
- if err := DB.Where("id = ?", user.Id).First(&createdUser).Error; err == nil {
- defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
- if defaultSidebarConfig != "" {
- currentSetting := createdUser.GetSetting()
- currentSetting.SidebarModules = defaultSidebarConfig
- createdUser.SetSetting(currentSetting)
- createdUser.Update(false)
- common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
- }
- }
- if common.QuotaForNewUser > 0 {
- RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(common.QuotaForNewUser)))
- }
- if inviterId != 0 {
- if common.QuotaForInvitee > 0 {
- _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
- RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
- }
- if common.QuotaForInviter > 0 {
- RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
- _ = inviteUser(inviterId)
- }
- }
- }
- func (user *User) Update(updatePassword bool) error {
- var err error
- if updatePassword {
- user.Password, err = common.Password2Hash(user.Password)
- if err != nil {
- return err
- }
- }
- newUser := *user
- DB.First(&user, user.Id)
- if err = DB.Model(user).Updates(newUser).Error; err != nil {
- return err
- }
- // Update cache
- return updateUserCache(*user)
- }
- func (user *User) Edit(updatePassword bool) error {
- var err error
- if updatePassword {
- user.Password, err = common.Password2Hash(user.Password)
- if err != nil {
- return err
- }
- }
- newUser := *user
- updates := map[string]interface{}{
- "username": newUser.Username,
- "display_name": newUser.DisplayName,
- "group": newUser.Group,
- "quota": newUser.Quota,
- "remark": newUser.Remark,
- }
- if updatePassword {
- updates["password"] = newUser.Password
- }
- DB.First(&user, user.Id)
- if err = DB.Model(user).Updates(updates).Error; err != nil {
- return err
- }
- // Update cache
- return updateUserCache(*user)
- }
- func (user *User) ClearBinding(bindingType string) error {
- if user.Id == 0 {
- return errors.New("user id is empty")
- }
- bindingColumnMap := map[string]string{
- "email": "email",
- "github": "github_id",
- "discord": "discord_id",
- "oidc": "oidc_id",
- "wechat": "wechat_id",
- "telegram": "telegram_id",
- "linuxdo": "linux_do_id",
- }
- column, ok := bindingColumnMap[bindingType]
- if !ok {
- return errors.New("invalid binding type")
- }
- if err := DB.Model(&User{}).Where("id = ?", user.Id).Update(column, "").Error; err != nil {
- return err
- }
- if err := DB.Where("id = ?", user.Id).First(user).Error; err != nil {
- return err
- }
- return updateUserCache(*user)
- }
- func (user *User) Delete() error {
- if user.Id == 0 {
- return errors.New("id 为空!")
- }
- if err := DB.Delete(user).Error; err != nil {
- return err
- }
- // 清除缓存
- return invalidateUserCache(user.Id)
- }
- func (user *User) HardDelete() error {
- if user.Id == 0 {
- return errors.New("id 为空!")
- }
- err := DB.Unscoped().Delete(user).Error
- return err
- }
- // ValidateAndFill check password & user status
- func (user *User) ValidateAndFill() (err error) {
- // When querying with struct, GORM will only query with non-zero fields,
- // that means if your field's value is 0, '', false or other zero values,
- // it won't be used to build query conditions
- password := user.Password
- username := strings.TrimSpace(user.Username)
- if username == "" || password == "" {
- return errors.New("用户名或密码为空")
- }
- // find buy username or email
- DB.Where("username = ? OR email = ?", username, username).First(user)
- okay := common.ValidatePasswordAndHash(password, user.Password)
- if !okay || user.Status != common.UserStatusEnabled {
- return errors.New("用户名或密码错误,或用户已被封禁")
- }
- return nil
- }
- func (user *User) FillUserById() error {
- if user.Id == 0 {
- return errors.New("id 为空!")
- }
- DB.Where(User{Id: user.Id}).First(user)
- return nil
- }
- func (user *User) FillUserByEmail() error {
- if user.Email == "" {
- return errors.New("email 为空!")
- }
- DB.Where(User{Email: user.Email}).First(user)
- return nil
- }
- func (user *User) FillUserByGitHubId() error {
- if user.GitHubId == "" {
- return errors.New("GitHub id 为空!")
- }
- DB.Where(User{GitHubId: user.GitHubId}).First(user)
- return nil
- }
- // UpdateGitHubId updates the user's GitHub ID (used for migration from login to numeric ID)
- func (user *User) UpdateGitHubId(newGitHubId string) error {
- if user.Id == 0 {
- return errors.New("user id is empty")
- }
- return DB.Model(user).Update("github_id", newGitHubId).Error
- }
- func (user *User) FillUserByDiscordId() error {
- if user.DiscordId == "" {
- return errors.New("discord id 为空!")
- }
- DB.Where(User{DiscordId: user.DiscordId}).First(user)
- return nil
- }
- func (user *User) FillUserByOidcId() error {
- if user.OidcId == "" {
- return errors.New("oidc id 为空!")
- }
- DB.Where(User{OidcId: user.OidcId}).First(user)
- return nil
- }
- func (user *User) FillUserByWeChatId() error {
- if user.WeChatId == "" {
- return errors.New("WeChat id 为空!")
- }
- DB.Where(User{WeChatId: user.WeChatId}).First(user)
- return nil
- }
- func (user *User) FillUserByTelegramId() error {
- if user.TelegramId == "" {
- return errors.New("Telegram id 为空!")
- }
- err := DB.Where(User{TelegramId: user.TelegramId}).First(user).Error
- if errors.Is(err, gorm.ErrRecordNotFound) {
- return errors.New("该 Telegram 账户未绑定")
- }
- return nil
- }
- func IsEmailAlreadyTaken(email string) bool {
- return DB.Unscoped().Where("email = ?", email).Find(&User{}).RowsAffected == 1
- }
- func IsWeChatIdAlreadyTaken(wechatId string) bool {
- return DB.Unscoped().Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1
- }
- func IsGitHubIdAlreadyTaken(githubId string) bool {
- return DB.Unscoped().Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1
- }
- func IsDiscordIdAlreadyTaken(discordId string) bool {
- return DB.Unscoped().Where("discord_id = ?", discordId).Find(&User{}).RowsAffected == 1
- }
- func IsOidcIdAlreadyTaken(oidcId string) bool {
- return DB.Where("oidc_id = ?", oidcId).Find(&User{}).RowsAffected == 1
- }
- func IsTelegramIdAlreadyTaken(telegramId string) bool {
- return DB.Unscoped().Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1
- }
- func ResetUserPasswordByEmail(email string, password string) error {
- if email == "" || password == "" {
- return errors.New("邮箱地址或密码为空!")
- }
- hashedPassword, err := common.Password2Hash(password)
- if err != nil {
- return err
- }
- err = DB.Model(&User{}).Where("email = ?", email).Update("password", hashedPassword).Error
- return err
- }
- func IsAdmin(userId int) bool {
- if userId == 0 {
- return false
- }
- var user User
- err := DB.Where("id = ?", userId).Select("role").Find(&user).Error
- if err != nil {
- common.SysLog("no such user " + err.Error())
- return false
- }
- return user.Role >= common.RoleAdminUser
- }
- //// IsUserEnabled checks user status from Redis first, falls back to DB if needed
- //func IsUserEnabled(id int, fromDB bool) (status bool, err error) {
- // defer func() {
- // // Update Redis cache asynchronously on successful DB read
- // if shouldUpdateRedis(fromDB, err) {
- // gopool.Go(func() {
- // if err := updateUserStatusCache(id, status); err != nil {
- // common.SysError("failed to update user status cache: " + err.Error())
- // }
- // })
- // }
- // }()
- // if !fromDB && common.RedisEnabled {
- // // Try Redis first
- // status, err := getUserStatusCache(id)
- // if err == nil {
- // return status == common.UserStatusEnabled, nil
- // }
- // // Don't return error - fall through to DB
- // }
- // fromDB = true
- // var user User
- // err = DB.Where("id = ?", id).Select("status").Find(&user).Error
- // if err != nil {
- // return false, err
- // }
- //
- // return user.Status == common.UserStatusEnabled, nil
- //}
- func ValidateAccessToken(token string) (user *User) {
- if token == "" {
- return nil
- }
- token = strings.Replace(token, "Bearer ", "", 1)
- user = &User{}
- if DB.Where("access_token = ?", token).First(user).RowsAffected == 1 {
- return user
- }
- return nil
- }
- // GetUserQuota gets quota from Redis first, falls back to DB if needed
- func GetUserQuota(id int, fromDB bool) (quota int, err error) {
- defer func() {
- // Update Redis cache asynchronously on successful DB read
- if shouldUpdateRedis(fromDB, err) {
- gopool.Go(func() {
- if err := updateUserQuotaCache(id, quota); err != nil {
- common.SysLog("failed to update user quota cache: " + err.Error())
- }
- })
- }
- }()
- if !fromDB && common.RedisEnabled {
- quota, err := getUserQuotaCache(id)
- if err == nil {
- return quota, nil
- }
- // Don't return error - fall through to DB
- }
- fromDB = true
- err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find("a).Error
- if err != nil {
- return 0, err
- }
- return quota, nil
- }
- func GetUserUsedQuota(id int) (quota int, err error) {
- err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find("a).Error
- return quota, err
- }
- func GetUserEmail(id int) (email string, err error) {
- err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
- return email, err
- }
- // GetUserGroup gets group from Redis first, falls back to DB if needed
- func GetUserGroup(id int, fromDB bool) (group string, err error) {
- defer func() {
- // Update Redis cache asynchronously on successful DB read
- if shouldUpdateRedis(fromDB, err) {
- gopool.Go(func() {
- if err := updateUserGroupCache(id, group); err != nil {
- common.SysLog("failed to update user group cache: " + err.Error())
- }
- })
- }
- }()
- if !fromDB && common.RedisEnabled {
- group, err := getUserGroupCache(id)
- if err == nil {
- return group, nil
- }
- // Don't return error - fall through to DB
- }
- fromDB = true
- err = DB.Model(&User{}).Where("id = ?", id).Select(commonGroupCol).Find(&group).Error
- if err != nil {
- return "", err
- }
- return group, nil
- }
- // GetUserSetting gets setting from Redis first, falls back to DB if needed
- func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error) {
- var setting string
- defer func() {
- // Update Redis cache asynchronously on successful DB read
- if shouldUpdateRedis(fromDB, err) {
- gopool.Go(func() {
- if err := updateUserSettingCache(id, setting); err != nil {
- common.SysLog("failed to update user setting cache: " + err.Error())
- }
- })
- }
- }()
- if !fromDB && common.RedisEnabled {
- setting, err := getUserSettingCache(id)
- if err == nil {
- return setting, nil
- }
- // Don't return error - fall through to DB
- }
- fromDB = true
- // can be nil setting
- var safeSetting sql.NullString
- err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&safeSetting).Error
- if err != nil {
- return settingMap, err
- }
- if safeSetting.Valid {
- setting = safeSetting.String
- } else {
- setting = ""
- }
- userBase := &UserBase{
- Setting: setting,
- }
- return userBase.GetSetting(), nil
- }
- func IncreaseUserQuota(id int, quota int, db bool) (err error) {
- if quota < 0 {
- return errors.New("quota 不能为负数!")
- }
- gopool.Go(func() {
- err := cacheIncrUserQuota(id, int64(quota))
- if err != nil {
- common.SysLog("failed to increase user quota: " + err.Error())
- }
- })
- if !db && common.BatchUpdateEnabled {
- addNewRecord(BatchUpdateTypeUserQuota, id, quota)
- return nil
- }
- return increaseUserQuota(id, quota)
- }
- func increaseUserQuota(id int, quota int) (err error) {
- err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
- if err != nil {
- return err
- }
- return err
- }
- func DecreaseUserQuota(id int, quota int) (err error) {
- if quota < 0 {
- return errors.New("quota 不能为负数!")
- }
- gopool.Go(func() {
- err := cacheDecrUserQuota(id, int64(quota))
- if err != nil {
- common.SysLog("failed to decrease user quota: " + err.Error())
- }
- })
- if common.BatchUpdateEnabled {
- addNewRecord(BatchUpdateTypeUserQuota, id, -quota)
- return nil
- }
- return decreaseUserQuota(id, quota)
- }
- func decreaseUserQuota(id int, quota int) (err error) {
- err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
- if err != nil {
- return err
- }
- return err
- }
- func DeltaUpdateUserQuota(id int, delta int) (err error) {
- if delta == 0 {
- return nil
- }
- if delta > 0 {
- return IncreaseUserQuota(id, delta, false)
- } else {
- return DecreaseUserQuota(id, -delta)
- }
- }
- //func GetRootUserEmail() (email string) {
- // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email)
- // return email
- //}
- func GetRootUser() (user *User) {
- DB.Where("role = ?", common.RoleRootUser).First(&user)
- return user
- }
- func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
- if common.BatchUpdateEnabled {
- addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
- addNewRecord(BatchUpdateTypeRequestCount, id, 1)
- return
- }
- updateUserUsedQuotaAndRequestCount(id, quota, 1)
- }
- func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
- err := DB.Model(&User{}).Where("id = ?", id).Updates(
- map[string]interface{}{
- "used_quota": gorm.Expr("used_quota + ?", quota),
- "request_count": gorm.Expr("request_count + ?", count),
- },
- ).Error
- if err != nil {
- common.SysLog("failed to update user used quota and request count: " + err.Error())
- return
- }
- //// 更新缓存
- //if err := invalidateUserCache(id); err != nil {
- // common.SysError("failed to invalidate user cache: " + err.Error())
- //}
- }
- func updateUserUsedQuota(id int, quota int) {
- err := DB.Model(&User{}).Where("id = ?", id).Updates(
- map[string]interface{}{
- "used_quota": gorm.Expr("used_quota + ?", quota),
- },
- ).Error
- if err != nil {
- common.SysLog("failed to update user used quota: " + err.Error())
- }
- }
- func updateUserRequestCount(id int, count int) {
- err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error
- if err != nil {
- common.SysLog("failed to update user request count: " + err.Error())
- }
- }
- // GetUsernameById gets username from Redis first, falls back to DB if needed
- func GetUsernameById(id int, fromDB bool) (username string, err error) {
- defer func() {
- // Update Redis cache asynchronously on successful DB read
- if shouldUpdateRedis(fromDB, err) {
- gopool.Go(func() {
- if err := updateUserNameCache(id, username); err != nil {
- common.SysLog("failed to update user name cache: " + err.Error())
- }
- })
- }
- }()
- if !fromDB && common.RedisEnabled {
- username, err := getUserNameCache(id)
- if err == nil {
- return username, nil
- }
- // Don't return error - fall through to DB
- }
- fromDB = true
- err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error
- if err != nil {
- return "", err
- }
- return username, nil
- }
- func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
- var user User
- err := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).First(&user).Error
- return !errors.Is(err, gorm.ErrRecordNotFound)
- }
- func (user *User) FillUserByLinuxDOId() error {
- if user.LinuxDOId == "" {
- return errors.New("linux do id is empty")
- }
- err := DB.Where("linux_do_id = ?", user.LinuxDOId).First(user).Error
- return err
- }
- func RootUserExists() bool {
- var user User
- err := DB.Where("role = ?", common.RoleRootUser).First(&user).Error
- if err != nil {
- return false
- }
- return true
- }
|