enkly-ai/services/iam/internal/infrastructure/postgres/verification_repo.go

package postgres

import (
	"context"
	"crypto/subtle"
	"database/sql"
	"errors"
	"time"

	"enkly/pkg/pgkit"
	"enkly/services/iam/internal/domain"
)

// maxVerificationAttempts caps how many wrong codes a single
// verification row accepts before it's dead even to the correct code —
// see domain.VerificationRepository's doc comment.
const maxVerificationAttempts = 8

// VerificationRepo implements domain.VerificationRepository.
type VerificationRepo struct {
	DB *pgkit.DB
}

func NewVerificationRepo(db *pgkit.DB) *VerificationRepo { return &VerificationRepo{DB: db} }

func (r *VerificationRepo) Create(ctx context.Context, userID int64, codeHash string, expiresAt time.Time) error {
	tx, err := r.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx,
		`DELETE FROM email_verification_codes WHERE user_id = $1 AND used_at IS NULL`, userID); err != nil {
		return err
	}
	if _, err := tx.ExecContext(ctx,
		`INSERT INTO email_verification_codes (user_id, code_hash, expires_at) VALUES ($1, $2, $3)`,
		userID, codeHash, expiresAt); err != nil {
		return err
	}
	return tx.Commit()
}

func (r *VerificationRepo) VerifyAndConsume(ctx context.Context, userID int64, codeHash string) error {
	tx, err := r.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	var verified bool
	if err := tx.QueryRowContext(ctx, `SELECT email_verified FROM users WHERE id = $1 FOR UPDATE`, userID).Scan(&verified); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return domain.ErrNotFound
		}
		return err
	}
	if verified {
		return nil // idempotent for an already-verified account
	}

	var (
		rowID      int64
		storedHash string
		attempts   int
		expiresAt  time.Time
	)
	err = tx.QueryRowContext(ctx, `
		SELECT id, code_hash, attempts, expires_at
		FROM email_verification_codes
		WHERE user_id = $1 AND used_at IS NULL
		ORDER BY created_at DESC LIMIT 1
		FOR UPDATE`, userID).Scan(&rowID, &storedHash, &attempts, &expiresAt)
	if errors.Is(err, sql.ErrNoRows) {
		return domain.ErrInvalidToken
	}
	if err != nil {
		return err
	}

	if attempts >= maxVerificationAttempts || time.Now().After(expiresAt) {
		return domain.ErrInvalidToken
	}

	if subtle.ConstantTimeCompare([]byte(storedHash), []byte(codeHash)) != 1 {
		if _, err := tx.ExecContext(ctx,
			`UPDATE email_verification_codes SET attempts = attempts + 1 WHERE id = $1`, rowID); err != nil {
			return err
		}
		if err := tx.Commit(); err != nil {
			return err
		}
		return domain.ErrInvalidToken
	}

	if _, err := tx.ExecContext(ctx,
		`UPDATE email_verification_codes SET used_at = now() WHERE id = $1`, rowID); err != nil {
		return err
	}
	if _, err := tx.ExecContext(ctx,
		`UPDATE users SET email_verified = true WHERE id = $1`, userID); err != nil {
		return err
	}
	return tx.Commit()
}

// ---------------------------------------------------------------------
// Password reset tokens
// ---------------------------------------------------------------------

// PasswordResetRepo implements domain.PasswordResetRepository.
type PasswordResetRepo struct {
	DB *pgkit.DB
}

func NewPasswordResetRepo(db *pgkit.DB) *PasswordResetRepo { return &PasswordResetRepo{DB: db} }

func (r *PasswordResetRepo) Create(ctx context.Context, userID int64, tokenHash string, expiresAt time.Time) error {
	tx, err := r.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx,
		`DELETE FROM password_reset_tokens WHERE user_id = $1 AND used_at IS NULL`, userID); err != nil {
		return err
	}
	if _, err := tx.ExecContext(ctx,
		`INSERT INTO password_reset_tokens (user_id, token_hash, expires_at) VALUES ($1, $2, $3)`,
		userID, tokenHash, expiresAt); err != nil {
		return err
	}
	return tx.Commit()
}

func (r *PasswordResetRepo) Valid(ctx context.Context, tokenHash string) (bool, error) {
	var ok bool
	err := r.DB.QueryRowContext(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM password_reset_tokens
			WHERE token_hash = $1 AND used_at IS NULL AND expires_at > now()
		)`, tokenHash).Scan(&ok)
	return ok, err
}

func (r *PasswordResetRepo) ConsumeAndResetPassword(ctx context.Context, tokenHash, newPasswordHash string) (int64, error) {
	tx, err := r.DB.BeginTx(ctx, nil)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	var userID int64
	err = tx.QueryRowContext(ctx, `
		UPDATE password_reset_tokens SET used_at = now()
		WHERE token_hash = $1 AND used_at IS NULL AND expires_at > now()
		RETURNING user_id`, tokenHash).Scan(&userID)
	if errors.Is(err, sql.ErrNoRows) {
		return 0, domain.ErrInvalidToken
	}
	if err != nil {
		return 0, err
	}

	if _, err := tx.ExecContext(ctx,
		`UPDATE users SET password_hash = $2 WHERE id = $1`, userID, newPasswordHash); err != nil {
		return 0, err
	}
	if _, err := tx.ExecContext(ctx, `DELETE FROM sessions WHERE user_id = $1`, userID); err != nil {
		return 0, err
	}
	if err := tx.Commit(); err != nil {
		return 0, err
	}
	return userID, nil
}