优化验证码失败次数统计

internal/nodes/upgrade_manager.go

@@ -13,7 +13,6 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/rpc"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/iwind/TeaGo/Tea"
-	"github.com/iwind/TeaGo/logs"
 	stringutil "github.com/iwind/TeaGo/utils/string"
 	"os"
 	"os/exec"
@@ -64,7 +63,7 @@ func (this *UpgradeManager) Start() {
 	goman.New(func() {
 		err = this.restart()
 		if err != nil {
-			logs.Println("UPGRADE_MANAGER", err.Error())
+			remotelogs.Error("UPGRADE_MANAGER", err.Error())
 		}
 	})
 }

internal/waf/action_captcha.go

@@ -98,17 +98,17 @@ func (this *CaptchaAction) WillChange() bool {
 	return true
 }
 
-func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool) {
+func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req requests.Request, writer http.ResponseWriter) (allow bool) {
 	// 是否在白名单中
-	if SharedIPWhiteList.Contains("set:"+types.String(set.Id), this.Scope, request.WAFServerId(), request.WAFRemoteIP()) {
+	if SharedIPWhiteList.Contains("set:"+types.String(set.Id), this.Scope, req.WAFServerId(), req.WAFRemoteIP()) {
 		return true
 	}
 
-	var refURL = request.WAFRaw().URL.String()
+	var refURL = req.WAFRaw().URL.String()
 
 	// 覆盖配置
 	if strings.HasPrefix(refURL, CaptchaPath) {
-		info := request.WAFRaw().URL.Query().Get("info")
+		info := req.WAFRaw().URL.Query().Get("info")
 		if len(info) > 0 {
 			m, err := utils.SimpleDecryptMap(info)
 			if err == nil && m != nil {
@@ -131,7 +131,10 @@ func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 		return true
 	}
 
-	http.Redirect(writer, request.WAFRaw(), CaptchaPath+"?info="+url.QueryEscape(info), http.StatusTemporaryRedirect)
+	// 占用一次失败次数
+	CaptchaIncreaseFails(req, this, waf.Id, group.Id, set.Id, CaptchaPageCodeInit)
+
+	http.Redirect(writer, req.WAFRaw(), CaptchaPath+"?info="+url.QueryEscape(info), http.StatusTemporaryRedirect)
 
 	return false
 }

internal/waf/captcha_counter.go

@@ -0,0 +1,54 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package waf
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
+	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
+	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
+	"github.com/iwind/TeaGo/types"
+	"time"
+)
+
+type CaptchaPageCode = string
+
+const (
+	CaptchaPageCodeInit   CaptchaPageCode = "init"
+	CaptchaPageCodeShow   CaptchaPageCode = "show"
+	CaptchaPageCodeSubmit CaptchaPageCode = "submit"
+)
+
+// CaptchaIncreaseFails 增加Captcha失败次数，以便后续操作
+func CaptchaIncreaseFails(req requests.Request, actionConfig *CaptchaAction, policyId int64, groupId int64, setId int64, pageCode CaptchaPageCode) (goNext bool) {
+	var maxFails = actionConfig.MaxFails
+	var failBlockTimeout = actionConfig.FailBlockTimeout
+	if maxFails > 0 && failBlockTimeout > 0 {
+		if maxFails <= 3 {
+			maxFails = 3 // 不能小于3，防止意外刷新出现
+		}
+		var countFails = ttlcache.SharedCache.IncreaseInt64(CaptchaCacheKey(req, pageCode), 1, time.Now().Unix()+300, true)
+		if int(countFails) >= maxFails {
+			var useLocalFirewall = false
+
+			if actionConfig.FailBlockScopeAll {
+				useLocalFirewall = true
+			}
+
+			SharedIPBlackList.RecordIP(IPTypeAll, firewallconfigs.FirewallScopeService, req.WAFServerId(), req.WAFRemoteIP(), time.Now().Unix()+int64(failBlockTimeout), policyId, useLocalFirewall, groupId, setId, "CAPTCHA验证连续失败")
+			return false
+		}
+	}
+	return true
+}
+
+// CaptchaDeleteCacheKey 清除计数
+func CaptchaDeleteCacheKey(req requests.Request) {
+	ttlcache.SharedCache.Delete(CaptchaCacheKey(req, CaptchaPageCodeInit))
+	ttlcache.SharedCache.Delete(CaptchaCacheKey(req, CaptchaPageCodeShow))
+	ttlcache.SharedCache.Delete(CaptchaCacheKey(req, CaptchaPageCodeSubmit))
+}
+
+// CaptchaCacheKey 获取Captcha缓存Key
+func CaptchaCacheKey(req requests.Request, pageCode CaptchaPageCode) string {
+	return "CAPTCHA:FAILS:" + pageCode + ":" + req.WAFRemoteIP() + ":" + types.String(req.WAFServerId())
+}

internal/waf/captcha_validator.go

@@ -3,8 +3,6 @@ package waf
 import (
 	"bytes"
 	"encoding/base64"
-	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
-	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
 	"github.com/dchest/captcha"
@@ -70,7 +68,7 @@ func (this *CaptchaValidator) Run(req requests.Request, writer http.ResponseWrit
 		this.validate(captchaActionConfig, policyId, groupId, setId, originURL, req, writer)
 	} else {
 		// 增加计数
-		this.IncreaseFails(req, captchaActionConfig, policyId, groupId, setId)
+		CaptchaIncreaseFails(req, captchaActionConfig, policyId, groupId, setId, CaptchaPageCodeShow)
 		this.show(captchaActionConfig, req, writer)
 	}
 }
@@ -129,6 +127,29 @@ func (this *CaptchaValidator) show(actionConfig *CaptchaAction, req requests.Req
 	var msgCss = ""
 	var requestIdBox = `<address>` + msgRequestId + `: ` + req.Format("${requestId}") + `</address>`
 	var msgFooter = ""
+
+	// 默认设置
+	if actionConfig.UIIsOn {
+		if len(actionConfig.UIPrompt) > 0 {
+			msgPrompt = actionConfig.UIPrompt
+		}
+		if len(actionConfig.UIButtonTitle) > 0 {
+			msgButtonTitle = actionConfig.UIButtonTitle
+		}
+		if len(actionConfig.UITitle) > 0 {
+			msgTitle = actionConfig.UITitle
+		}
+		if len(actionConfig.UICss) > 0 {
+			msgCss = actionConfig.UICss
+		}
+		if !actionConfig.UIShowRequestId {
+			requestIdBox = ""
+		}
+		if len(actionConfig.UIFooter) > 0 {
+			msgFooter = actionConfig.UIFooter
+		}
+	}
+
 	var body = `<form method="POST">
 	<input type="hidden" name="GOEDGE_WAF_CAPTCHA_ID" value="` + captchaId + `"/>
 	<div class="ui-image">
@@ -145,26 +166,8 @@ func (this *CaptchaValidator) show(actionConfig *CaptchaAction, req requests.Req
 ` + requestIdBox + `
 ` + msgFooter + ``
 
-	// 默认设置
+	// Body
 	if actionConfig.UIIsOn {
-		if len(actionConfig.UITitle) > 0 {
-			msgTitle = actionConfig.UITitle
-		}
-		if len(actionConfig.UIPrompt) > 0 {
-			msgPrompt = actionConfig.UIPrompt
-		}
-		if len(actionConfig.UIButtonTitle) > 0 {
-			msgButtonTitle = actionConfig.UIButtonTitle
-		}
-		if len(actionConfig.UICss) > 0 {
-			msgCss = actionConfig.UICss
-		}
-		if !actionConfig.UIShowRequestId {
-			requestIdBox = ""
-		}
-		if len(actionConfig.UIFooter) > 0 {
-			msgFooter = actionConfig.UIFooter
-		}
 		if len(actionConfig.UIBody) > 0 {
 			var index = strings.Index(actionConfig.UIBody, "${body}")
 			if index < 0 {
@@ -208,7 +211,7 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, policyId int
 		var captchaCode = req.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_CODE")
 		if captcha.VerifyString(captchaId, captchaCode) {
 			// 清除计数
-			ttlcache.SharedCache.Delete(this.cacheKey(req))
+			CaptchaDeleteCacheKey(req)
 
 			var life = CaptchaSeconds
 			if actionConfig.Life > 0 {
@@ -223,7 +226,7 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, policyId int
 			return false
 		} else {
 			// 增加计数
-			if !this.IncreaseFails(req, actionConfig, policyId, groupId, setId) {
+			if !CaptchaIncreaseFails(req, actionConfig, policyId, groupId, setId, CaptchaPageCodeSubmit) {
 				return false
 			}
 
@@ -233,30 +236,3 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, policyId int
 
 	return true
 }
-
-// IncreaseFails 增加失败次数，以便后续操作
-func (this *CaptchaValidator) IncreaseFails(req requests.Request, actionConfig *CaptchaAction, policyId int64, groupId int64, setId int64) (goNext bool) {
-	var maxFails = actionConfig.MaxFails
-	var failBlockTimeout = actionConfig.FailBlockTimeout
-	if maxFails > 0 && failBlockTimeout > 0 {
-		// 加上展示的计数
-		maxFails *= 2
-
-		var countFails = ttlcache.SharedCache.IncreaseInt64(this.cacheKey(req), 1, time.Now().Unix()+300, true)
-		if int(countFails) >= maxFails {
-			var useLocalFirewall = false
-
-			if actionConfig.FailBlockScopeAll {
-				useLocalFirewall = true
-			}
-
-			SharedIPBlackList.RecordIP(IPTypeAll, firewallconfigs.FirewallScopeService, req.WAFServerId(), req.WAFRemoteIP(), time.Now().Unix()+int64(failBlockTimeout), policyId, useLocalFirewall, groupId, setId, "CAPTCHA验证连续失败")
-			return false
-		}
-	}
-	return true
-}
-
-func (this *CaptchaValidator) cacheKey(req requests.Request) string {
-	return "CAPTCHA:FAILS:" + req.WAFRemoteIP() + ":" + types.String(req.WAFServerId())
-}