go-common/captcha/captcha_test.go

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