Refactor SessionsController

Also adds test case for providing an invalid 2FA code and then a valid
one without re-entering username and password.

app/controllers/sessions_controller.rb

@@ -29,10 +29,7 @@ class SessionsController < Devise::SessionsController
 
   def create
     super do |resource|
-      # Remove any lingering user data from login
-      session.delete(:user)
-
-      # User has successfully signed in, so clear any unused reset tokens
+      # User has successfully signed in, so clear any unused reset token
       if resource.reset_password_token.present?
         resource.update_attributes(reset_password_token: nil,
                                    reset_password_sent_at: nil)
@@ -46,34 +43,40 @@ class SessionsController < Devise::SessionsController
     params.require(:user).permit(:login, :password, :remember_me, :otp_attempt)
   end
 
-  def authenticate_with_two_factor
-    @user = User.by_login(user_params[:login])
+  def find_user
+    if user_params[:login]
+      User.by_login(user_params[:login])
+    elsif user_params[:otp_attempt] && session[:otp_user_id]
+      User.find(session[:otp_user_id])
+    end
+  end
 
-    if user_params[:otp_attempt].present? && session[:user]
-      if valid_otp_attempt?
-        # Insert the saved params from the session into the request parameters
-        # so they're available to Devise::Strategies::DatabaseAuthenticatable
-        request.params[:user].merge!(session[:user])
+  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]
+      if valid_otp_attempt?(user)
+        # Remove any lingering user data from login
+        session.delete(:otp_user_id)
+
+        sign_in(user)
       else
         @error = 'Invalid two-factor code'
         render :two_factor and return
       end
     else
-      if @user && @user.valid_password?(user_params[:password])
-        self.resource = @user
-
-        if resource.otp_required_for_login
-          # Login is valid, save the values to the session so we can prompt the
-          # user for a one-time password.
-          session[:user] = user_params
-          render :two_factor and return
-        end
+      if user && user.valid_password?(user_params[:password])
+        # Save the user's ID to session so we can ask for a one-time password
+        session[:otp_user_id] = user.id
+        render :two_factor and return
       end
     end
   end
 
-  def valid_otp_attempt?
-    @user.valid_otp?(user_params[:otp_attempt]) ||
-    @user.invalidate_otp_backup_code!(user_params[:otp_attempt])
+  def valid_otp_attempt?(user)
+    user.valid_otp?(user_params[:otp_attempt]) ||
+    user.invalidate_otp_backup_code!(user_params[:otp_attempt])
   end
 end

app/views/devise/sessions/two_factor.html.haml

@@ -7,7 +7,6 @@
         - if @error
           .alert.alert-danger
             = @error
-        = f.hidden_field :login
         = f.text_field :otp_attempt, class: 'form-control', placeholder: 'Two-factor authentication code', required: true, autofocus: true
         .prepend-top-20
           = f.submit "Verify code", class: "btn btn-save"

spec/features/login_spec.rb

@@ -31,6 +31,14 @@ feature 'Login' do
           enter_code('foo')
           expect(page).to have_content('Invalid two-factor code')
         end
+
+        it 'allows login with invalid code, then valid code' do
+          enter_code('foo')
+          expect(page).to have_content('Invalid two-factor code')
+
+          enter_code(user.current_otp)
+          expect(current_path).to eq root_path
+        end
       end
 
       context 'using backup code' do