usernode/internal/auth/otp.go

// Package auth 提供认证相关能力：OTP（双通道）、会话、bcrypt、图形验证码。
//
// OTP 双通道对齐：邮件发送与 CLI 获取共用同一 OTPStore（同一验证码、同一
// 有效期、同一冷却与失败限速），邮件失败不阻断 CLI 通道。邮件通道经
// Send 生成验证码，CLI 通道优先经 Current 复用同一验证码，无有效码时才
// 触发 Send（仍受同一冷却约束）。验证码落 DB（otp_codes 表），单实例部署
// 即可保证跨进程（HTTP 服务与 CLI 子命令）共享；多实例需改 DB 行锁/Redis
// （PLAN §6）。
package auth

import (
	"context"
	"errors"
	"sync"
	"time"

	"gorm.io/gorm"

	"ws_usernode/internal/model"
	"ws_usernode/internal/pkg"
)

// OTP 错误。
var (
	ErrCooldown     = errors.New("auth: otp 发送冷却中")
	ErrInvalidCode  = errors.New("auth: 验证码错误")
	ErrTooManyFails = errors.New("auth: 失败次数过多，请稍后再试")
)

const (
	otpCodeLen = 6
)

// OTP 失败限速默认参数（单账号连续失败阈值与计数窗口）。
const (
	DefaultMaxFailures = 5
	DefaultFailureWin  = 10 * time.Minute
)

// OTPStore 为 OTP 验证码存储。DB 实现（DBOTPStore）为生产默认，
// MemoryOTPStore 供测试与单进程内嵌场景使用。
type OTPStore interface {
	// Send 为 username 生成新验证码并覆盖旧码（邮件通道）。冷却期内返回 ErrCooldown。
	Send(ctx context.Context, username string, ttl, cooldown time.Duration) (string, error)
	// Current 返回当前有效（未过期、未消费）验证码，供 CLI 通道复用同一验证码。
	// 无有效验证码返回 ErrInvalidCode。
	Current(ctx context.Context, username string) (string, error)
	// Verify 校验验证码并一次性消费。失败累计计数（达到阈值返回 ErrTooManyFails）。
	Verify(ctx context.Context, username, code string) (bool, error)
	// Failures 返回 username 当前失败计数。
	Failures(ctx context.Context, username string) (int, error)
}

// DBOTPStore 基于 model.OTPCode 的存储实现，每用户一行（username 唯一）。
type DBOTPStore struct {
	db          *gorm.DB
	maxFailures int
	failureWin  time.Duration
}

// NewDBOTPStore 创建 DB OTP 存储。
func NewDBOTPStore(db *gorm.DB, maxFailures int, failureWin time.Duration) *DBOTPStore {
	return &DBOTPStore{db: db, maxFailures: maxFailures, failureWin: failureWin}
}

func (s *DBOTPStore) Send(ctx context.Context, username string, ttl, cooldown time.Duration) (string, error) {
	now := time.Now()
	var row model.OTPCode
	err := s.db.WithContext(ctx).Where("username = ?", username).First(&row).Error
	// 注意：First 的 err 不能复用给后续语句，避免被覆盖导致走错分支
	exists := err == nil
	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
		return "", err
	}
	if exists && now.Before(row.CooldownUntil) {
		return "", ErrCooldown
	}
	code, randErr := pkg.RandomDigits(otpCodeLen)
	if randErr != nil {
		return "", randErr
	}
	updates := map[string]any{
		"code":           code,
		"expires_at":     now.Add(ttl),
		"cooldown_until": now.Add(cooldown),
		"failures":       0,
		"failed_at":      nil,
		"consumed_at":    nil,
	}
	if exists {
		err = s.db.WithContext(ctx).Model(&row).Updates(updates).Error
	} else {
		err = s.db.WithContext(ctx).Create(&model.OTPCode{
			Username:      username,
			Code:          code,
			ExpiresAt:     now.Add(ttl),
			CooldownUntil: now.Add(cooldown),
		}).Error
	}
	if err != nil {
		return "", err
	}
	return code, nil
}

func (s *DBOTPStore) Current(ctx context.Context, username string) (string, error) {
	var row model.OTPCode
	if err := s.db.WithContext(ctx).Where("username = ?", username).First(&row).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return "", ErrInvalidCode
		}
		return "", err
	}
	if row.ConsumedAt != nil || time.Now().After(row.ExpiresAt) {
		return "", ErrInvalidCode
	}
	return row.Code, nil
}

func (s *DBOTPStore) Verify(ctx context.Context, username, code string) (bool, error) {
	now := time.Now()
	var row model.OTPCode
	if err := s.db.WithContext(ctx).Where("username = ?", username).First(&row).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return false, ErrInvalidCode
		}
		return false, err
	}
	if row.ConsumedAt != nil || now.After(row.ExpiresAt) {
		return false, ErrInvalidCode
	}
	// 失败计数窗口：超出窗口则重置计数
	if row.FailedAt != nil && now.Sub(*row.FailedAt) > s.failureWin {
		row.Failures = 0
		row.FailedAt = nil
	}
	if row.Failures >= s.maxFailures {
		return false, ErrTooManyFails
	}
	if row.Code != code {
		row.Failures++
		f := now
		row.FailedAt = &f
		_ = s.db.WithContext(ctx).Model(&row).Updates(map[string]any{"failures": row.Failures, "failed_at": row.FailedAt}).Error
		if row.Failures >= s.maxFailures {
			return false, ErrTooManyFails
		}
		return false, ErrInvalidCode
	}
	consumed := now
	if err := s.db.WithContext(ctx).Model(&row).Update("consumed_at", &consumed).Error; err != nil {
		return false, err
	}
	return true, nil
}

func (s *DBOTPStore) Failures(ctx context.Context, username string) (int, error) {
	var row model.OTPCode
	if err := s.db.WithContext(ctx).Where("username = ?", username).First(&row).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return 0, nil
		}
		return 0, err
	}
	return row.Failures, nil
}

type otpEntry struct {
	code       string
	expiresAt  time.Time
	cooldownAt time.Time
	failures   int
}

// MemoryOTPStore 为单进程内存实现（测试/内嵌场景）。
type MemoryOTPStore struct {
	mu      sync.Mutex
	entries map[string]*otpEntry
}

// NewMemoryOTPStore 创建内存 OTP 存储。
func NewMemoryOTPStore() *MemoryOTPStore {
	return &MemoryOTPStore{entries: make(map[string]*otpEntry)}
}

func (s *MemoryOTPStore) Send(ctx context.Context, username string, ttl, cooldown time.Duration) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	_ = ctx

	now := time.Now()
	if e, ok := s.entries[username]; ok && now.Before(e.cooldownAt) {
		return "", ErrCooldown
	}
	code, err := pkg.RandomDigits(otpCodeLen)
	if err != nil {
		return "", err
	}
	s.entries[username] = &otpEntry{
		code:       code,
		expiresAt:  now.Add(ttl),
		cooldownAt: now.Add(cooldown),
		failures:   0,
	}
	return code, nil
}

func (s *MemoryOTPStore) Current(ctx context.Context, username string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	_ = ctx
	e, ok := s.entries[username]
	if !ok || time.Now().After(e.expiresAt) {
		return "", ErrInvalidCode
	}
	return e.code, nil
}

func (s *MemoryOTPStore) Verify(ctx context.Context, username, code string) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	_ = ctx

	e, ok := s.entries[username]
	if !ok {
		return false, ErrInvalidCode
	}
	now := time.Now()
	if now.After(e.expiresAt) {
		delete(s.entries, username)
		return false, ErrInvalidCode
	}
	if e.failures >= DefaultMaxFailures {
		return false, ErrTooManyFails
	}
	if e.code != code {
		e.failures++
		if e.failures >= DefaultMaxFailures {
			return false, ErrTooManyFails
		}
		return false, ErrInvalidCode
	}
	delete(s.entries, username) // 一次性
	return true, nil
}

func (s *MemoryOTPStore) Failures(ctx context.Context, username string) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	_ = ctx
	if e, ok := s.entries[username]; ok {
		return e.failures, nil
	}
	return 0, nil
}