增加图形验证码的功能

captcha/captcha_test.go

@@ -0,0 +1,174 @@
+package captcha_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"git.toowon.com/jimmy/go-common/captcha"
+	"git.toowon.com/jimmy/go-common/config"
+	"github.com/alicebob/miniredis/v2"
+	"github.com/redis/go-redis/v9"
+)
+
+func TestDisabledManager(t *testing.T) {
+	m := captcha.NewManager(nil, nil)
+	if m.Enabled() {
+		t.Fatal("expected disabled")
+	}
+	_, err := m.Generate(context.Background(), captcha.SceneRegister)
+	if err != captcha.ErrDisabled {
+		t.Fatalf("Generate err = %v, want ErrDisabled", err)
+	}
+}
+
+func TestGenerateAndVerify(t *testing.T) {
+	s, err := miniredis.Run()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+
+	rdb := redis.NewClient(&redis.Options{Addr: s.Addr()})
+	cfg := &config.CaptchaConfig{
+		Enabled: true,
+		Length:  4,
+		Width:   120,
+		Height:  40,
+		TTLSec:  300,
+		Login:   &config.LoginCaptchaConfig{},
+	}
+	m := captcha.NewManager(rdb, cfg)
+	if !m.Enabled() {
+		t.Fatal("expected enabled")
+	}
+
+	ctx := context.Background()
+	result, err := m.Generate(ctx, captcha.SceneRegister)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if result.ID == "" || result.ImageBase64 == "" {
+		t.Fatal("expected id and image")
+	}
+
+	// 从 redis 读取答案用于测试
+	answer, err := rdb.Get(ctx, "captcha:code:"+result.ID).Result()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := m.Verify(ctx, captcha.VerifyRequest{
+		Scene:  captcha.SceneRegister,
+		ID:     result.ID,
+		Answer: answer,
+	}); err != nil {
+		t.Fatalf("Verify: %v", err)
+	}
+	if err := m.Verify(ctx, captcha.VerifyRequest{
+		Scene:  captcha.SceneRegister,
+		ID:     result.ID,
+		Answer: answer,
+	}); err != captcha.ErrInvalid {
+		t.Fatalf("second Verify err = %v, want ErrInvalid", err)
+	}
+}
+
+func TestLoginGuard(t *testing.T) {
+	s, err := miniredis.Run()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+
+	rdb := redis.NewClient(&redis.Options{Addr: s.Addr()})
+	cfg := &config.CaptchaConfig{
+		Enabled: true,
+		Length:  4,
+		Width:   120,
+		Height:  40,
+		TTLSec:  300,
+		Login: &config.LoginCaptchaConfig{
+			ShowAfterFailures: 2,
+			LockAfterFailures: 4,
+			LockDurationMin:   1,
+		},
+	}
+	m := captcha.NewManager(rdb, cfg)
+	ctx := context.Background()
+	key := "alice"
+
+	status, err := m.RecordLoginFailure(ctx, key)
+	if err != nil || status.FailCount != 1 || status.NeedCaptcha {
+		t.Fatalf("first failure: status=%+v err=%v", status, err)
+	}
+
+	status, err = m.RecordLoginFailure(ctx, key)
+	if err != nil || status.FailCount != 2 || !status.NeedCaptcha {
+		t.Fatalf("second failure: status=%+v err=%v", status, err)
+	}
+
+	need, err := m.LoginNeedCaptcha(ctx, key)
+	if err != nil || !need {
+		t.Fatalf("LoginNeedCaptcha = %v, err = %v", need, err)
+	}
+
+	for i := 0; i < 2; i++ {
+		_, _ = m.RecordLoginFailure(ctx, key)
+	}
+	status, err = m.RecordLoginFailure(ctx, key)
+	if err != captcha.ErrLoginLocked || !status.Locked {
+		t.Fatalf("lock: status=%+v err=%v", status, err)
+	}
+
+	locked, remaining, err := m.IsLoginLocked(ctx, key)
+	if err != nil || !locked || remaining <= 0 {
+		t.Fatalf("IsLoginLocked = %v, %d, err = %v", locked, remaining, err)
+	}
+
+	if err := m.ClearLoginFailures(ctx, key); err != nil {
+		t.Fatal(err)
+	}
+	locked, _, _ = m.IsLoginLocked(ctx, key)
+	if locked {
+		t.Fatal("expected unlocked after clear")
+	}
+}
+
+func TestSceneRequiresCaptcha(t *testing.T) {
+	cfg := &config.CaptchaConfig{
+		Enabled: true,
+		Scenes: map[string]config.CaptchaSceneConfig{
+			"register": {Mode: "always"},
+			"login":    {Mode: "never"},
+		},
+	}
+	m := captcha.NewManager(redis.NewClient(&redis.Options{Addr: "127.0.0.1:1"}), cfg)
+	if !m.SceneRequiresCaptcha(captcha.SceneRegister) {
+		t.Fatal("register should require captcha")
+	}
+	if m.SceneRequiresCaptcha(captcha.SceneLogin) {
+		t.Fatal("login should not always require captcha")
+	}
+}
+
+func TestRedisStoreTTL(t *testing.T) {
+	s, err := miniredis.Run()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+
+	rdb := redis.NewClient(&redis.Options{Addr: s.Addr()})
+	cfg := &config.CaptchaConfig{Enabled: true, TTLSec: 1, Login: &config.LoginCaptchaConfig{}}
+	m := captcha.NewManager(rdb, cfg)
+	ctx := context.Background()
+
+	result, err := m.Generate(ctx, captcha.SceneRegister)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s.FastForward(2 * time.Second)
+	if err := m.Verify(ctx, captcha.VerifyRequest{ID: result.ID, Answer: "x"}); err != captcha.ErrInvalid {
+		t.Fatalf("expired verify err = %v", err)
+	}
+}