优化验证码性能

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

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

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)