增加图形验证码的功能
This commit is contained in:
307
captcha/captcha.go
Normal file
307
captcha/captcha.go
Normal file
@@ -0,0 +1,307 @@
|
||||
package captcha
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.toowon.com/jimmy/go-common/config"
|
||||
"github.com/mojocn/base64Captcha"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// Scene 验证码业务场景。
|
||||
type Scene string
|
||||
|
||||
const (
|
||||
SceneRegister Scene = "register"
|
||||
SceneChangePhone Scene = "change_phone"
|
||||
SceneLogin Scene = "login"
|
||||
SceneSendSMS Scene = "send_sms"
|
||||
)
|
||||
|
||||
// GenerateResult 生成图形验证码的返回。
|
||||
type GenerateResult struct {
|
||||
ID string `json:"id"`
|
||||
ImageBase64 string `json:"imageBase64"`
|
||||
}
|
||||
|
||||
// VerifyRequest 校验图形验证码的请求。
|
||||
type VerifyRequest struct {
|
||||
Scene Scene
|
||||
ID string
|
||||
Answer string
|
||||
}
|
||||
|
||||
// LoginGuardStatus 登录防暴力状态。
|
||||
type LoginGuardStatus struct {
|
||||
FailCount int `json:"failCount"`
|
||||
NeedCaptcha bool `json:"needCaptcha"`
|
||||
Locked bool `json:"locked"`
|
||||
RemainingSec int `json:"remainingSec"`
|
||||
}
|
||||
|
||||
// Manager 图形验证码子系统（`app.Captcha()` 取得）。
|
||||
//
|
||||
// 职责：
|
||||
// - 生成/校验图形验证码（Redis 存储，一次性消费）；
|
||||
// - 登录失败计数、触发验证码、临时锁定；
|
||||
// - Enabled() 反映 config.captcha.enabled 且 Redis 可用。
|
||||
//
|
||||
// 不含路由与业务编排：消费方在 handler 中按需调用 Generate / Verify / RecordLoginFailure。
|
||||
type Manager interface {
|
||||
Enabled() bool
|
||||
Generate(ctx context.Context, scene Scene) (*GenerateResult, error)
|
||||
Verify(ctx context.Context, req VerifyRequest) error
|
||||
SceneRequiresCaptcha(scene Scene) bool
|
||||
LoginNeedCaptcha(ctx context.Context, key string) (bool, error)
|
||||
IsLoginLocked(ctx context.Context, key string) (bool, int, error)
|
||||
RecordLoginFailure(ctx context.Context, key string) (LoginGuardStatus, error)
|
||||
ClearLoginFailures(ctx context.Context, key string) error
|
||||
}
|
||||
|
||||
type manager struct {
|
||||
enabled bool
|
||||
cfg *config.CaptchaConfig
|
||||
rdb *redis.Client
|
||||
captcha *base64Captcha.Captcha
|
||||
}
|
||||
|
||||
// NewManager 根据配置与 Redis 构建图形验证码子系统。
|
||||
// cfg 为 nil 或 enabled=false 时返回 Enabled()=false 的安全实例，其余方法返回 ErrDisabled。
|
||||
// enabled=true 但 rdb=nil 时同样视为不可用（Enabled()=false）。
|
||||
func NewManager(rdb *redis.Client, cfg *config.CaptchaConfig) Manager {
|
||||
m := &manager{cfg: cfg}
|
||||
if cfg == nil || !cfg.Enabled || rdb == nil {
|
||||
return m
|
||||
}
|
||||
|
||||
normalizeCaptchaConfig(cfg)
|
||||
|
||||
ttl := time.Duration(cfg.TTLSec) * time.Second
|
||||
store := newRedisStore(rdb, ttl)
|
||||
driver := base64Captcha.NewDriverDigit(cfg.Height, cfg.Width, cfg.Length, 0.7, 80)
|
||||
m.enabled = true
|
||||
m.rdb = rdb
|
||||
m.captcha = base64Captcha.NewCaptcha(driver, store)
|
||||
return m
|
||||
}
|
||||
|
||||
func normalizeCaptchaConfig(cfg *config.CaptchaConfig) {
|
||||
if cfg.Length == 0 {
|
||||
cfg.Length = 4
|
||||
}
|
||||
if cfg.Width == 0 {
|
||||
cfg.Width = 120
|
||||
}
|
||||
if cfg.Height == 0 {
|
||||
cfg.Height = 40
|
||||
}
|
||||
if cfg.TTLSec == 0 {
|
||||
cfg.TTLSec = 300
|
||||
}
|
||||
if cfg.Login == nil {
|
||||
cfg.Login = &config.LoginCaptchaConfig{}
|
||||
}
|
||||
if cfg.Login.ShowAfterFailures == 0 {
|
||||
cfg.Login.ShowAfterFailures = 3
|
||||
}
|
||||
if cfg.Login.LockAfterFailures == 0 {
|
||||
cfg.Login.LockAfterFailures = 10
|
||||
}
|
||||
if cfg.Login.LockDurationMin == 0 {
|
||||
cfg.Login.LockDurationMin = 30
|
||||
}
|
||||
}
|
||||
|
||||
func (m *manager) Enabled() bool {
|
||||
return m != nil && m.enabled
|
||||
}
|
||||
|
||||
func (m *manager) Generate(ctx context.Context, scene Scene) (*GenerateResult, error) {
|
||||
if err := m.requireEnabled(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = ctx
|
||||
|
||||
id, b64s, _, err := m.captcha.Generate()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("captcha generate: %w", err)
|
||||
}
|
||||
return &GenerateResult{
|
||||
ID: id,
|
||||
ImageBase64: b64s,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *manager) Verify(ctx context.Context, req VerifyRequest) error {
|
||||
if err := m.requireEnabled(); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = ctx
|
||||
|
||||
id := strings.TrimSpace(req.ID)
|
||||
answer := strings.TrimSpace(req.Answer)
|
||||
if id == "" || answer == "" {
|
||||
return ErrInvalid
|
||||
}
|
||||
if !m.captcha.Verify(id, answer, true) {
|
||||
return ErrInvalid
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *manager) SceneRequiresCaptcha(scene Scene) bool {
|
||||
if !m.Enabled() {
|
||||
return false
|
||||
}
|
||||
if scene == SceneLogin {
|
||||
return false
|
||||
}
|
||||
mode := m.sceneMode(scene)
|
||||
return mode == "always"
|
||||
}
|
||||
|
||||
func (m *manager) LoginNeedCaptcha(ctx context.Context, key string) (bool, error) {
|
||||
if !m.Enabled() {
|
||||
return false, nil
|
||||
}
|
||||
loginCfg := m.loginConfig()
|
||||
if loginCfg.ShowAfterFailures <= 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
locked, _, err := m.IsLoginLocked(ctx, key)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if locked {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
count, err := m.getLoginFailCount(ctx, key)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count >= loginCfg.ShowAfterFailures, nil
|
||||
}
|
||||
|
||||
func (m *manager) IsLoginLocked(ctx context.Context, key string) (bool, int, error) {
|
||||
if !m.Enabled() {
|
||||
return false, 0, nil
|
||||
}
|
||||
ttl, err := m.rdb.TTL(ctx, loginLockKey(key)).Result()
|
||||
if err != nil {
|
||||
return false, 0, err
|
||||
}
|
||||
if ttl <= 0 {
|
||||
return false, 0, nil
|
||||
}
|
||||
return true, int(ttl.Seconds()), nil
|
||||
}
|
||||
|
||||
func (m *manager) RecordLoginFailure(ctx context.Context, key string) (LoginGuardStatus, error) {
|
||||
status := LoginGuardStatus{}
|
||||
if !m.Enabled() {
|
||||
return status, nil
|
||||
}
|
||||
|
||||
loginCfg := m.loginConfig()
|
||||
lockKey := loginLockKey(key)
|
||||
failKey := loginFailKey(key)
|
||||
|
||||
if ttl, err := m.rdb.TTL(ctx, lockKey).Result(); err != nil {
|
||||
return status, err
|
||||
} else if ttl > 0 {
|
||||
status.Locked = true
|
||||
status.NeedCaptcha = true
|
||||
status.RemainingSec = int(ttl.Seconds())
|
||||
return status, ErrLoginLocked
|
||||
}
|
||||
|
||||
count, err := m.rdb.Incr(ctx, failKey).Result()
|
||||
if err != nil {
|
||||
return status, err
|
||||
}
|
||||
if count == 1 {
|
||||
_ = m.rdb.Expire(ctx, failKey, time.Duration(loginCfg.LockDurationMin)*time.Minute).Err()
|
||||
}
|
||||
|
||||
status.FailCount = int(count)
|
||||
status.NeedCaptcha = loginCfg.ShowAfterFailures > 0 && int(count) >= loginCfg.ShowAfterFailures
|
||||
|
||||
if loginCfg.LockAfterFailures > 0 && int(count) >= loginCfg.LockAfterFailures {
|
||||
lockTTL := time.Duration(loginCfg.LockDurationMin) * time.Minute
|
||||
if err := m.rdb.Set(ctx, lockKey, "1", lockTTL).Err(); err != nil {
|
||||
return status, err
|
||||
}
|
||||
_ = m.rdb.Del(ctx, failKey).Err()
|
||||
status.Locked = true
|
||||
status.NeedCaptcha = true
|
||||
status.RemainingSec = int(lockTTL.Seconds())
|
||||
return status, ErrLoginLocked
|
||||
}
|
||||
|
||||
return status, nil
|
||||
}
|
||||
|
||||
func (m *manager) ClearLoginFailures(ctx context.Context, key string) error {
|
||||
if !m.Enabled() {
|
||||
return nil
|
||||
}
|
||||
return m.rdb.Del(ctx, loginFailKey(key), loginLockKey(key)).Err()
|
||||
}
|
||||
|
||||
func (m *manager) requireEnabled() error {
|
||||
if !m.Enabled() {
|
||||
return ErrDisabled
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *manager) sceneMode(scene Scene) string {
|
||||
if m.cfg == nil || m.cfg.Scenes == nil {
|
||||
return defaultSceneMode(scene)
|
||||
}
|
||||
if sc, ok := m.cfg.Scenes[string(scene)]; ok && sc.Mode != "" {
|
||||
return sc.Mode
|
||||
}
|
||||
return defaultSceneMode(scene)
|
||||
}
|
||||
|
||||
func defaultSceneMode(scene Scene) string {
|
||||
switch scene {
|
||||
case SceneRegister, SceneChangePhone, SceneSendSMS:
|
||||
return "always"
|
||||
default:
|
||||
return "never"
|
||||
}
|
||||
}
|
||||
|
||||
func (m *manager) loginConfig() config.LoginCaptchaConfig {
|
||||
if m.cfg == nil || m.cfg.Login == nil {
|
||||
return config.LoginCaptchaConfig{
|
||||
ShowAfterFailures: 3,
|
||||
LockAfterFailures: 10,
|
||||
LockDurationMin: 30,
|
||||
}
|
||||
}
|
||||
return *m.cfg.Login
|
||||
}
|
||||
|
||||
func (m *manager) getLoginFailCount(ctx context.Context, key string) (int, error) {
|
||||
n, err := m.rdb.Get(ctx, loginFailKey(key)).Int()
|
||||
if err == redis.Nil {
|
||||
return 0, nil
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func loginFailKey(key string) string {
|
||||
return "captcha:login:fail:" + key
|
||||
}
|
||||
|
||||
func loginLockKey(key string) string {
|
||||
return "captcha:login:lock:" + key
|
||||
}
|
||||
174
captcha/captcha_test.go
Normal file
174
captcha/captcha_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
17
captcha/errors.go
Normal file
17
captcha/errors.go
Normal file
@@ -0,0 +1,17 @@
|
||||
package captcha
|
||||
|
||||
import "errors"
|
||||
|
||||
var (
|
||||
// ErrDisabled 图形验证码子系统未启用。
|
||||
ErrDisabled = errors.New("captcha disabled")
|
||||
|
||||
// ErrInvalid 验证码错误或已失效。
|
||||
ErrInvalid = errors.New("captcha invalid")
|
||||
|
||||
// ErrRequired 当前场景要求验证码但未提供或校验失败。
|
||||
ErrRequired = errors.New("captcha required")
|
||||
|
||||
// ErrLoginLocked 登录失败次数过多，账号或 IP 已被临时锁定。
|
||||
ErrLoginLocked = errors.New("login locked")
|
||||
)
|
||||
42
captcha/redis_store.go
Normal file
42
captcha/redis_store.go
Normal file
@@ -0,0 +1,42 @@
|
||||
package captcha
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
const codeKeyPrefix = "captcha:code:"
|
||||
|
||||
// redisStore 实现 base64Captcha.Store，将验证码答案存入 Redis。
|
||||
type redisStore struct {
|
||||
client *redis.Client
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
func newRedisStore(client *redis.Client, ttl time.Duration) *redisStore {
|
||||
return &redisStore{client: client, ttl: ttl}
|
||||
}
|
||||
|
||||
func (s *redisStore) Set(id string, value string) error {
|
||||
return s.client.Set(context.Background(), codeKeyPrefix+id, value, s.ttl).Err()
|
||||
}
|
||||
|
||||
func (s *redisStore) Get(id string, clear bool) string {
|
||||
ctx := context.Background()
|
||||
key := codeKeyPrefix + id
|
||||
val, err := s.client.Get(ctx, key).Result()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
if clear {
|
||||
_ = s.client.Del(ctx, key).Err()
|
||||
}
|
||||
return val
|
||||
}
|
||||
|
||||
func (s *redisStore) Verify(id, answer string, clear bool) bool {
|
||||
val := s.Get(id, clear)
|
||||
return val != "" && val == answer
|
||||
}
|
||||
Reference in New Issue
Block a user