You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

94 lines
2.4 KiB

  1. package middleware
  2. import (
  3. "context"
  4. "errors"
  5. "net/http"
  6. "time"
  7. "go.uber.org/zap"
  8. "github.com/flipped-aurora/gin-vue-admin/server/global"
  9. "github.com/flipped-aurora/gin-vue-admin/server/model/common/response"
  10. "github.com/gin-gonic/gin"
  11. )
  12. type LimitConfig struct {
  13. // GenerationKey 根据业务生成key 下面CheckOrMark查询生成
  14. GenerationKey func(c *gin.Context) string
  15. // 检查函数,用户可修改具体逻辑,更加灵活
  16. CheckOrMark func(key string, expire int, limit int) error
  17. // Expire key 过期时间
  18. Expire int
  19. // Limit 周期时间
  20. Limit int
  21. }
  22. func (l LimitConfig) LimitWithTime() gin.HandlerFunc {
  23. return func(c *gin.Context) {
  24. if err := l.CheckOrMark(l.GenerationKey(c), l.Expire, l.Limit); err != nil {
  25. c.JSON(http.StatusOK, gin.H{"code": response.ERROR, "msg": err})
  26. c.Abort()
  27. return
  28. } else {
  29. c.Next()
  30. }
  31. }
  32. }
  33. // DefaultGenerationKey 默认生成key
  34. func DefaultGenerationKey(c *gin.Context) string {
  35. return "GVA_Limit" + c.ClientIP()
  36. }
  37. func DefaultCheckOrMark(key string, expire int, limit int) (err error) {
  38. // 判断是否开启redis
  39. if global.GVA_REDIS == nil {
  40. return err
  41. }
  42. if err = SetLimitWithTime(key, limit, time.Duration(expire)*time.Second); err != nil {
  43. global.GVA_LOG.Error("limit", zap.Error(err))
  44. }
  45. return err
  46. }
  47. func DefaultLimit() gin.HandlerFunc {
  48. return LimitConfig{
  49. GenerationKey: DefaultGenerationKey,
  50. CheckOrMark: DefaultCheckOrMark,
  51. Expire: global.GVA_CONFIG.System.LimitTimeIP,
  52. Limit: global.GVA_CONFIG.System.LimitCountIP,
  53. }.LimitWithTime()
  54. }
  55. // SetLimitWithTime 设置访问次数
  56. func SetLimitWithTime(key string, limit int, expiration time.Duration) error {
  57. count, err := global.GVA_REDIS.Exists(context.Background(), key).Result()
  58. if err != nil {
  59. return err
  60. }
  61. if count == 0 {
  62. pipe := global.GVA_REDIS.TxPipeline()
  63. pipe.Incr(context.Background(), key)
  64. pipe.Expire(context.Background(), key, expiration)
  65. _, err = pipe.Exec(context.Background())
  66. return err
  67. } else {
  68. //次数
  69. if times, err := global.GVA_REDIS.Get(context.Background(), key).Int(); err != nil {
  70. return err
  71. } else {
  72. if times >= limit {
  73. if t, err := global.GVA_REDIS.PTTL(context.Background(), key).Result(); err != nil {
  74. return errors.New("请求太过频繁,请稍后再试")
  75. } else {
  76. return errors.New("请求太过频繁, 请 " + t.String() + " 秒后尝试")
  77. }
  78. } else {
  79. return global.GVA_REDIS.Incr(context.Background(), key).Err()
  80. }
  81. }
  82. }
  83. }