2 files changed, 106 insertions, 7 deletions

diff --git a/app/controllers/sessions_controller.rb b/app/controllers/sessions_controller.rb
index 8bd8fbb692f..21505442e35 100644
--- a/app/controllers/sessions_controller.rb
+++ b/app/controllers/sessions_controller.rb
@@ -1,5 +1,5 @@
 class SessionsController < Devise::SessionsController
-  prepend_before_action :two_factor_enabled?, only: :create
+  prepend_before_action :authenticate_with_two_factor, only: :create
 
   def new
     redirect_path =
@@ -29,6 +29,9 @@ 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
       if resource.reset_password_token.present?
         resource.update_attributes(reset_password_token: nil,
@@ -39,24 +42,38 @@ class SessionsController < Devise::SessionsController
 
   private
 
-  def two_factor_enabled?
-    user_params = params[:user]
+  def user_params
+    params.require(:user).permit(:login, :password, :remember_me, :otp_attempt)
+  end
+
+  def authenticate_with_two_factor
     @user = User.by_login(user_params[:login])
 
-    if user_params[:otp_attempt].present?
-      unless @user.valid_otp?(user_params[:otp_attempt]) ||
-        @user.recovery_code?(user_params[:otp_attempt])
+    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])
+      else
         @error = 'Invalid two-factor code'
         render :two_factor and return
       end
     else
-      if @user && @user.valid_password?(params[:user][:password])
+      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
       end
     end
   end
+
+  def valid_otp_attempt?
+    @user.valid_otp?(user_params[:otp_attempt]) ||
+    @user.invalidate_otp_backup_code!(user_params[:otp_attempt])
+  end
 end
diff --git a/spec/features/login_spec.rb b/spec/features/login_spec.rb
new file mode 100644
index 00000000000..5ffaae54766
--- /dev/null
+++ b/spec/features/login_spec.rb
@@ -0,0 +1,82 @@
+require 'spec_helper'
+
+feature 'Login' do
+  let(:user) { create(:user) }
+
+  context 'with two-factor authentication' do
+    before do
+      user.otp_required_for_login = true
+      user.otp_secret = User.generate_otp_secret
+      user.save!
+    end
+
+    context 'with valid username/password' do
+      before do
+        login_with(user)
+        expect(page).to have_content('Two-factor Authentication')
+      end
+
+      def enter_code(code)
+        fill_in 'Two-factor authentication code', with: code
+        click_button 'Verify code'
+      end
+
+      context 'using one-time code' do
+        it 'allows login with valid code' do
+          enter_code(user.current_otp)
+          expect(current_path).to eq root_path
+        end
+
+        it 'blocks login with invalid code' do
+          enter_code('foo')
+          expect(page).to have_content('Invalid two-factor code')
+        end
+      end
+
+      context 'using backup code' do
+        let(:codes) { user.generate_otp_backup_codes! }
+
+        before do
+          expect(codes.size).to eq 5
+
+          # Because `generate_otp_backup_codes!` doesn't actually do this...
+          user.save
+        end
+
+        context 'with valid code' do
+          it 'allows login' do
+            enter_code(codes.sample)
+            expect(current_path).to eq root_path
+          end
+
+          it 'invalidates the used code' do
+            # FIXME (rspeicher): Broken library is broken
+            expect { enter_code(codes.sample) }.to change { user.otp_backup_codes.size }.by(-1)
+          end
+        end
+
+        context 'with invalid code' do
+          it 'blocks login' do
+            # FIXME (rspeicher): Broken library is broken
+            code = codes.sample
+            expect(user.invalidate_otp_backup_code!(code)).to eq true
+            expect(user.otp_backup_codes.size).to eq 4 # Passes
+            user.save!
+            user.reload
+            expect(user.otp_backup_codes.size).to eq 4 # Fails... WAT?!
+
+            enter_code(code)
+            expect(page).to have_content('Invalid two-factor code')
+          end
+        end
+      end
+    end
+  end
+
+  context 'without two-factor authentication' do
+    it 'allows basic login' do
+      login_with(user)
+      expect(current_path).to eq root_path
+    end
+  end
+end