Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

464 linhas
14 KiB

  1. package model
  2. import (
  3. "errors"
  4. "fmt"
  5. "strings"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/QuantumNous/new-api/setting/operation_setting"
  8. "github.com/bytedance/gopkg/util/gopool"
  9. "gorm.io/gorm"
  10. )
  11. type Token struct {
  12. Id int `json:"id"`
  13. UserId int `json:"user_id" gorm:"index"`
  14. Key string `json:"key" gorm:"type:char(48);uniqueIndex"`
  15. Status int `json:"status" gorm:"default:1"`
  16. Name string `json:"name" gorm:"index" `
  17. CreatedTime int64 `json:"created_time" gorm:"bigint"`
  18. AccessedTime int64 `json:"accessed_time" gorm:"bigint"`
  19. ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
  20. RemainQuota int `json:"remain_quota" gorm:"default:0"`
  21. UnlimitedQuota bool `json:"unlimited_quota"`
  22. ModelLimitsEnabled bool `json:"model_limits_enabled"`
  23. ModelLimits string `json:"model_limits" gorm:"type:varchar(1024);default:''"`
  24. AllowIps *string `json:"allow_ips" gorm:"default:''"`
  25. UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota
  26. Group string `json:"group" gorm:"default:''"`
  27. CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试,仅auto分组有效
  28. BoundChannelId *int `json:"bound_channel_id" gorm:"index"` // 绑定的渠道ID,nil表示不绑定
  29. DeletedAt gorm.DeletedAt `gorm:"index"`
  30. }
  31. func (token *Token) Clean() {
  32. token.Key = ""
  33. }
  34. func (token *Token) GetIpLimits() []string {
  35. // delete empty spaces
  36. //split with \n
  37. ipLimits := make([]string, 0)
  38. if token.AllowIps == nil {
  39. return ipLimits
  40. }
  41. cleanIps := strings.ReplaceAll(*token.AllowIps, " ", "")
  42. if cleanIps == "" {
  43. return ipLimits
  44. }
  45. ips := strings.Split(cleanIps, "\n")
  46. for _, ip := range ips {
  47. ip = strings.TrimSpace(ip)
  48. ip = strings.ReplaceAll(ip, ",", "")
  49. if ip != "" {
  50. ipLimits = append(ipLimits, ip)
  51. }
  52. }
  53. return ipLimits
  54. }
  55. func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
  56. var tokens []*Token
  57. var err error
  58. err = DB.Where("user_id = ?", userId).Order("id desc").Limit(num).Offset(startIdx).Find(&tokens).Error
  59. return tokens, err
  60. }
  61. // sanitizeLikePattern 校验并清洗用户输入的 LIKE 搜索模式。
  62. // 规则:
  63. // 1. 转义 ! 和 _(使用 ! 作为 ESCAPE 字符,兼容 MySQL/PostgreSQL/SQLite)
  64. // 2. 连续的 % 合并为单个 %
  65. // 3. 最多允许 2 个 %
  66. // 4. 含 % 时(模糊搜索),去掉 % 后关键词长度必须 >= 2
  67. // 5. 不含 % 时按精确匹配
  68. func sanitizeLikePattern(input string) (string, error) {
  69. // 1. 先转义 ESCAPE 字符 ! 自身,再转义 _
  70. // 使用 ! 而非 \ 作为 ESCAPE 字符,避免 MySQL 中反斜杠的字符串转义问题
  71. input = strings.ReplaceAll(input, "!", "!!")
  72. input = strings.ReplaceAll(input, `_`, `!_`)
  73. // 2. 连续的 % 直接拒绝
  74. if strings.Contains(input, "%%") {
  75. return "", errors.New("搜索模式中不允许包含连续的 % 通配符")
  76. }
  77. // 3. 统计 % 数量,不得超过 2
  78. count := strings.Count(input, "%")
  79. if count > 2 {
  80. return "", errors.New("搜索模式中最多允许包含 2 个 % 通配符")
  81. }
  82. // 4. 含 % 时,去掉 % 后关键词长度必须 >= 2
  83. if count > 0 {
  84. stripped := strings.ReplaceAll(input, "%", "")
  85. if len(stripped) < 2 {
  86. return "", errors.New("使用模糊搜索时,关键词长度至少为 2 个字符")
  87. }
  88. return input, nil
  89. }
  90. // 5. 无 % 时,精确全匹配
  91. return input, nil
  92. }
  93. const searchHardLimit = 100
  94. func SearchUserTokens(userId int, keyword string, token string, offset int, limit int) (tokens []*Token, total int64, err error) {
  95. // model 层强制截断
  96. if limit <= 0 || limit > searchHardLimit {
  97. limit = searchHardLimit
  98. }
  99. if offset < 0 {
  100. offset = 0
  101. }
  102. if token != "" {
  103. token = strings.TrimPrefix(token, "sk-")
  104. }
  105. // 超量用户(令牌数超过上限)只允许精确搜索,禁止模糊搜索
  106. maxTokens := operation_setting.GetMaxUserTokens()
  107. hasFuzzy := strings.Contains(keyword, "%") || strings.Contains(token, "%")
  108. if hasFuzzy {
  109. count, err := CountUserTokens(userId)
  110. if err != nil {
  111. common.SysLog("failed to count user tokens: " + err.Error())
  112. return nil, 0, errors.New("获取令牌数量失败")
  113. }
  114. if int(count) > maxTokens {
  115. return nil, 0, errors.New("令牌数量超过上限,仅允许精确搜索,请勿使用 % 通配符")
  116. }
  117. }
  118. baseQuery := DB.Model(&Token{}).Where("user_id = ?", userId)
  119. // 非空才加 LIKE 条件,空则跳过(不过滤该字段)
  120. if keyword != "" {
  121. keywordPattern, err := sanitizeLikePattern(keyword)
  122. if err != nil {
  123. return nil, 0, err
  124. }
  125. baseQuery = baseQuery.Where("name LIKE ? ESCAPE '!'", keywordPattern)
  126. }
  127. if token != "" {
  128. tokenPattern, err := sanitizeLikePattern(token)
  129. if err != nil {
  130. return nil, 0, err
  131. }
  132. baseQuery = baseQuery.Where(commonKeyCol+" LIKE ? ESCAPE '!'", tokenPattern)
  133. }
  134. // 先查匹配总数(用于分页,受 maxTokens 上限保护,避免全表 COUNT)
  135. err = baseQuery.Limit(maxTokens).Count(&total).Error
  136. if err != nil {
  137. common.SysError("failed to count search tokens: " + err.Error())
  138. return nil, 0, errors.New("搜索令牌失败")
  139. }
  140. // 再分页查数据
  141. err = baseQuery.Order("id desc").Offset(offset).Limit(limit).Find(&tokens).Error
  142. if err != nil {
  143. common.SysError("failed to search tokens: " + err.Error())
  144. return nil, 0, errors.New("搜索令牌失败")
  145. }
  146. return tokens, total, nil
  147. }
  148. func ValidateUserToken(key string) (token *Token, err error) {
  149. if key == "" {
  150. return nil, errors.New("未提供令牌")
  151. }
  152. token, err = GetTokenByKey(key, false)
  153. if err == nil {
  154. if token.Status == common.TokenStatusExhausted {
  155. keyPrefix := key[:3]
  156. keySuffix := key[len(key)-3:]
  157. return token, errors.New("该令牌额度已用尽 TokenStatusExhausted[sk-" + keyPrefix + "***" + keySuffix + "]")
  158. } else if token.Status == common.TokenStatusExpired {
  159. return token, errors.New("该令牌已过期")
  160. }
  161. if token.Status != common.TokenStatusEnabled {
  162. return token, errors.New("该令牌状态不可用")
  163. }
  164. if token.ExpiredTime != -1 && token.ExpiredTime < common.GetTimestamp() {
  165. if !common.RedisEnabled {
  166. token.Status = common.TokenStatusExpired
  167. err := token.SelectUpdate()
  168. if err != nil {
  169. common.SysLog("failed to update token status" + err.Error())
  170. }
  171. }
  172. return token, errors.New("该令牌已过期")
  173. }
  174. if !token.UnlimitedQuota && token.RemainQuota <= 0 {
  175. if !common.RedisEnabled {
  176. // in this case, we can make sure the token is exhausted
  177. token.Status = common.TokenStatusExhausted
  178. err := token.SelectUpdate()
  179. if err != nil {
  180. common.SysLog("failed to update token status" + err.Error())
  181. }
  182. }
  183. keyPrefix := key[:3]
  184. keySuffix := key[len(key)-3:]
  185. return token, errors.New(fmt.Sprintf("[sk-%s***%s] 该令牌额度已用尽 !token.UnlimitedQuota && token.RemainQuota = %d", keyPrefix, keySuffix, token.RemainQuota))
  186. }
  187. return token, nil
  188. }
  189. common.SysLog("ValidateUserToken: failed to get token: " + err.Error())
  190. if errors.Is(err, gorm.ErrRecordNotFound) {
  191. return nil, errors.New("无效的令牌")
  192. } else {
  193. return nil, errors.New("无效的令牌,数据库查询出错,请联系管理员")
  194. }
  195. }
  196. func GetTokenByIds(id int, userId int) (*Token, error) {
  197. if id == 0 || userId == 0 {
  198. return nil, errors.New("id 或 userId 为空!")
  199. }
  200. token := Token{Id: id, UserId: userId}
  201. var err error = nil
  202. err = DB.First(&token, "id = ? and user_id = ?", id, userId).Error
  203. return &token, err
  204. }
  205. func GetTokenById(id int) (*Token, error) {
  206. if id == 0 {
  207. return nil, errors.New("id 为空!")
  208. }
  209. token := Token{Id: id}
  210. var err error = nil
  211. err = DB.First(&token, "id = ?", id).Error
  212. if shouldUpdateRedis(true, err) {
  213. gopool.Go(func() {
  214. if err := cacheSetToken(token); err != nil {
  215. common.SysLog("failed to update user status cache: " + err.Error())
  216. }
  217. })
  218. }
  219. return &token, err
  220. }
  221. func GetTokenByKey(key string, fromDB bool) (token *Token, err error) {
  222. defer func() {
  223. // Update Redis cache asynchronously on successful DB read
  224. if shouldUpdateRedis(fromDB, err) && token != nil {
  225. gopool.Go(func() {
  226. if err := cacheSetToken(*token); err != nil {
  227. common.SysLog("failed to update user status cache: " + err.Error())
  228. }
  229. })
  230. }
  231. }()
  232. if !fromDB && common.RedisEnabled {
  233. // Try Redis first
  234. token, err := cacheGetTokenByKey(key)
  235. if err == nil {
  236. return token, nil
  237. }
  238. // Don't return error - fall through to DB
  239. }
  240. fromDB = true
  241. err = DB.Where(commonKeyCol+" = ?", key).First(&token).Error
  242. return token, err
  243. }
  244. func (token *Token) Insert() error {
  245. var err error
  246. err = DB.Create(token).Error
  247. return err
  248. }
  249. // Update Make sure your token's fields is completed, because this will update non-zero values
  250. func (token *Token) Update() (err error) {
  251. defer func() {
  252. if shouldUpdateRedis(true, err) {
  253. gopool.Go(func() {
  254. err := cacheSetToken(*token)
  255. if err != nil {
  256. common.SysLog("failed to update token cache: " + err.Error())
  257. }
  258. })
  259. }
  260. }()
  261. err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
  262. "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error
  263. return err
  264. }
  265. func (token *Token) SelectUpdate() (err error) {
  266. defer func() {
  267. if shouldUpdateRedis(true, err) {
  268. gopool.Go(func() {
  269. err := cacheSetToken(*token)
  270. if err != nil {
  271. common.SysLog("failed to update token cache: " + err.Error())
  272. }
  273. })
  274. }
  275. }()
  276. // This can update zero values
  277. return DB.Model(token).Select("accessed_time", "status").Updates(token).Error
  278. }
  279. func (token *Token) Delete() (err error) {
  280. defer func() {
  281. if shouldUpdateRedis(true, err) {
  282. gopool.Go(func() {
  283. err := cacheDeleteToken(token.Key)
  284. if err != nil {
  285. common.SysLog("failed to delete token cache: " + err.Error())
  286. }
  287. })
  288. }
  289. }()
  290. err = DB.Delete(token).Error
  291. return err
  292. }
  293. func (token *Token) IsModelLimitsEnabled() bool {
  294. return token.ModelLimitsEnabled
  295. }
  296. func (token *Token) GetModelLimits() []string {
  297. if token.ModelLimits == "" {
  298. return []string{}
  299. }
  300. return strings.Split(token.ModelLimits, ",")
  301. }
  302. func (token *Token) GetModelLimitsMap() map[string]bool {
  303. limits := token.GetModelLimits()
  304. limitsMap := make(map[string]bool)
  305. for _, limit := range limits {
  306. limitsMap[limit] = true
  307. }
  308. return limitsMap
  309. }
  310. func DisableModelLimits(tokenId int) error {
  311. token, err := GetTokenById(tokenId)
  312. if err != nil {
  313. return err
  314. }
  315. token.ModelLimitsEnabled = false
  316. token.ModelLimits = ""
  317. return token.Update()
  318. }
  319. func DeleteTokenById(id int, userId int) (err error) {
  320. // Why we need userId here? In case user want to delete other's token.
  321. if id == 0 || userId == 0 {
  322. return errors.New("id 或 userId 为空!")
  323. }
  324. token := Token{Id: id, UserId: userId}
  325. err = DB.Where(token).First(&token).Error
  326. if err != nil {
  327. return err
  328. }
  329. return token.Delete()
  330. }
  331. func IncreaseTokenQuota(tokenId int, key string, quota int) (err error) {
  332. if quota < 0 {
  333. return errors.New("quota 不能为负数!")
  334. }
  335. if common.RedisEnabled {
  336. gopool.Go(func() {
  337. err := cacheIncrTokenQuota(key, int64(quota))
  338. if err != nil {
  339. common.SysLog("failed to increase token quota: " + err.Error())
  340. }
  341. })
  342. }
  343. if common.BatchUpdateEnabled {
  344. addNewRecord(BatchUpdateTypeTokenQuota, tokenId, quota)
  345. return nil
  346. }
  347. return increaseTokenQuota(tokenId, quota)
  348. }
  349. func increaseTokenQuota(id int, quota int) (err error) {
  350. err = DB.Model(&Token{}).Where("id = ?", id).Updates(
  351. map[string]interface{}{
  352. "remain_quota": gorm.Expr("remain_quota + ?", quota),
  353. "used_quota": gorm.Expr("used_quota - ?", quota),
  354. "accessed_time": common.GetTimestamp(),
  355. },
  356. ).Error
  357. return err
  358. }
  359. func DecreaseTokenQuota(id int, key string, quota int) (err error) {
  360. if quota < 0 {
  361. return errors.New("quota 不能为负数!")
  362. }
  363. if common.RedisEnabled {
  364. gopool.Go(func() {
  365. err := cacheDecrTokenQuota(key, int64(quota))
  366. if err != nil {
  367. common.SysLog("failed to decrease token quota: " + err.Error())
  368. }
  369. })
  370. }
  371. if common.BatchUpdateEnabled {
  372. addNewRecord(BatchUpdateTypeTokenQuota, id, -quota)
  373. return nil
  374. }
  375. return decreaseTokenQuota(id, quota)
  376. }
  377. func decreaseTokenQuota(id int, quota int) (err error) {
  378. err = DB.Model(&Token{}).Where("id = ?", id).Updates(
  379. map[string]interface{}{
  380. "remain_quota": gorm.Expr("remain_quota - ?", quota),
  381. "used_quota": gorm.Expr("used_quota + ?", quota),
  382. "accessed_time": common.GetTimestamp(),
  383. },
  384. ).Error
  385. return err
  386. }
  387. // CountUserTokens returns total number of tokens for the given user, used for pagination
  388. func CountUserTokens(userId int) (int64, error) {
  389. var total int64
  390. err := DB.Model(&Token{}).Where("user_id = ?", userId).Count(&total).Error
  391. return total, err
  392. }
  393. // BatchDeleteTokens 删除指定用户的一组令牌,返回成功删除数量
  394. func BatchDeleteTokens(ids []int, userId int) (int, error) {
  395. if len(ids) == 0 {
  396. return 0, errors.New("ids 不能为空!")
  397. }
  398. tx := DB.Begin()
  399. var tokens []Token
  400. if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Find(&tokens).Error; err != nil {
  401. tx.Rollback()
  402. return 0, err
  403. }
  404. if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Delete(&Token{}).Error; err != nil {
  405. tx.Rollback()
  406. return 0, err
  407. }
  408. if err := tx.Commit().Error; err != nil {
  409. return 0, err
  410. }
  411. if common.RedisEnabled {
  412. gopool.Go(func() {
  413. for _, t := range tokens {
  414. _ = cacheDeleteToken(t.Key)
  415. }
  416. })
  417. }
  418. return len(tokens), nil
  419. }