[POC] Add 2FA interface

app/controllers/devise/passwords_controller.rb

@@ -37,6 +37,15 @@ def update
     if resource.errors.empty?
       resource.unlock_access! if unlockable?(resource)
       if sign_in_after_reset_password?
+        if resource.respond_to?(:two_factor_enabled?) && resource.two_factor_enabled?
+          session["devise.two_factor.resource_id"] = resource.id
+          session["devise.two_factor.remember_me"] = false
+          default_method = resource.enabled_two_factors.first
+          set_flash_message!(:notice, :updated_two_factor_required)
+          respond_with resource, location: new_two_factor_challenge_path(resource_name, default_method)
+          return
+        end
+
         flash_message = resource.active_for_authentication? ? :updated : :updated_not_active
         set_flash_message!(:notice, flash_message)
         resource.after_database_authentication

app/controllers/devise/two_factor_controller.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+class Devise::TwoFactorController < DeviseController
+  prepend_before_action :require_no_authentication
+  prepend_before_action :set_authenticating_resource
+
+  # Extensions can inject custom actions or override defaults via on_load
+  ActiveSupport.run_load_hooks(:devise_two_factor_controller, self)
+
+  # Auto-generate default new_<method> actions for each registered 2FA module.
+  # Extensions that injected a custom action via on_load won't be overwritten.
+  Devise.two_factor_method_configs.each_key do |mod|
+    define_method(:"new_#{mod}") {} unless method_defined?(:"new_#{mod}")
+  end
+
+  # POST /users/two_factor
+  # All methods POST here. Warden picks the right strategy via valid?.
+  def create
+    self.resource = warden.authenticate!(auth_options)
+    set_flash_message!(:notice, :signed_in, scope: :"devise.sessions")
+    sign_in(resource_name, resource)
+    yield resource if block_given?
+    respond_with resource, location: after_sign_in_path_for(resource)
+  end
+
+  protected
+
+  def auth_options
+    default_method = @authenticating_resource.enabled_two_factors.first
+    { scope: resource_name, recall: "#{controller_path}#new_#{default_method}" }
+  end
+
+  def translation_scope
+    'devise.two_factor'
+  end
+
+  private
+
+  def set_authenticating_resource
+    resource_id = session["devise.two_factor.resource_id"]
+    @authenticating_resource = resource_class.where(id: resource_id).first if resource_id
+    return if @authenticating_resource
+
+    set_flash_message!(:alert, :sign_in_not_initiated, scope: :"devise.failure")
+    redirect_to new_session_path(resource_name)
+  end
+end

app/controllers/devise_controller.rb

@@ -6,6 +6,7 @@ class DeviseController < Devise.parent_controller.constantize
 
   if respond_to?(:helper)
     helper DeviseHelper
+    helper Devise::TwoFactorHelper
   end
 
   if respond_to?(:helper_method)

app/helpers/devise/two_factor_helper.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Devise
+  module TwoFactorHelper
+    # Renders the link partials provided by each registered 2FA extension,
+    # excluding the method currently being challenged. Each extension is
+    # expected to ship a `devise/two_factor/<method>_link` partial.
+    def two_factor_method_links(resource, current_method)
+      methods = resource.enabled_two_factors - [current_method]
+      safe_join(methods.map { |method| render "devise/two_factor/#{method}_link" })
+    end
+  end
+end

config/locales/en.yml

@@ -16,6 +16,8 @@ en:
       timeout: "Your session expired. Please sign in again to continue."
       unauthenticated: "You need to sign in or sign up before continuing."
       unconfirmed: "You have to confirm your email address before continuing."
+      two_factor_session_expired: "Your two-factor authentication session has expired. Please sign in again."
+      sign_in_not_initiated: "Please sign in first."
     mailer:
       confirmation_instructions:
         subject: "Confirmation instructions"
@@ -36,6 +38,7 @@ en:
       send_paranoid_instructions: "If your email address exists in our database, you will receive a password recovery link at your email address in a few minutes."
       updated: "Your password has been changed successfully. You are now signed in."
       updated_not_active: "Your password has been changed successfully."
+      updated_two_factor_required: "Your password has been changed successfully. Please complete two-factor authentication."
     registrations:
       destroyed: "Bye! Your account has been successfully cancelled. We hope to see you again soon."
       signed_up: "Welcome! You have signed up successfully."

lib/devise.rb

@@ -18,6 +18,7 @@ module Devise
   autoload :ParameterSanitizer, 'devise/parameter_sanitizer'
   autoload :TimeInflector,      'devise/time_inflector'
   autoload :TokenGenerator,     'devise/token_generator'
+  autoload :TwoFactor,          'devise/two_factor'
 
   module Controllers
     autoload :Helpers,        'devise/controllers/helpers'
@@ -40,6 +41,7 @@ module Mailers
   module Strategies
     autoload :Base,            'devise/strategies/base'
     autoload :Authenticatable, 'devise/strategies/authenticatable'
+    autoload :TwoFactor,       'devise/strategies/two_factor'
   end
 
   module Test
@@ -312,6 +314,14 @@ def self.mappings
   mattr_accessor :sign_in_after_change_password
   @@sign_in_after_change_password = true
 
+  # Global default for two_factor_methods per-model config.
+  mattr_accessor :two_factor_methods
+  @@two_factor_methods = []
+
+  # Registry of two-factor method configs set via register_two_factor_method.
+  mattr_reader :two_factor_method_configs
+  @@two_factor_method_configs = {}
+
   # Default way to set up Devise. Run rails generate devise_install to create
   # a fresh initializer with all configuration values.
   def self.setup
@@ -439,6 +449,59 @@ def self.add_module(module_name, options = {})
     Devise::Mapping.add_module module_name
   end
 
+
+  # Register available devise two factor methods.
+  # Third-party modules that intend to add a 2FA method need to be added explicitly using this method.
+  #
+  # Note that adding a module using this method does not cause it to be used in the authentication
+  # process. That requires the `:two_factor_authenticatable` module to be listed in the arguments passed
+  # to the 'devise' method in the model class definition along with the two factor method name listed under
+  # the `:two_factor_methods` argument passed to the 'devise' method.
+  #
+  # == Options:
+  #
+  #   +name+       - String representing the name of the 2FA method. This will be used to identify it.
+  #   +model+      - String representing the load path to a custom *model* for this 2FA method (to autoload.)
+  #   +strategy+   - Symbol representing if this module got a custom *strategy*.
+  #   +route+      - Generates extension-specific routes and URL helpers (e.g., credential management
+  #                  endpoints). This is separate from the core challenge/create routes that Devise
+  #                  generates automatically from +two_factor_methods+. Accepts true (defaults route
+  #                  name to the method name), a Symbol, or a Hash. Works the same as the +:route+
+  #                  option in +add_module+.
+  #
+  # == Examples:
+  #
+  #   Devise.register_two_factor_method(:my_two_factor_method)
+  #   Devise.register_two_factor_method(:my_two_factor_method, model: 'my_two_factor_method/model')
+  #   Devise.register_two_factor_method(:my_two_factor_method, model: 'my_two_factor_method/model', strategy: :my_two_factor_method, route: true)
+  #
+  def self.register_two_factor_method(name, options = {})
+    options.assert_valid_keys(:model, :strategy, :route)
+
+    two_factor_method_configs[name.to_sym] = options
+
+    STRATEGIES[name.to_sym] = options[:strategy] if options[:strategy]
+
+    if route = options[:route]
+      case route
+      when TrueClass
+        key, value = name, []
+      when Symbol
+        key, value = route, []
+      when Hash
+        key, value = route.keys.first, route.values.flatten
+      else
+        raise ArgumentError, ":route should be true, a Symbol or a Hash"
+      end
+
+      URL_HELPERS[key] ||= []
+      URL_HELPERS[key].concat(value)
+      URL_HELPERS[key].uniq!
+
+      ROUTES[name.to_sym] = key
+    end
+  end
+
   # Sets warden configuration using a block that will be invoked on warden
   # initialization.
   #

lib/devise/mapping.rb

@@ -84,15 +84,27 @@ def to
     end
 
     def strategies
-      @strategies ||= STRATEGIES.values_at(*self.modules).compact.uniq.reverse
+      @strategies ||= begin
+        keys = self.modules
+        if to.respond_to?(:two_factor_methods) && to.two_factor_methods
+          keys = keys + Array(to.two_factor_methods)
+        end
+        STRATEGIES.values_at(*keys).compact.uniq.reverse
+      end
     end
 
     def no_input_strategies
       self.strategies & Devise::NO_INPUT
     end
 
     def routes
-      @routes ||= ROUTES.values_at(*self.modules).compact.uniq
+      @routes ||= begin
+        keys = self.modules
+        if to.respond_to?(:two_factor_methods) && to.two_factor_methods
+          keys = keys + Array(to.two_factor_methods)
+        end
+        ROUTES.values_at(*keys).compact.uniq
+      end
     end
 
     def authenticatable?

lib/devise/models.rb

@@ -120,3 +120,4 @@ def devise_modules_hook!
 end
 
 require 'devise/models/authenticatable'
+require 'devise/models/two_factor_authenticatable'

lib/devise/models/two_factor_authenticatable.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Devise
+  module Models
+    module TwoFactorAuthenticatable
+      extend ActiveSupport::Concern
+
+      def self.required_fields(klass)
+        []
+      end
+
+      module ClassMethods
+        Devise::Models.config(self, :two_factor_methods)
+
+        def two_factor_methods=(methods)
+          @two_factor_methods = methods
+          Array(methods).each do |method_name|
+            config = Devise.two_factor_method_configs[method_name]
+            raise "Unknown two-factor method: #{method_name}. " \
+              "Did you call Devise.register_two_factor_method?" unless config
+            begin
+              require config[:model]
+            rescue LoadError
+              raise unless config[:model].camelize.safe_constantize
+            end
+            mod = config[:model].camelize.constantize
+            include mod
+          end
+        end
+
+        def two_factor_modules
+          Array(two_factor_methods)
+        end
+      end
+
+      def enabled_two_factors
+        self.class.two_factor_modules.select do |method_name|
+          send(:"#{method_name}_two_factor_enabled?")
+        end
+      end
+
+      def two_factor_enabled?
+        enabled_two_factors.any?
+      end
+    end
+  end
+end

lib/devise/modules.rb

@@ -13,6 +13,9 @@
   # Other authentications
   d.add_module :omniauthable, controller: :omniauth_callbacks,  route: :omniauth_callback
 
+  # Two-factor authentication
+  d.add_module :two_factor_authenticatable, controller: :two_factor, route: :two_factor
+
   # Misc after
   routes = [nil, :new, :edit]
   d.add_module :recoverable,  controller: :passwords,     route: { password: routes }

lib/devise/rails.rb

@@ -37,6 +37,14 @@ class Engine < ::Rails::Engine
       end
     end
 
+    initializer "devise.two_factor" do
+      config.after_initialize do
+        if Devise.two_factor_method_configs.any?
+          Devise.include_helpers(Devise::TwoFactor)
+        end
+      end
+    end
+
     initializer "devise.secret_key" do |app|
       Devise.secret_key ||= app.secret_key_base
 

lib/devise/rails/routes.rb

@@ -398,6 +398,25 @@ def devise_unlock(mapping, controllers) #:nodoc:
         end
       end
 
+      def devise_two_factor(mapping, controllers) #:nodoc:
+        return unless mapping.to.respond_to?(:two_factor_methods) && mapping.to.two_factor_methods.present?
+
+        controller = controllers[:two_factor] || "devise/two_factor"
+        two_factor_path = mapping.path_names[:two_factor] || "two_factor"
+
+        # Central POST endpoint — all methods submit here
+        post two_factor_path,
+          to: "#{controller}#create",
+          as: "two_factor"
+
+        # Per-method challenge routes
+        Array(mapping.to.two_factor_methods).each do |method_name|
+          get "#{two_factor_path}/#{method_name}/new",
+            to: "#{controller}#new_#{method_name}",
+            as: "new_two_factor_#{method_name}"
+        end
+      end
+
       def devise_registration(mapping, controllers) #:nodoc:
         path_names = {
           new: mapping.path_names[:sign_up],

lib/devise/strategies/database_authenticatable.rb

@@ -11,9 +11,13 @@ def authenticate!
         hashed = false
 
         if validate(resource){ hashed = true; resource.valid_password?(password) }
-          remember_me(resource)
-          resource.after_database_authentication
-          success!(resource)
+          if resource.respond_to?(:two_factor_enabled?) && resource.two_factor_enabled?
+            initiate_two_factor_authentication!(resource)
+          else
+            remember_me(resource)
+            resource.after_database_authentication
+            success!(resource)
+          end
         end
 
         # In paranoid mode, hash the password even when a resource doesn't exist for the given authentication key.
@@ -24,6 +28,20 @@ def authenticate!
           Devise.paranoid ? fail(:invalid) : fail(:not_found_in_database)
         end
       end
+
+      private
+
+      def initiate_two_factor_authentication!(resource)
+        session["devise.two_factor.resource_id"] = resource.id
+        session["devise.two_factor.remember_me"] = remember_me?
+        default_method = resource.enabled_two_factors.first
+        redirect!(new_two_factor_challenge_path(scope, default_method))
+      end
+
+      def new_two_factor_challenge_path(scope, method)
+        Rails.application.routes.url_helpers
+          .send(:"#{scope}_new_two_factor_#{method}_path")
+      end
     end
   end
 end