From 37d46eda0eacd905da160cf010937b41581e45fe Mon Sep 17 00:00:00 2001
From: Shikong <919411476@qq.com>
Date: Tue, 25 Oct 2022 16:23:35 +0800
Subject: [PATCH] =?UTF-8?q?=E8=AE=BF=E9=97=AE=E9=A2=91=E7=8E=87=E9=99=90?=
 =?UTF-8?q?=E5=88=B6=20=E4=BC=98=E5=8C=96=20=E6=B7=BB=E5=8A=A0=E9=BB=91?=
 =?UTF-8?q?=E5=90=8D=E5=8D=95=E6=9C=BA=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 services/waf/waf.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/services/waf/waf.go b/services/waf/waf.go
index ebfaef0..2560f09 100644
--- a/services/waf/waf.go
+++ b/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
 		}
 	}