gin结合go-redis根据ip完成访问频率限制

import (
	"time"

	"github.com/go-redis/redis"
)

// Limiter 定义属性
type Limiter struct {
	// Redis client connection.
	rc *redis.Client
}

// 根据redisURL创建新的limiter并返回
func NewLimiter(redisURL string) (*Limiter, error) {
	opts, err := redis.ParseURL(redisURL)
	if err != nil {
		return nil, err
	}

	rc := redis.NewClient(opts)
	if err := rc.Ping().Err(); err != nil {
		return nil, err
	}

	return &Limiter{rc: rc}, nil
}

// 通过redis的value判断第几次访问并返回是否允许访问
func (l *Limiter) Allow(key string, events int64, per time.Duration) bool {
	curr := l.rc.LLen(key).Val()
	if curr >= events {
		return false
	}

	if v := l.rc.Exists(key).Val(); v == 0 {
		pipe := l.rc.TxPipeline()
		pipe.RPush(key, key)
		//设置过期时间
		pipe.Expire(key, per)
		_, _ = pipe.Exec()
	} else {
		l.rc.RPushX(key, key)
	}

	return true
}

main.go

func main(){
	router := gin.Default()
	//创建一个新limiter.
	limiter, err := util.NewLimiter("redis://localhost:6379")
	if err != nil {
		log.Fatal(err)
	}
	router.Use(middlewares.RateMiddleware(limiter))
	router.run()
}

meddlewares.go

//限制访问
func RateMiddleware(limiter *util.Limiter) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 如果ip请求连接数在两秒内超过5次，返回429并抛出error
		if !limiter.Allow(c.ClientIP(), 5, 2*time.Second) {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"error": "too many requests",
			})
			log.Println("too many requests")
			return
		}
		c.Next()
	}
}