From e2f9c87600e34a415d43c981e0182094b123771f Mon Sep 17 00:00:00 2001
From: Patricio Cano <suprnova32@gmail.com>
Date: Fri, 12 Aug 2016 16:16:12 -0500
Subject: [PATCH] Added checks for 2FA to the API `/sessions` endpoint and the
 Resource Owner Password Credentials flow.

---
 app/services/user_retrieval_service.rb       | 13 ++++++++
 config/initializers/doorkeeper.rb            |  2 +-
 lib/api/helpers.rb                           |  4 +++
 lib/api/session.rb                           |  1 +
 spec/requests/api/oauth_tokens_spec.rb       | 31 ++++++++++++++++++++
 spec/requests/api/session_spec.rb            | 10 +++++++
 spec/services/user_retrieval_service_spec.rb | 19 ++++++++++++
 7 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 app/services/user_retrieval_service.rb
 create mode 100644 spec/requests/api/oauth_tokens_spec.rb
 create mode 100644 spec/services/user_retrieval_service_spec.rb

diff --git a/app/services/user_retrieval_service.rb b/app/services/user_retrieval_service.rb
new file mode 100644
index 00000000000..299fafca32d
--- /dev/null
+++ b/app/services/user_retrieval_service.rb
@@ -0,0 +1,13 @@
+class UserRetrievalService
+  attr_accessor :login, :password
+
+  def initialize(login, password)
+    @login = login
+    @password = password
+  end
+
+  def execute
+    user = Gitlab::Auth.find_with_user_password(login, password)
+    user unless user.two_factor_enabled?
+  end
+end
\ No newline at end of file
diff --git a/config/initializers/doorkeeper.rb b/config/initializers/doorkeeper.rb
index 618dba74151..f78f0cf7c5c 100644
--- a/config/initializers/doorkeeper.rb
+++ b/config/initializers/doorkeeper.rb
@@ -12,7 +12,7 @@ Doorkeeper.configure do
   end
 
   resource_owner_from_credentials do |routes|
-    Gitlab::Auth.find_with_user_password(params[:username], params[:password])
+    UserRetrievalService.new(params[:username], params[:password]).execute
   end
 
   # If you want to restrict access to the web interface for adding oauth authorized applications, you need to declare the block below.
diff --git a/lib/api/helpers.rb b/lib/api/helpers.rb
index d0469d6602d..bbd647684a4 100644
--- a/lib/api/helpers.rb
+++ b/lib/api/helpers.rb
@@ -275,6 +275,10 @@ module API
       end
     end
 
+    def render_2fa_error!
+      render_api_error!('401 You have 2FA enabled. Please use a personal access token to access the API', 401)
+    end
+
     def render_api_error!(message, status)
       error!({ 'message' => message }, status)
     end
diff --git a/lib/api/session.rb b/lib/api/session.rb
index 56c202f1294..b26be3be22e 100644
--- a/lib/api/session.rb
+++ b/lib/api/session.rb
@@ -14,6 +14,7 @@ module API
       user = Gitlab::Auth.find_with_user_password(params[:email] || params[:login], params[:password])
 
       return unauthorized! unless user
+      return render_2fa_error! if user.two_factor_enabled?
       present user, with: Entities::UserLogin
     end
   end
diff --git a/spec/requests/api/oauth_tokens_spec.rb b/spec/requests/api/oauth_tokens_spec.rb
new file mode 100644
index 00000000000..4730e9aa13c
--- /dev/null
+++ b/spec/requests/api/oauth_tokens_spec.rb
@@ -0,0 +1,31 @@
+require 'spec_helper'
+
+describe API::API, api: true  do
+  include ApiHelpers
+
+  context 'Resource Owner Password Credentials' do
+    def request_oauth_token(user)
+      post '/oauth/token', username: user.username, password: user.password, grant_type: 'password'
+    end
+
+    context 'when user has 2FA enabled' do
+      it 'does not create an access token' do
+        user = create(:user, :two_factor)
+        request_oauth_token(user)
+
+        expect(response).to have_http_status(401)
+        expect(json_response['error']).to eq('invalid_grant')
+      end
+    end
+
+    context 'when user does not have 2FA enabled' do
+      it 'creates an access token' do
+        user = create(:user)
+        request_oauth_token(user)
+
+        expect(response).to have_http_status(200)
+        expect(json_response['access_token']).not_to be_nil
+      end
+    end
+  end
+end
\ No newline at end of file
diff --git a/spec/requests/api/session_spec.rb b/spec/requests/api/session_spec.rb
index 519e7ce12ad..09f9192e7a8 100644
--- a/spec/requests/api/session_spec.rb
+++ b/spec/requests/api/session_spec.rb
@@ -17,6 +17,16 @@ describe API::API, api: true  do
         expect(json_response['can_create_project']).to eq(user.can_create_project?)
         expect(json_response['can_create_group']).to eq(user.can_create_group?)
       end
+
+      context 'with 2FA enabled' do
+        it 'rejects sign in attempts' do
+          user = create(:user, :two_factor)
+
+          post api('/session'), email: user.email, password: user.password
+
+          expect(response).to have_http_status(401)
+        end
+      end
     end
 
     context 'when email has case-typo and password is valid' do
diff --git a/spec/services/user_retrieval_service_spec.rb b/spec/services/user_retrieval_service_spec.rb
new file mode 100644
index 00000000000..82a00e44345
--- /dev/null
+++ b/spec/services/user_retrieval_service_spec.rb
@@ -0,0 +1,19 @@
+require 'spec_helper'
+
+describe UserRetrievalService, services: true do
+  context 'user retrieval' do
+    it 'retrieves the correct user' do
+      user = create(:user)
+      retrieved_user = described_class.new(user.username, user.password).execute
+
+      expect(retrieved_user).to eq(user)
+    end
+
+    it 'returns nil when 2FA is enabled' do
+      user = create(:user, :two_factor)
+      retrieved_user = described_class.new(user.username, user.password).execute
+
+      expect(retrieved_user).to be_nil
+    end
+  end
+end
\ No newline at end of file
-- 
2.30.9