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.
50 lines
829 B
50 lines
829 B
package middleware |
|
|
|
import ( |
|
"server/common" |
|
"server/db" |
|
"server/modules/web/app" |
|
"time" |
|
|
|
"github.com/gin-gonic/gin" |
|
"github.com/liangdas/mqant/log" |
|
"golang.org/x/time/rate" |
|
) |
|
|
|
const ( |
|
UserLimitPerSec = 10 |
|
) |
|
|
|
func RateLimitMiddleware() gin.HandlerFunc { |
|
l := rate.NewLimiter(5000, 20000) |
|
return func(c *gin.Context) { |
|
if !l.Allow() { |
|
c.Abort() |
|
return |
|
} |
|
c.Next() |
|
} |
|
} |
|
|
|
func UserLimitMiddleware() gin.HandlerFunc { |
|
return func(c *gin.Context) { |
|
a := app.NewApp(c) |
|
ip := a.GetRemoteIP() |
|
key := common.GetRedisLockKeyIP(ip) |
|
if db.Redis().Exist(key) { |
|
res, err := db.Redis().Incr(key, 1) |
|
if err != nil { |
|
log.Error("err:%v", err) |
|
c.Abort() |
|
return |
|
} |
|
if res > UserLimitPerSec { |
|
c.Abort() |
|
return |
|
} |
|
} else { |
|
db.Redis().SetData(key, 0, time.Second) |
|
} |
|
c.Next() |
|
} |
|
}
|
|
|