访问频率限制 优化

添加黑名单机制

services/waf/waf.go

@@ -6,6 +6,7 @@ import (
 	"gofiber.study.skcks.cn/common/config"
 	"gofiber.study.skcks.cn/common/logger"
 	"gofiber.study.skcks.cn/global"
+	"gofiber.study.skcks.cn/model/generic/models"
 	"time"
 	"xorm.io/xorm"
 )
@@ -43,6 +44,19 @@ func InitService() {
 
 func (w *Waf) Access(ip string) bool {
 	key := StoreName + Separator + "access" + Separator + ip
+
+	wafModel := &models.Waf{Ip: ip}
+	ban, err := w.db.Get(wafModel)
+	if err != nil {
+		logger.Log.Errorf("[waf] access 出错 %s", err)
+		return false
+	}
+
+	if ban {
+		logger.Log.Infof("[waf] 阻止黑名单 ip:%s 访问", ip)
+		return false
+	}
+
 	ctx := context.Background()
 	num, err := w.store.LLen(ctx, key).Result()
 	if err != nil {
@@ -52,15 +66,17 @@ func (w *Waf) Access(ip string) bool {
 
 	if num < w.config.RateLimit {
 		w.store.LPush(ctx, key, time.Now().Unix())
+		w.store.Expire(ctx, key, 300*time.Second)
 		return true
 	} else {
 		last, _ := w.store.LIndex(ctx, key, -1).Int64()
 		if time.Now().Unix()-last < 60 {
-			logger.Log.Infof("[waf] ip:%s 访问频率超过限制 %d", ip, w.config.RateLimit)
+			logger.Log.Warnf("[waf] ip:%s 访问频率超过限制 %d", ip, w.config.RateLimit)
 			return false
 		} else {
 			w.store.LPush(ctx, key, time.Now().Unix())
 			w.store.LTrim(ctx, key, 0, w.config.RateLimit-1)
+			w.store.Expire(ctx, key, 300*time.Second)
 			return true
 		}
 	}