From b29559d112ab2a30b78bdfcff5841c5d41db3ff4 Mon Sep 17 00:00:00 2001 From: GoEdgeLab Date: Thu, 30 Nov 2023 17:25:41 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E9=AA=8C=E8=AF=81=E7=A0=81?= =?UTF-8?q?=E6=80=A7=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/waf/captcha_generator.go | 71 +++++++++++++++++++++ internal/waf/captcha_generator_test.go | 87 ++++++++++++++++++++++++++ internal/waf/captcha_validator.go | 8 +-- 3 files changed, 162 insertions(+), 4 deletions(-) create mode 100644 internal/waf/captcha_generator.go create mode 100644 internal/waf/captcha_generator_test.go diff --git a/internal/waf/captcha_generator.go b/internal/waf/captcha_generator.go new file mode 100644 index 0000000..d5ba03b --- /dev/null +++ b/internal/waf/captcha_generator.go @@ -0,0 +1,71 @@ +// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn . + +package waf + +import ( + "bytes" + "github.com/dchest/captcha" + "github.com/iwind/TeaGo/rands" + "io" + "time" +) + +// CaptchaGenerator captcha generator +type CaptchaGenerator struct { + store captcha.Store +} + +func NewCaptchaGenerator() *CaptchaGenerator { + return &CaptchaGenerator{ + store: captcha.NewMemoryStore(100_000, 5*time.Minute), + } +} + +// NewCaptcha create new captcha +func (this *CaptchaGenerator) NewCaptcha(length int) (captchaId string) { + captchaId = rands.HexString(16) + + if length <= 0 || length > 20 { + length = 4 + } + + this.store.Set(captchaId, captcha.RandomDigits(length)) + return +} + +// WriteImage write image to front writer +func (this *CaptchaGenerator) WriteImage(w io.Writer, id string, width, height int) error { + var d = this.store.Get(id, false) + if d == nil { + return captcha.ErrNotFound + } + _, err := captcha.NewImage(id, d, width, height).WriteTo(w) + return err +} + +// Verify user input +func (this *CaptchaGenerator) Verify(id string, digits string) bool { + var countDigits = len(digits) + if countDigits == 0 { + return false + } + var value = this.store.Get(id, true) + if len(value) != countDigits { + return false + } + + var nb = make([]byte, countDigits) + for i := 0; i < countDigits; i++ { + var d = digits[i] + if d >= '0' && d <= '9' { + nb[i] = d - '0' + } + } + + return bytes.Equal(nb, value) +} + +// Get captcha data +func (this *CaptchaGenerator) Get(id string) []byte { + return this.store.Get(id, false) +} diff --git a/internal/waf/captcha_generator_test.go b/internal/waf/captcha_generator_test.go new file mode 100644 index 0000000..d32eeeb --- /dev/null +++ b/internal/waf/captcha_generator_test.go @@ -0,0 +1,87 @@ +// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn . + +package waf_test + +import ( + "github.com/TeaOSLab/EdgeNode/internal/utils/testutils" + "github.com/TeaOSLab/EdgeNode/internal/waf" + "github.com/iwind/TeaGo/assert" + "github.com/iwind/TeaGo/types" + "runtime" + "strings" + "testing" + "time" +) + +func TestCaptchaGenerator_NewCaptcha(t *testing.T) { + var a = assert.NewAssertion(t) + + var generator = waf.NewCaptchaGenerator() + var captchaId = generator.NewCaptcha(6) + t.Log("captchaId:", captchaId) + + var digits = generator.Get(captchaId) + var s []string + for _, digit := range digits { + s = append(s, types.String(digit)) + } + t.Log(strings.Join(s, " ")) + + a.IsTrue(generator.Verify(captchaId, strings.Join(s, ""))) + a.IsFalse(generator.Verify(captchaId, strings.Join(s, ""))) +} + +func TestCaptchaGenerator_NewCaptcha_UTF8(t *testing.T) { + var a = assert.NewAssertion(t) + + var generator = waf.NewCaptchaGenerator() + var captchaId = generator.NewCaptcha(6) + t.Log("captchaId:", captchaId) + + var digits = generator.Get(captchaId) + var s []string + for _, digit := range digits { + s = append(s, types.String(digit)) + } + t.Log(strings.Join(s, " ")) + + a.IsFalse(generator.Verify(captchaId, "中文真的很长")) +} + +func TestCaptchaGenerator_NewCaptcha_Memory(t *testing.T) { + runtime.GC() + + var stat1 = &runtime.MemStats{} + runtime.ReadMemStats(stat1) + + var generator = waf.NewCaptchaGenerator() + for i := 0; i < 1_000_000; i++ { + generator.NewCaptcha(6) + } + + if testutils.IsSingleTesting() { + time.Sleep(1 * time.Second) + } + + runtime.GC() + + var stat2 = &runtime.MemStats{} + runtime.ReadMemStats(stat2) + + t.Log((stat2.HeapInuse-stat1.HeapInuse)>>10, "KiB") + + _ = generator +} + +func BenchmarkNewCaptchaGenerator(b *testing.B) { + runtime.GOMAXPROCS(4) + + var generator = waf.NewCaptchaGenerator() + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + generator.NewCaptcha(6) + } + }) +} diff --git a/internal/waf/captcha_validator.go b/internal/waf/captcha_validator.go index 1c1bfdc..8ead8b1 100644 --- a/internal/waf/captcha_validator.go +++ b/internal/waf/captcha_validator.go @@ -12,7 +12,6 @@ import ( "github.com/TeaOSLab/EdgeNode/internal/utils/fasttime" "github.com/TeaOSLab/EdgeNode/internal/waf/requests" wafutils "github.com/TeaOSLab/EdgeNode/internal/waf/utils" - "github.com/dchest/captcha" "github.com/iwind/TeaGo/logs" "github.com/iwind/TeaGo/maps" "github.com/iwind/TeaGo/rands" @@ -29,6 +28,7 @@ import ( const captchaIdName = "GOEDGE_WAF_CAPTCHA_ID" var captchaValidator = NewCaptchaValidator() +var captchaGenerator = NewCaptchaGenerator() type CaptchaValidator struct { } @@ -162,7 +162,7 @@ func (this *CaptchaValidator) showVerifyCodesForm(actionConfig *CaptchaAction, r if actionConfig.CountLetters > 0 && actionConfig.CountLetters <= 10 { countLetters = int(actionConfig.CountLetters) } - var captchaId = captcha.NewLen(countLetters) + var captchaId = captchaGenerator.NewCaptcha(countLetters) var lang = actionConfig.Lang if len(lang) == 0 { @@ -305,7 +305,7 @@ func (this *CaptchaValidator) showVerifyImage(actionConfig *CaptchaAction, req r } writer.Header().Set("Content-Type", "image/png") - err := captcha.WriteImage(writer, captchaId, 200, 100) + err := captchaGenerator.WriteImage(writer, captchaId, 200, 100) if err != nil { logs.Error(err) return @@ -316,7 +316,7 @@ func (this *CaptchaValidator) validateVerifyCodeForm(actionConfig *CaptchaAction var captchaId = req.WAFRaw().FormValue(captchaIdName) if len(captchaId) > 0 { var captchaCode = req.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_CODE") - if captcha.VerifyString(captchaId, captchaCode) { + if captchaGenerator.Verify(captchaId, captchaCode) { // 清除计数 CaptchaDeleteCacheKey(req)