rateLimitWithRedis.go

package main

import (
	"errors"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis"
	"sync"
	"time"
)

var rdb *redis.Client

func incRequestCount(key string,rateLimit int64,second int64) error {
	err := rdb.Watch(func(tx *redis.Tx) error {
		_ = tx.SetNX(key,0,time.Duration(second)*time.Second)
		count,err := tx.Incr(key).Result()
		if count > rateLimit {
			err  = errors.New("rate limited")
		}
		if err != nil {
		    return err
		}
		return nil
	},key)
	return  err
}

var wg sync.WaitGroup

func useRateLimit(rateLimit int64,second int64) gin.HandlerFunc{
	return func(c *gin.Context){
		clientIp := c.ClientIP()
		key := "RATE_LIMIT_COUNT_"+clientIp
		err := incRequestCount(key,rateLimit,second)
		if err != nil {
			c.AbortWithStatus(403)
			return
		}
		c.Next()
	}
}

func main(){
	app :=gin.Default()
	app.Use(useRateLimit(10,60))
	app.GET("/", func(c *gin.Context) {
		c.JSON(200,"Hello World")
	})
	rdb = redis.NewClient(&redis.Options{
	    Addr:     "localhost:6379", // use default Addr
	    Password: "",               // no password set
	    DB:       0,                // use default database
	    PoolTimeout:time.Minute,    // since we user transaction so it can take a long time
	})
	app.Run(":8080")
}

Nếu mình thêm chức năng bật tắt nó thì mình cần làm gì anh . em cũng mới học sơ về golang nên em cần hỏi Ý ạ thanks a mong a giải đáp