Initial commit

middleware/limit/limit.go

@@ -0,0 +1,25 @@
+package limit
+
+import (
+	"fmt"
+	"stat/common/redis"
+	"stat/statg"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+func GetKey(name string) func() string {
+	return func() string {
+		return redis.GetIpLimitDayKey(statg.Conf.Host, statg.Conf.Port) + "_" + name
+	}
+}
+
+func Limit(name string, capacity int64, rate int64) gin.HandlerFunc {
+	limiter := redis.NewTokenBucket(redis.StatRds.Client(), GetKey(name), capacity, rate, 48*time.Hour)
+	return func(ctx *gin.Context) {
+		if err := limiter.TakeTokens(ctx.ClientIP(), 1); err != nil {
+			ctx.AbortWithStatusJSON(200, gin.H{"msg": fmt.Sprintf("%+v", err)})
+		}
+	}
+}