oci-portal/internal/service/totp.go

package service

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/pquerna/otp/totp"
	"golang.org/x/crypto/bcrypt"

	"oci-portal/internal/model"
)

// TOTP(RFC 6238)两步验证相关错误;api 层据此映射状态码。
var (
	// ErrTotpRequired 表示账号已启用两步验证但登录未携带验证码(密码已通过)。
	ErrTotpRequired = errors.New("需要两步验证码")
	// ErrTotpInvalid 表示验证码校验失败(设置流程中)。
	ErrTotpInvalid = errors.New("两步验证码错误")
	// ErrTotpNotSetup 表示激活前未执行 setup 或暂存已过期。
	ErrTotpNotSetup = errors.New("请先重新发起两步验证设置")
	// ErrTotpAlreadyOn 表示重复启用。
	ErrTotpAlreadyOn = errors.New("两步验证已启用,如需重置请先停用")
	// ErrTotpConfirm 表示停用时未提供密码或验证码确认。
	ErrTotpConfirm = errors.New("停用需要当前验证码或登录密码确认")
)

// totpPendingTTL 是 setup 暂存密钥的有效期;过期须重新 setup。
const totpPendingTTL = 10 * time.Minute

// totpIssuer 展示在验证器 App 中的签发方名称。
const totpIssuer = "oci-portal"

// pendingTotp 是 setup 后待激活的临时密钥(内存,重启丢弃)。
type pendingTotp struct {
	secret  string
	expires time.Time
}

// TotpStatus 返回用户是否已启用两步验证。
func (s *AuthService) TotpStatus(ctx context.Context, username string) (bool, error) {
	user, err := s.findUser(ctx, username)
	if err != nil {
		return false, err
	}
	return user.TotpSecretEnc != "", nil
}

// findUser 按用户名加载账号。
func (s *AuthService) findUser(ctx context.Context, username string) (*model.User, error) {
	var user model.User
	if err := s.db.WithContext(ctx).Where("username = ?", username).First(&user).Error; err != nil {
		return nil, fmt.Errorf("find user: %w", err)
	}
	return &user, nil
}

// SetupTotp 生成新 TOTP 密钥暂存(10 分钟内待激活),返回手动输入码与 otpauth URI;
// 重复调用覆盖旧暂存,已启用时拒绝(先停用再重置)。
func (s *AuthService) SetupTotp(ctx context.Context, username string) (secret, uri string, err error) {
	if s.cipher == nil {
		return "", "", errors.New("totp: cipher not configured")
	}
	enabled, err := s.TotpStatus(ctx, username)
	if err != nil {
		return "", "", err
	}
	if enabled {
		return "", "", ErrTotpAlreadyOn
	}
	key, err := totp.Generate(totp.GenerateOpts{Issuer: totpIssuer, AccountName: username})
	if err != nil {
		return "", "", fmt.Errorf("generate totp key: %w", err)
	}
	s.totpMu.Lock()
	s.totpPending[username] = pendingTotp{secret: key.Secret(), expires: time.Now().Add(totpPendingTTL)}
	s.totpMu.Unlock()
	return key.Secret(), key.URL(), nil
}

// ActivateTotp 校验暂存密钥的验证码,通过后加密落库启用。
func (s *AuthService) ActivateTotp(ctx context.Context, username, code string) error {
	s.totpMu.Lock()
	pending, ok := s.totpPending[username]
	s.totpMu.Unlock()
	if !ok || time.Now().After(pending.expires) {
		return ErrTotpNotSetup
	}
	if !totp.Validate(code, pending.secret) {
		return ErrTotpInvalid
	}
	enc, err := s.cipher.EncryptString(pending.secret)
	if err != nil {
		return fmt.Errorf("encrypt totp secret: %w", err)
	}
	err = s.db.WithContext(ctx).Model(&model.User{}).
		Where("username = ?", username).Update("totp_secret_enc", enc).Error
	if err != nil {
		return fmt.Errorf("save totp secret: %w", err)
	}
	s.totpMu.Lock()
	delete(s.totpPending, username)
	s.totpMu.Unlock()
	return nil
}

// DisableTotp 停用两步验证;需当前验证码或登录密码任一确认。
func (s *AuthService) DisableTotp(ctx context.Context, username, password, code string) error {
	user, err := s.findUser(ctx, username)
	if err != nil {
		return err
	}
	if user.TotpSecretEnc == "" {
		return nil
	}
	if !s.confirmDisable(user, password, code) {
		return ErrTotpConfirm
	}
	err = s.db.WithContext(ctx).Model(&model.User{}).
		Where("username = ?", username).Update("totp_secret_enc", "").Error
	if err != nil {
		return fmt.Errorf("clear totp secret: %w", err)
	}
	return nil
}

// confirmDisable 校验停用凭证:验证码或密码任一通过即可。
func (s *AuthService) confirmDisable(user *model.User, password, code string) bool {
	if code != "" && s.verifyTotp(user, code) {
		return true
	}
	if password != "" && bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(password)) == nil {
		return true
	}
	return false
}

// verifyTotp 解密密钥并校验验证码(±1 时间步容差)。
func (s *AuthService) verifyTotp(user *model.User, code string) bool {
	if s.cipher == nil || user.TotpSecretEnc == "" {
		return false
	}
	secret, err := s.cipher.DecryptString(user.TotpSecretEnc)
	if err != nil {
		return false
	}
	return totp.Validate(code, secret)
}