GovAI/server/internal/service/twofa/store.go

package twofa

import (
	"context"

	"github.com/enterprise-ai-platform/server/pkg/auth"
	"github.com/jackc/pgx/v5/pgxpool"
)

// pgxStore 基于 pgx 连接池的 Store 实现。
type pgxStore struct {
	pool *pgxpool.Pool
}

func NewPgxStore(pool *pgxpool.Pool) Store {
	return &pgxStore{pool: pool}
}

func (s *pgxStore) GetStatus(ctx context.Context, userID string) (bool, int, error) {
	var enabled bool
	var remaining int
	err := s.pool.QueryRow(ctx,
		`SELECT u.totp_enabled,
		        (SELECT COUNT(*) FROM user_backup_codes b WHERE b.user_id = u.id AND b.used_at IS NULL)
		 FROM users u WHERE u.id = $1`, userID).Scan(&enabled, &remaining)
	return enabled, remaining, err
}

func (s *pgxStore) GetSecret(ctx context.Context, userID string) (string, error) {
	var secret *string
	if err := s.pool.QueryRow(ctx,
		`SELECT totp_secret FROM users WHERE id = $1`, userID).Scan(&secret); err != nil {
		return "", err
	}
	if secret == nil {
		return "", nil
	}
	return *secret, nil
}

func (s *pgxStore) GetSecretAndEnabled(ctx context.Context, userID string) (string, bool, error) {
	var secret *string
	var enabled bool
	if err := s.pool.QueryRow(ctx,
		`SELECT totp_secret, totp_enabled FROM users WHERE id = $1`, userID).Scan(&secret, &enabled); err != nil {
		return "", false, err
	}
	if secret == nil {
		return "", enabled, nil
	}
	return *secret, enabled, nil
}

func (s *pgxStore) SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)

	if _, err = tx.Exec(ctx,
		`UPDATE users SET totp_secret = $2, totp_enabled = false WHERE id = $1`, userID, secret); err != nil {
		return err
	}
	if _, err = tx.Exec(ctx, `DELETE FROM user_backup_codes WHERE user_id = $1`, userID); err != nil {
		return err
	}
	for _, h := range codeHashes {
		if _, err = tx.Exec(ctx,
			`INSERT INTO user_backup_codes (user_id, code_hash) VALUES ($1, $2)`, userID, h); err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}

func (s *pgxStore) Enable(ctx context.Context, userID string) error {
	_, err := s.pool.Exec(ctx, `UPDATE users SET totp_enabled = true WHERE id = $1`, userID)
	return err
}

func (s *pgxStore) Disable(ctx context.Context, userID string) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)
	if _, err = tx.Exec(ctx,
		`UPDATE users SET totp_enabled = false, totp_secret = NULL WHERE id = $1`, userID); err != nil {
		return err
	}
	if _, err = tx.Exec(ctx, `DELETE FROM user_backup_codes WHERE user_id = $1`, userID); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *pgxStore) ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT id, code_hash FROM user_backup_codes WHERE user_id = $1 AND used_at IS NULL`, userID)
	if err != nil {
		return false, err
	}
	type bc struct{ id, hash string }
	var list []bc
	for rows.Next() {
		var x bc
		if rows.Scan(&x.id, &x.hash) == nil {
			list = append(list, x)
		}
	}
	rows.Close() // 先释放连接再执行更新

	for _, x := range list {
		if auth.CheckBackupCode(code, x.hash) {
			_, _ = s.pool.Exec(ctx, `UPDATE user_backup_codes SET used_at = NOW() WHERE id = $1`, x.id)
			return true, nil
		}
	}
	return false, nil
}