rond-with-you/internal/auth/totp.go

package auth

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"
)

// TOTP 实现（RFC 6238，30 秒步长、6 位数字、HMAC-SHA1），
// 与 Google Authenticator / 1Password / Microsoft Authenticator 兼容。

const (
	totpStep   = 30 * time.Second
	totpStepS  = 30 // 步长（秒），计数器 = unix秒 / 30
	totpWindow = 1  // 允许前后偏差的步数
)

// counterOf 把时刻换算成 TOTP 计数器。
func counterOf(unixSec int64) uint64 {
	if unixSec < 0 {
		return 0
	}
	return uint64(unixSec / totpStepS)
}

var ErrTOTPNotEnabled = errors.New("未开启两步验证")

// GenerateTOTPSecret 生成 20 字节随机秘钥的 base32 编码（无填充）。
func GenerateTOTPSecret() (string, error) {
	raw := make([]byte, 20)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(raw), nil
}

func totpCode(secret string, counter uint64) (string, error) {
	key, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(
		strings.ToUpper(strings.ReplaceAll(secret, " ", "")))
	if err != nil {
		return "", fmt.Errorf("秘钥格式无效: %w", err)
	}
	var msg [8]byte
	binary.BigEndian.PutUint64(msg[:], counter)
	m := hmac.New(sha1.New, key)
	m.Write(msg[:])
	sum := m.Sum(nil)
	off := sum[len(sum)-1] & 0x0f
	code := (binary.BigEndian.Uint32(sum[off:off+4]) & 0x7fffffff) % 1_000_000
	return fmt.Sprintf("%06d", code), nil
}

// TOTPCode 返回指定时刻的验证码（用于展示与测试）。
func TOTPCode(secret string, t time.Time) (string, error) {
	return totpCode(secret, counterOf(t.Unix()))
}

// VerifyTOTP 校验验证码，允许前后各一个步长的时钟偏差。
func VerifyTOTP(secret, code string) bool {
	now := counterOf(time.Now().Unix())
	for _, c := range []uint64{now - totpWindow, now, now + totpWindow} {
		want, err := totpCode(secret, c)
		if err != nil {
			return false
		}
		if subtle.ConstantTimeCompare([]byte(want), []byte(strings.TrimSpace(code))) == 1 {
			return true
		}
	}
	return false
}

// OTPAuthURL 生成验证器 App 扫码用的 URI。
func OTPAuthURL(username, secret string) string {
	label := urlEscape("rond-with-you:" + username)
	return "otpauth://totp/" + label + "?secret=" + secret + "&issuer=" + urlEscape("rond-with-you") +
		"&algorithm=SHA1&digits=6&period=30"
}

func urlEscape(s string) string {
	var b strings.Builder
	for _, r := range s {
		switch {
		case r >= 'A' && r <= 'Z', r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '-', r == '_', r == '.', r == '~', r == ':':
			b.WriteRune(r)
		default:
			fmt.Fprintf(&b, "%%%02X", r)
		}
	}
	return b.String()
}

// ---------- 验证码重放 ----------

// replayTTL 要盖过 VerifyTOTP 的整个接受窗口（前后各一步 = 90 秒），
// 否则同一步长的码会在条目过期后又被收下。
const replayTTL = 3 * totpStep

// CodeReplay 记住最近用过的验证码，拦住「同一个码在窗口内被用第二次」。
// VerifyTOTP 只判断「这个码此刻对不对」，不含一次性语义；被偷看到或被日志
// 截获的码，在它过期前都能再换一个会话。
type CodeReplay struct {
	mu   sync.Mutex
	used map[string]time.Time
	now  func() time.Time
}

func NewCodeReplay() *CodeReplay {
	return &CodeReplay{used: map[string]time.Time{}, now: time.Now}
}

// Consume 在验证码**已经校验通过**之后登记一次使用；已被用过则返回 false。
// 键里带用户号：不同账号的同一个数字码互不影响。
func (r *CodeReplay) Consume(userID int64, code string) bool {
	now := r.now()
	key := fmt.Sprintf("%d|%s", userID, code)
	r.mu.Lock()
	defer r.mu.Unlock()
	for k, t := range r.used {
		if now.Sub(t) > replayTTL {
			delete(r.used, k)
		}
	}
	if _, ok := r.used[key]; ok {
		return false
	}
	r.used[key] = now
	return true
}

// ---------- 用户秘钥存取 ----------

// TOTPState 返回用户的两步验证状态。
func (s *Store) TOTPState(ctx context.Context, userID int64) (secret string, enabled bool, err error) {
	err = s.DB.QueryRowContext(ctx, `SELECT COALESCE(totp_secret,''), totp_enabled
		FROM users WHERE id=$1`, userID).Scan(&secret, &enabled)
	return
}

// SetTOTPSecret 保存（或覆盖）秘钥；enabled 保持不变，供「先展示二维码再确认」流程。
func (s *Store) SetTOTPSecret(ctx context.Context, userID int64, secret string) error {
	_, err := s.DB.ExecContext(ctx, `UPDATE users SET totp_secret=$1 WHERE id=$2`, secret, userID)
	return err
}

// EnableTOTP 确认启用两步验证。
func (s *Store) EnableTOTP(ctx context.Context, userID int64) error {
	_, err := s.DB.ExecContext(ctx, `UPDATE users SET totp_enabled=TRUE WHERE id=$1 AND totp_secret<>''`, userID)
	return err
}

// DisableTOTP 关闭两步验证并清除秘钥。
func (s *Store) DisableTOTP(ctx context.Context, userID int64) error {
	_, err := s.DB.ExecContext(ctx, `UPDATE users SET totp_enabled=FALSE, totp_secret='' WHERE id=$1`, userID)
	return err
}

// ---------- 两步登录的中间凭据 ----------

// SignPending2FA 在口令校验通过但未过两步验证时签发短时中间令牌，
// 绑定口令摘要（改密即失效），有效期 5 分钟。
func (s *Store) SignPending2FA(userID int64, hash string) string {
	exp := time.Now().Add(5 * time.Minute).Unix()
	payload := fmt.Sprintf("2fa|%d|%d", userID, exp)
	return payload + "|" + s.mac("p"+payload, hash)
}

// ParsePending2FA 还原中间令牌。
func (s *Store) ParsePending2FA(token, hash string) (int64, bool) {
	parts := strings.Split(token, "|")
	if len(parts) != 4 || parts[0] != "2fa" {
		return 0, false
	}
	var uid, exp int64
	if _, err := fmt.Sscanf(parts[1], "%d", &uid); err != nil {
		return 0, false
	}
	if _, err := fmt.Sscanf(parts[2], "%d", &exp); err != nil {
		return 0, false
	}
	if time.Now().Unix() > exp {
		return 0, false
	}
	payload := "2fa|" + parts[1] + "|" + parts[2]
	if !hmac.Equal([]byte(parts[3]), []byte(s.mac("p"+payload, hash))) {
		return 0, false
	}
	return uid, true
}