Add multi factor authentication with TOTP

Gemfile

@@ -10,6 +10,8 @@ gem "devise"
 gem "devise-security"
 gem "devise_token_auth", "~> 1.2"
 gem "omniauth-azure-activedirectory-v2", "~> 2.1.0"
+gem "rotp", "~> 6.3" # TOTP authentication
+gem "rqrcode", "~> 2.2" # QR code generation for TOTP
 gem "fog-aws"
 gem "jsonapi-serializer"
 gem "kaminari"

Gemfile.lock

@@ -92,8 +92,8 @@ GEM
     brakeman (6.1.1)
       racc
     builder (3.3.0)
-    bundler-audit (0.9.2)
-      bundler (>= 1.2.0, < 3)
+    bundler-audit (0.9.3)
+      bundler (>= 1.2.0)
       thor (~> 1.0)
     byebug (11.1.3)
     capybara (3.39.2)
@@ -106,6 +106,7 @@ GEM
       regexp_parser (>= 1.5, < 3.0)
       xpath (~> 3.2)
     childprocess (5.0.0)
+    chunky_png (1.4.0)
     clockwork (3.0.2)
       activesupport
       tzinfo
@@ -371,7 +372,12 @@ GEM
       actionpack (>= 7.0)
       railties (>= 7.0)
     rexml (3.4.2)
+    rotp (6.3.0)
     rouge (4.2.0)
+    rqrcode (2.2.0)
+      chunky_png (~> 1.0)
+      rqrcode_core (~> 1.0)
+    rqrcode_core (1.2.0)
     rspec-core (3.12.2)
       rspec-support (~> 3.12.0)
     rspec-expectations (3.12.3)
@@ -510,6 +516,8 @@ DEPENDENCIES
   pundit
   rack-cors
   rails (~> 8.0.2)
+  rotp (~> 6.3)
+  rqrcode (~> 2.2)
   rspec-rails
   secure_headers (>= 7.1.0)
   shoulda-matchers
@@ -521,7 +529,7 @@ DEPENDENCIES
   with_advisory_lock
 
 RUBY VERSION
-   ruby 3.3.7p123
+  ruby 3.3.7p123
 
 BUNDLED WITH
-   4.0.4
+  4.0.4

app/controllers/multi_factor_controller.rb

@@ -16,38 +16,17 @@ class MultiFactorController < ApplicationController
   ##
   # Verifies a multi-factor authentication code and issues auth tokens.
   #
-  # After a user successfully authenticates with email/password,
-  # they receive a temp_token and an OTP via email. This endpoint
-  # validates the OTP and exchanges the temp_token for real auth tokens.
+  # Supports both email OTP and TOTP verification, as well as backup codes.
+  # The user's MFA method is detected automatically.
   #
   # @api POST /auth/verify_multi_factor
   # @param temp_token [String] temporary token from initial sign-in
-  # @param otp_code [String] 6-digit code from email
+  # @param otp_code [String] 6-digit code from email/app or 8-char backup code
   # @return [JSON] user data and auth tokens in headers
   # @status 200 OTP verified, auth tokens issued
   # @status 401 Invalid or expired OTP/temp_token
+  # @status 403 Account locked due to failed attempts
   # @status 422 Missing required parameters
-  #
-  # @example Request
-  #   POST /auth/verify_multi_factor
-  #   {
-  #     "temp_token": "abc123...",
-  #     "otp_code": "123456"
-  #   }
-  #
-  # @example Response Headers
-  #   access-token: "xyz789..."
-  #   client: "client_id"
-  #   uid: "user@example.com"
-  #
-  # @example Response Body
-  #   {
-  #     "data": {
-  #       "id": 1,
-  #       "email": "user@example.com",
-  #       ...
-  #     }
-  #   }
   def verify
     if params[:temp_token].blank? || params[:otp_code].blank?
       return render json: {errors: ["temp_token and otp_code are required"]}, status: 422
@@ -57,40 +36,32 @@ def verify
     return render json: {errors: ["Invalid or expired temp token"]}, status: :unauthorized unless user_id
 
     @resource = User.find_by(id: user_id)
-
     return render json: {errors: ["User not found"]}, status: :unauthorized unless @resource
 
-    if @resource.multi_factor_email_code_expired?
-      return render json: {errors: ["Multi-factor code has expired"]}, status: :unauthorized
+    # Check if account is locked due to failed MFA attempts
+    if @resource.mfa_locked?
+      return render json: {
+        errors: ["Account temporarily locked due to too many failed attempts. Try again in 30 minutes."]
+      }, status: :forbidden
     end
 
-    unless @resource.validate_multi_factor_email_code(params[:otp_code])
-      return render json: {errors: ["Invalid multi-factor code"]}, status: :unauthorized
+    # Verify code based on user's MFA method
+    code_valid = case @resource.mfa_method
+    when :totp
+      verify_totp_code
+    when :email_otp
+      verify_email_otp_code
+    else
+      false
     end
 
-    # OTP verified! Generate auth tokens
-    Rails.cache.delete("otp_temp_token:#{params[:temp_token]}")
-    @resource.update_columns(multi_factor_email_code: nil, multi_factor_email_code_sent_at: nil)
-
-    @client_id = SecureRandom.urlsafe_base64(nil, false)
-    @token = SecureRandom.urlsafe_base64(nil, false)
-    @expiry = (Time.current + DeviseTokenAuth.token_lifespan).to_i
-
-    @resource.tokens ||= {}
-    @resource.tokens[@client_id] = {
-      token: BCrypt::Password.create(@token),
-      expiry: @expiry
-    }
-
-    @resource.save!
-
-    response.headers["access-token"] = @token
-    response.headers["client"] = @client_id
-    response.headers["uid"] = @resource.uid
-    response.headers["expiry"] = @expiry.to_s
-    response.headers["token-type"] = "Bearer"
+    unless code_valid
+      @resource.increment_mfa_failed_attempts!
+      return render json: {errors: ["Invalid multi-factor code"]}, status: :unauthorized
+    end
 
-    render json: {data: @resource.token_validation_response}
+    # Code verified! Generate auth tokens and reset failed attempts
+    complete_authentication
   end
 
   ##
@@ -131,4 +102,61 @@ def resend
 
     render json: {message: "Multi-factor code re-sent to your email"}, status: :ok
   end
+
+  private
+
+  ##
+  # Verifies TOTP code or backup code.
+  def verify_totp_code
+    # Check if it's a backup code (8 characters, alphanumeric)
+    if params[:otp_code].length == 8 && params[:otp_code].match?(/^[a-z0-9]+$/i)
+      return BackupCode.use_code(@resource, params[:otp_code])
+    end
+
+    # Otherwise treat as TOTP code
+    @resource.validate_totp_code(params[:otp_code])
+  end
+
+  ##
+  # Verifies email OTP code.
+  def verify_email_otp_code
+    return false if @resource.multi_factor_email_code_expired?
+    @resource.validate_multi_factor_email_code(params[:otp_code])
+  end
+
+  ##
+  # Completes authentication by issuing tokens.
+  def complete_authentication
+    # Clear temp token and OTP data
+    Rails.cache.delete("otp_temp_token:#{params[:temp_token]}")
+
+    if @resource.mfa_method == :email_otp
+      @resource.update_columns(multi_factor_email_code: nil, multi_factor_email_code_sent_at: nil)
+    end
+
+    # Reset failed attempts
+    @resource.reset_mfa_failed_attempts!
+
+    # Generate auth tokens
+    @client_id = SecureRandom.urlsafe_base64(nil, false)
+    @token = SecureRandom.urlsafe_base64(nil, false)
+    @expiry = (Time.current + DeviseTokenAuth.token_lifespan).to_i
+
+    @resource.tokens ||= {}
+    @resource.tokens[@client_id] = {
+      token: BCrypt::Password.create(@token),
+      expiry: @expiry
+    }
+
+    @resource.save!
+
+    # Set response headers
+    response.headers["access-token"] = @token
+    response.headers["client"] = @client_id
+    response.headers["uid"] = @resource.uid
+    response.headers["expiry"] = @expiry.to_s
+    response.headers["token-type"] = "Bearer"
+
+    render json: {data: @resource.token_validation_response}
+  end
 end

app/controllers/passwords_controller.rb

@@ -3,7 +3,7 @@ def redirect_options
     trusted_host = URI.parse(ENV["CLIENT_URL"].to_s.strip).host
     redirect_host = begin
       URI.parse(params[:redirect_url]).host
-    rescue
+    rescue URI::InvalidURIError, ArgumentError
       nil
     end
 

app/controllers/sessions_controller.rb

@@ -21,16 +21,18 @@ class SessionsController < DeviseTokenAuth::SessionsController
   #
   # When a user with MFA enabled attempts to sign in, this endpoint:
   # 1. Validates email/password credentials
-  # 2. Generates and sends an OTP via email
-  # 3. Returns a temp_token for OTP verification
+  # 2. Detects the user's MFA method (TOTP or email OTP)
+  # 3. For email OTP: generates and sends an OTP via email
+  # 4. For TOTP: prompts for authenticator app code (no email sent)
+  # 5. Returns a temp_token for OTP verification
   #
   # For users without MFA, it proceeds with normal token authentication.
   #
   # @api POST /auth/sign_in
   # @param email [String] user's email address
   # @param password [String] user's password
   # @return [JSON] either temp_token (MFA) or auth tokens (no MFA)
-  # @status 202 MFA required, OTP sent
+  # @status 202 MFA required, OTP sent or TOTP prompt
   # @status 200 Signed in successfully (no MFA)
   # @status 401 Invalid credentials
   #
@@ -41,14 +43,22 @@ class SessionsController < DeviseTokenAuth::SessionsController
   #     "password": "password123"
   #   }
   #
-  # @example Response (MFA required)
+  # @example Response (Email OTP required)
   #   HTTP 202 Accepted
   #   {
   #     "otp_required": true,
   #     "temp_token": "abc123...",
   #     "message": "Multi-factor code sent to your email"
   #   }
   #
+  # @example Response (TOTP required)
+  #   HTTP 202 Accepted
+  #   {
+  #     "otp_required": true,
+  #     "temp_token": "abc123...",
+  #     "message": "Enter code from your authenticator app"
+  #   }
+  #
   # @example Response (no MFA)
   #   HTTP 200 OK
   #   Headers:
@@ -85,13 +95,29 @@ def create
       # Password is valid - reset failed attempts
       @resource.update_column(:failed_attempts, 0) if @resource.failed_attempts > 0
 
-      # STEP 3: MFA FLOW - check if MFA is enabled
-      if Rails.application.config.enable_mfa
-        @resource.generate_and_send_multi_factor_email!
+      # STEP 3: MFA FLOW - check if MFA is required
+      if Rails.application.config.require_mfa
+        # Generate temp token for MFA verification
         temp_token = SecureRandom.urlsafe_base64(32)
         Rails.cache.write("otp_temp_token:#{temp_token}", @resource.id, expires_in: 5.minutes)
 
-        render json: {otp_required: true, temp_token:, message: "Multi-factor code sent to your email"}, status: :accepted
+        case @resource.mfa_method
+        when :totp
+          # TOTP users use their authenticator app - no email sent
+          render json: {
+            otp_required: true,
+            temp_token:,
+            message: "Enter code from your authenticator app"
+          }, status: :accepted
+        else
+          # Email OTP users and users without MFA get a code via email
+          @resource.generate_and_send_multi_factor_email!
+          render json: {
+            otp_required: true,
+            temp_token:,
+            message: "Multi-factor code sent to your email"
+          }, status: :accepted
+        end
       else
         super
       end