authenciates_with_two_factor.rb

# frozen_string_literal: true

# == AuthenticatesWithTwoFactor
#
# Controller concern to handle two-factor authentication
module AuthenticatesWithTwoFactor
  extend ActiveSupport::Concern

    def prompt_for_two_factor(user)
      @user = user

      # Save the user's ID to session so we can ask for a one-time password
      session[:otp_user_id] = user.id
      render 'users/sessions/two_factor'
    end

    def authenticate_with_two_factor
      user = self.resource = find_user
      return unless user && user.otp_required_for_login

      if user_params[:otp_attempt].present? && session[:otp_user_id]
        authenticate_with_two_factor_via_otp(user)
      elsif user && user.valid_password?(user_params[:password])
        prompt_for_two_factor(user)
      end
    end

    def authenticate_with_two_factor_via_otp(user)
      if valid_otp_attempt?(user)
        # Remove any lingering user data from login
        session.delete(:otp_user_id)

        sign_in(user)
      else
        flash[:alert] = 'Invalid two-factor code.'
        prompt_for_two_factor(user)
      end
    end
end

sessions_controller.rb

# frozen_string_literal: true

class Users::SessionsController < Devise::SessionsController
  include AuthenticatesWithTwoFactor

  skip_before_action :check_two_factor_requirement, only: [:destroy]

  # This action comes from DeviseController, but because we call `sign_in`
  # manually inside `authenticate_with_two_factor`, not skipping this action
  # would cause a "You are already signed in." error message to be shown upon
  # successful login.
  skip_before_action :require_no_authentication, only: [:new, :create]
  before_action :configure_sign_in_params, only: [:new, :create]
  prepend_before_action :authenticate_with_two_factor,
    if: -> { action_name == 'create' && two_factor_enabled? }

  protect_from_forgery with: :exception, prepend: true, except: :destroy

  layout 'login'

  protected

    def user_params
      params.require(:user).permit(:email, :password, :otp_attempt)
    end

    def find_user
      if session[:otp_user_id]
        User.find(session[:otp_user_id])
      elsif user_params[:email]
        User.find_by_email(user_params[:email])
      end
    end

    def two_factor_enabled?
      find_user&.two_factor_enabled?
    end

    def valid_otp_attempt?(user)
      user.validate_and_consume_otp!(user_params[:otp_attempt]) ||
        user.invalidate_otp_backup_code!(user_params[:otp_attempt])
    end

    def after_sign_in_path_for(resource)
      if session[:after_sign_in_path].present?
        session[:after_sign_in_path]
      else
        stored_location_for(resource) || root_path
      end
    end

    def after_sign_out_path_for(resource)
      stored_location_for(resource) || root_path
    end

    def configure_sign_in_params
      devise_parameter_sanitizer.permit(:sign_in, keys: [:otp_attempt])
    end
end

two_factor.html.erb

- @no_header = true

%section.login.w-full
  %div
    .branding.block
    %div
      %h2.mt-6.text-3xl.leading-9.font-extrabold.text-gray-900
        = t('users.two_factor.two_factor_authentication')

    - if @user.two_factor_otp_enabled?
      = form_for(resource, as: resource_name, url: session_path(resource_name), method: :post) do |f|
        .rounded-md.shadow-sm
          %div
            = f.label :otp_attempt do
              Two-factor authentication code
            = f.text_field :otp_attempt,
              autofocus: true,
              class: "p-3 w-full"

        .mt-4
          %p.text-sm.text-gray-700.leading-5 Enter the code from the two-factor app on your mobile device. If you've lost your device, you may enter one of your recovery codes.

        .mt-4
          = f.submit t('users.two_factor.verify_code'),
            class: "btn btn-primary border border-red-500"