nexarch/internal/totp/store.go

package totp

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

var (
	ErrInvalidCode      = errors.New("totp: code ungueltig")
	ErrNotConfirmed     = errors.New("totp: einrichtung noch nicht bestaetigt")
	ErrRecoveryCodeUsed = errors.New("totp: wiederherstellungscode bereits verwendet oder ungueltig")
)

const RecoveryCodeCount = 10

type Store struct {
	pool *pgxpool.Pool
}

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

// BeginSetup erzeugt ein neues Secret fuer einen Benutzer und speichert es
// als NOCH NICHT bestaetigt — der Code muss vor Aktivierung einmal
// verifiziert werden (Akzeptanzkriterium 1).
func (s *Store) BeginSetup(ctx context.Context, userID, issuer, accountName string) (secret, provisioningURI string, err error) {
	secret, err = GenerateSecret()
	if err != nil {
		return "", "", err
	}

	_, err = s.pool.Exec(ctx, `
		INSERT INTO totp_credentials (user_id, secret, confirmed)
		VALUES ($1, $2, false)
		ON CONFLICT (user_id) DO UPDATE SET secret = $2, confirmed = false
	`, userID, secret)
	if err != nil {
		return "", "", fmt.Errorf("totp-einrichtung speichern: %w", err)
	}

	return secret, ProvisioningURI(issuer, accountName, secret), nil
}

// ConfirmSetup verifiziert den ersten eingegebenen Code und aktiviert 2FA
// erst dann (Akzeptanzkriterium 1). Bei Erfolg werden einmalig
// Wiederherstellungscodes erzeugt und im Klartext zurueckgegeben — gespeichert
// wird nur deren Hash.
func (s *Store) ConfirmSetup(ctx context.Context, userID, code string) (recoveryCodes []string, err error) {
	var secret string
	if err := s.pool.QueryRow(ctx, `
		SELECT secret FROM totp_credentials WHERE user_id = $1
	`, userID).Scan(&secret); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotConfirmed
		}
		return nil, fmt.Errorf("secret lesen: %w", err)
	}

	ok, err := Validate(secret, code, time.Now())
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, ErrInvalidCode
	}

	if _, err := s.pool.Exec(ctx, `UPDATE totp_credentials SET confirmed = true WHERE user_id = $1`, userID); err != nil {
		return nil, fmt.Errorf("bestaetigung speichern: %w", err)
	}

	recoveryCodes, err = s.generateRecoveryCodes(ctx, userID)
	if err != nil {
		return nil, err
	}
	return recoveryCodes, nil
}

func (s *Store) generateRecoveryCodes(ctx context.Context, userID string) ([]string, error) {
	codes := make([]string, RecoveryCodeCount)
	for i := range codes {
		buf := make([]byte, 8)
		if _, err := rand.Read(buf); err != nil {
			return nil, fmt.Errorf("wiederherstellungscode erzeugen: %w", err)
		}
		code := hex.EncodeToString(buf)
		codes[i] = code

		hash := hashRecoveryCode(code)
		if _, err := s.pool.Exec(ctx, `
			INSERT INTO totp_recovery_codes (user_id, code_hash) VALUES ($1, $2)
		`, userID, hash); err != nil {
			return nil, fmt.Errorf("wiederherstellungscode speichern: %w", err)
		}
	}
	return codes, nil
}

// IsEnabled prueft, ob 2FA fuer einen Benutzer aktiv (bestaetigt) ist.
func (s *Store) IsEnabled(ctx context.Context, userID string) (bool, error) {
	var confirmed bool
	err := s.pool.QueryRow(ctx, `SELECT confirmed FROM totp_credentials WHERE user_id = $1`, userID).Scan(&confirmed)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return false, nil
		}
		return false, fmt.Errorf("status lesen: %w", err)
	}
	return confirmed, nil
}

// VerifyLoginCode akzeptiert entweder einen gueltigen TOTP-Code oder einen
// noch nicht verwendeten Wiederherstellungscode (Akzeptanzkriterium 2 + 3).
func (s *Store) VerifyLoginCode(ctx context.Context, userID, code string) (bool, error) {
	var secret string
	var confirmed bool
	err := s.pool.QueryRow(ctx, `SELECT secret, confirmed FROM totp_credentials WHERE user_id = $1`, userID).Scan(&secret, &confirmed)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) || !confirmed {
			return false, nil
		}
		return false, fmt.Errorf("secret lesen: %w", err)
	}
	if !confirmed {
		return false, nil
	}

	if ok, err := Validate(secret, code, time.Now()); err != nil {
		return false, err
	} else if ok {
		return true, nil
	}

	return s.consumeRecoveryCode(ctx, userID, code)
}

// consumeRecoveryCode markiert einen Wiederherstellungscode atomar als
// verwendet (WHERE used_at IS NULL) — ein bereits verbrauchter Code kann
// nicht doppelt eingeloest werden (Akzeptanzkriterium 3 / Pruefung 2).
func (s *Store) consumeRecoveryCode(ctx context.Context, userID, code string) (bool, error) {
	hash := hashRecoveryCode(code)
	tag, err := s.pool.Exec(ctx, `
		UPDATE totp_recovery_codes SET used_at = now()
		WHERE user_id = $1 AND code_hash = $2 AND used_at IS NULL
	`, userID, hash)
	if err != nil {
		return false, fmt.Errorf("wiederherstellungscode einloesen: %w", err)
	}
	return tag.RowsAffected() == 1, nil
}

func hashRecoveryCode(code string) []byte {
	sum := sha256.Sum256([]byte(code))
	return sum[:]
}