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.
51 lines
829 B
51 lines
829 B
1 year ago
|
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()
|
||
|
|
}
|
||
|
|
}
|