oci-portal/internal/service/totp.go

package service

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

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

	"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, oldToken string, meta SessionMeta, proof TokenProof,
) (string, time.Time, error) {
	s.totpMu.Lock()
	pending, ok := s.totpPending[username]
	s.totpMu.Unlock()
	if !ok || time.Now().After(pending.expires) {
		return "", time.Time{}, ErrTotpNotSetup
	}
	if !totp.Validate(code, pending.secret) {
		return "", time.Time{}, ErrTotpInvalid
	}
	enc, err := s.cipher.EncryptString(pending.secret)
	if err != nil {
		return "", time.Time{}, fmt.Errorf("encrypt totp secret: %w", err)
	}
	token, expires, err := s.changeTotp(ctx, username, enc, oldToken, meta, proof, nil)
	if err == nil {
		s.deleteTotpPending(username, pending.secret)
	}
	return token, expires, err
}

// DisableTotp 停用两步验证;需当前验证码或登录密码任一确认。
func (s *AuthService) DisableTotp(
	ctx context.Context, username, password, code, oldToken string, meta SessionMeta, proof TokenProof,
) (string, time.Time, error) {
	confirm := func(user *model.User) error {
		if user.TotpSecretEnc != "" && !s.confirmDisable(user, password, code) {
			return ErrTotpConfirm
		}
		return nil
	}
	return s.changeTotp(ctx, username, "", oldToken, meta, proof, confirm)
}

func (s *AuthService) changeTotp(
	ctx context.Context, username, secret, oldToken string, meta SessionMeta,
	proof TokenProof, confirm func(*model.User) error,
) (string, time.Time, error) {
	var token string
	var expires time.Time
	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		user, err := lockUserForAuthChange(tx, username)
		if err != nil {
			return err
		}
		if err := s.ensureTokenCurrentTx(tx, user, proof); err != nil {
			return err
		}
		if confirm != nil {
			if err := confirm(user); err != nil {
				return err
			}
		}
		return s.persistTotpTx(tx, user, secret, oldToken, meta, &token, &expires)
	})
	return token, expires, err
}

func (s *AuthService) persistTotpTx(
	tx *gorm.DB, user *model.User, secret, oldToken string, meta SessionMeta,
	token *string, expires *time.Time,
) error {
	if err := tx.Model(user).Update("totp_secret_enc", secret).Error; err != nil {
		return fmt.Errorf("save totp secret: %w", err)
	}
	if err := bumpTokenVersionTx(tx, user.Username); err != nil {
		return err
	}
	user.TokenVersion++
	var err error
	*token, *expires, err = s.renewSessionTx(tx, user, oldToken, meta)
	return err
}

func (s *AuthService) deleteTotpPending(username, secret string) {
	s.totpMu.Lock()
	defer s.totpMu.Unlock()
	if p, ok := s.totpPending[username]; ok && p.secret == secret {
		delete(s.totpPending, username)
	}
}

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