308 lines
7.3 KiB
Go
308 lines
7.3 KiB
Go
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
|
||
}
|