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")
}