jarl/internal/store/totp.go

package store

import (
	"context"
	"database/sql"
	"time"
)

// SetTOTPPending stores a secret being set up (or clears it with ""). It
// becomes the user's secret once a code confirms it (EnableTOTP).
func (s *Store) SetTOTPPending(ctx context.Context, userID int64, secret string) error {
	res, err := s.db.ExecContext(ctx, `UPDATE users SET totp_pending = ? WHERE id = ?`, secret, userID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// EnableTOTP turns two-factor login on: the secret, the step of the code
// that confirmed it (so that code can't be used again), and the hashes of
// new recovery codes, replacing any old ones.
func (s *Store) EnableTOTP(ctx context.Context, userID int64, secret string, step int64, codeHashes []string) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	now := ts(time.Now())
	res, err := tx.ExecContext(ctx, `UPDATE users SET totp_secret = ?, totp_pending = '', totp_enabled_at = ?, totp_last_step = ?, updated_at = ? WHERE id = ?`,
		secret, now, step, now, userID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	if err := replaceRecoveryCodes(ctx, tx, userID, codeHashes); err != nil {
		return err
	}
	return tx.Commit()
}

// DisableTOTP turns two-factor login off and deletes the recovery codes.
func (s *Store) DisableTOTP(ctx context.Context, userID int64) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	res, err := tx.ExecContext(ctx, `UPDATE users SET totp_secret = '', totp_pending = '', totp_enabled_at = NULL, totp_last_step = 0, updated_at = ? WHERE id = ?`, ts(time.Now()), userID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	if _, err := tx.ExecContext(ctx, `DELETE FROM recovery_codes WHERE user_id = ?`, userID); err != nil {
		return err
	}
	return tx.Commit()
}

// UseTOTPStep records that the code of step was used. It returns false if
// that step (or a later one) was used already: each code works once, even
// with two requests at the same moment.
func (s *Store) UseTOTPStep(ctx context.Context, userID, step int64) (bool, error) {
	res, err := s.db.ExecContext(ctx, `UPDATE users SET totp_last_step = ? WHERE id = ? AND totp_last_step < ? AND totp_secret != ''`, step, userID, step)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n == 1, nil
}

// ReplaceRecoveryCodes gives a user new recovery codes (their hashes),
// making the old ones useless.
func (s *Store) ReplaceRecoveryCodes(ctx context.Context, userID int64, codeHashes []string) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if err := replaceRecoveryCodes(ctx, tx, userID, codeHashes); err != nil {
		return err
	}
	return tx.Commit()
}

func replaceRecoveryCodes(ctx context.Context, tx *sql.Tx, userID int64, hashes []string) error {
	if _, err := tx.ExecContext(ctx, `DELETE FROM recovery_codes WHERE user_id = ?`, userID); err != nil {
		return err
	}
	for _, h := range hashes {
		if _, err := tx.ExecContext(ctx, `INSERT INTO recovery_codes (user_id, code_hash) VALUES (?, ?)`, userID, h); err != nil {
			return err
		}
	}
	return nil
}

// UseRecoveryCode marks an unused recovery code (by its hash) used. It
// returns false if the user has no such unused code.
func (s *Store) UseRecoveryCode(ctx context.Context, userID int64, codeHash string) (bool, error) {
	res, err := s.db.ExecContext(ctx, `UPDATE recovery_codes SET used_at = ? WHERE user_id = ? AND code_hash = ? AND used_at IS NULL`, ts(time.Now()), userID, codeHash)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n == 1, nil
}

// RecoveryCodesLeft counts a user's unused recovery codes.
func (s *Store) RecoveryCodesLeft(ctx context.Context, userID int64) (int, error) {
	var n int
	err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM recovery_codes WHERE user_id = ? AND used_at IS NULL`, userID).Scan(&n)
	return n, err
}

// SetSessionTOTP records that a session proved a TOTP code at at.
func (s *Store) SetSessionTOTP(ctx context.Context, idHash string, at time.Time) error {
	_, err := s.db.ExecContext(ctx, `UPDATE sessions SET totp_at = ? WHERE id_hash = ?`, ts(at), idHash)
	return err
}