diff --git a/internal/waf/captcha_counter.go b/internal/waf/captcha_counter.go index e67b6f6..abaf389 100644 --- a/internal/waf/captcha_counter.go +++ b/internal/waf/captcha_counter.go @@ -15,6 +15,7 @@ type CaptchaPageCode = string const ( CaptchaPageCodeInit CaptchaPageCode = "init" CaptchaPageCodeShow CaptchaPageCode = "show" + CaptchaPageCodeImage CaptchaPageCode = "image" CaptchaPageCodeSubmit CaptchaPageCode = "submit" ) @@ -39,6 +40,7 @@ func CaptchaIncreaseFails(req requests.Request, actionConfig *CaptchaAction, pol func CaptchaDeleteCacheKey(req requests.Request) { counters.SharedCounter.ResetKey(CaptchaCacheKey(req, CaptchaPageCodeInit)) counters.SharedCounter.ResetKey(CaptchaCacheKey(req, CaptchaPageCodeShow)) + counters.SharedCounter.ResetKey(CaptchaCacheKey(req, CaptchaPageCodeImage)) counters.SharedCounter.ResetKey(CaptchaCacheKey(req, CaptchaPageCodeSubmit)) } diff --git a/internal/waf/captcha_test.go b/internal/waf/captcha_test.go new file mode 100644 index 0000000..ea14bdb --- /dev/null +++ b/internal/waf/captcha_test.go @@ -0,0 +1,70 @@ +// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn . + +package waf_test + +import ( + "bytes" + "fmt" + "github.com/TeaOSLab/EdgeNode/internal/utils/testutils" + "github.com/dchest/captcha" + "runtime" + "testing" + "time" +) + +func TestCaptchaMemory(t *testing.T) { + if !testutils.IsSingleTesting() { + return + } + + var stat1 = &runtime.MemStats{} + runtime.ReadMemStats(stat1) + + var count = 5_000 + var before = time.Now() + + for i := 0; i < count; i++ { + var id = captcha.NewLen(6) + var writer = &bytes.Buffer{} + err := captcha.WriteImage(writer, id, 200, 100) + if err != nil { + t.Fatal(err) + } + captcha.VerifyString(id, "abc") + } + + var stat2 = &runtime.MemStats{} + runtime.ReadMemStats(stat2) + t.Log((stat2.HeapInuse-stat1.HeapInuse)>>20, "MB", fmt.Sprintf("%.0f QPS", float64(count)/time.Since(before).Seconds())) +} + +func BenchmarkCaptcha_VerifyCode_100_50(b *testing.B) { + runtime.GOMAXPROCS(4) + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + var id = captcha.NewLen(6) + var writer = &bytes.Buffer{} + err := captcha.WriteImage(writer, id, 100, 50) + if err != nil { + b.Fatal(err) + } + } + }) +} + +func BenchmarkCaptcha_VerifyCode_200_100(b *testing.B) { + runtime.GOMAXPROCS(4) + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + var id = captcha.NewLen(6) + var writer = &bytes.Buffer{} + err := captcha.WriteImage(writer, id, 200, 100) + if err != nil { + b.Fatal(err) + } + _ = id + } + }) +} diff --git a/internal/waf/captcha_validator.go b/internal/waf/captcha_validator.go index 78705df..7346192 100644 --- a/internal/waf/captcha_validator.go +++ b/internal/waf/captcha_validator.go @@ -1,8 +1,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" @@ -19,6 +17,8 @@ import ( "time" ) +const captchaIdName = "GOEDGE_WAF_CAPTCHA_ID" + var captchaValidator = NewCaptchaValidator() type CaptchaValidator struct { @@ -81,7 +81,7 @@ func (this *CaptchaValidator) Run(req requests.Request, writer http.ResponseWrit captchaType = defaultCaptchaType } - if req.WAFRaw().Method == http.MethodPost && len(req.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_ID")) > 0 { + if req.WAFRaw().Method == http.MethodPost && len(req.WAFRaw().FormValue(captchaIdName)) > 0 { switch captchaType { case firewallconfigs.CaptchaTypeOneClick: this.validateOneClickForm(captchaActionConfig, policyId, groupId, setId, originURL, req, writer) @@ -91,9 +91,16 @@ func (this *CaptchaValidator) Run(req requests.Request, writer http.ResponseWrit this.validateVerifyCodeForm(captchaActionConfig, policyId, groupId, setId, originURL, req, writer) } } else { - // 增加计数 - CaptchaIncreaseFails(req, captchaActionConfig, policyId, groupId, setId, CaptchaPageCodeShow) - this.show(captchaActionConfig, req, writer, captchaType) + var captchaId = req.WAFRaw().URL.Query().Get(captchaIdName) + if len(captchaId) > 0 { + // 增加计数 + CaptchaIncreaseFails(req, captchaActionConfig, policyId, groupId, setId, CaptchaPageCodeImage) + this.showImage(captchaActionConfig, req, writer, captchaType) + } else { + // 增加计数 + CaptchaIncreaseFails(req, captchaActionConfig, policyId, groupId, setId, CaptchaPageCodeShow) + this.show(captchaActionConfig, req, writer, captchaType) + } } } @@ -108,6 +115,17 @@ func (this *CaptchaValidator) show(actionConfig *CaptchaAction, req requests.Req } } +func (this *CaptchaValidator) showImage(actionConfig *CaptchaAction, req requests.Request, writer http.ResponseWriter, captchaType firewallconfigs.ServerCaptchaType) { + switch captchaType { + case firewallconfigs.CaptchaTypeOneClick: + // stub + case firewallconfigs.CaptchaTypeSlide: + // stub + default: + this.showVerifyImage(actionConfig, req, writer) + } +} + func (this *CaptchaValidator) showVerifyCodesForm(actionConfig *CaptchaAction, req requests.Request, writer http.ResponseWriter) { // show captcha var countLetters = 6 @@ -115,12 +133,6 @@ func (this *CaptchaValidator) showVerifyCodesForm(actionConfig *CaptchaAction, r countLetters = int(actionConfig.CountLetters) } var captchaId = captcha.NewLen(countLetters) - var buf = bytes.NewBuffer([]byte{}) - err := captcha.WriteImage(buf, captchaId, 200, 100) - if err != nil { - logs.Error(err) - return - } var lang = actionConfig.Lang if len(lang) == 0 { @@ -191,9 +203,10 @@ func (this *CaptchaValidator) showVerifyCodesForm(actionConfig *CaptchaAction, r } var body = `
- +
- ` + ` +

loading ...

+

` + msgPrompt + `

@@ -226,9 +239,15 @@ func (this *CaptchaValidator) showVerifyCodesForm(actionConfig *CaptchaAction, r