2fa after auth

app/controllers/sessions_controller.rb

@@ -99,20 +99,6 @@ class SessionsController < ApplicationController
       @handoff = h
     end
 
-    session[:tfa_query] = pre_auth && pre_auth['need_2fa']
-    if session[:tfa_query] && params[:totp].present? && params[:totp].length == 6
-      session[:tfa_query] = false
-    end
-
-    if session[:tfa_query]
-      if pre_auth['need_2fa'] == 'totp'
-        flash[:notice] = nil
-        session[:tfa_query] = true
-        load_params
-        return render '2fa'
-      end
-    end
-
     begin
       res = if !@handoff && params[:unlock]
         ApiBackend::auth(
@@ -166,7 +152,13 @@ class SessionsController < ApplicationController
         # successful login
         return
       end
-    rescue ApiBackend::ApiError
+    rescue ApiBackend::ApiError => e
+      if e.api_msg == 'missing_2fa'
+        flash[:notice] = nil
+        session[:tfa_query] = true
+        load_params
+        return render '2fa'
+      end
     end
     return login_failed
   end