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 } }