token_cache.go 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. package model
  2. import (
  3. "fmt"
  4. "one-api/common"
  5. "one-api/constant"
  6. "time"
  7. )
  8. func cacheSetToken(token Token) error {
  9. key := common.GenerateHMAC(token.Key)
  10. token.Clean()
  11. err := common.RedisHSetObj(fmt.Sprintf("token:%s", key), &token, time.Duration(common.RedisKeyCacheSeconds())*time.Second)
  12. if err != nil {
  13. return err
  14. }
  15. return nil
  16. }
  17. func cacheDeleteToken(key string) error {
  18. key = common.GenerateHMAC(key)
  19. err := common.RedisDelKey(fmt.Sprintf("token:%s", key))
  20. if err != nil {
  21. return err
  22. }
  23. return nil
  24. }
  25. func cacheIncrTokenQuota(key string, increment int64) error {
  26. key = common.GenerateHMAC(key)
  27. err := common.RedisHIncrBy(fmt.Sprintf("token:%s", key), constant.TokenFiledRemainQuota, increment)
  28. if err != nil {
  29. return err
  30. }
  31. return nil
  32. }
  33. func cacheDecrTokenQuota(key string, decrement int64) error {
  34. return cacheIncrTokenQuota(key, -decrement)
  35. }
  36. func cacheSetTokenField(key string, field string, value string) error {
  37. key = common.GenerateHMAC(key)
  38. err := common.RedisHSetField(fmt.Sprintf("token:%s", key), field, value)
  39. if err != nil {
  40. return err
  41. }
  42. return nil
  43. }
  44. // CacheGetTokenByKey 从缓存中获取 token,如果缓存中不存在,则从数据库中获取
  45. func cacheGetTokenByKey(key string) (*Token, error) {
  46. hmacKey := common.GenerateHMAC(key)
  47. if !common.RedisEnabled {
  48. return nil, fmt.Errorf("redis is not enabled")
  49. }
  50. var token Token
  51. err := common.RedisHGetObj(fmt.Sprintf("token:%s", hmacKey), &token)
  52. if err != nil {
  53. return nil, err
  54. }
  55. token.Key = key
  56. return &token, nil
  57. }