Minor edits to two_factor_recovery_codes API error catching

1 files changed, 9 insertions, 5 deletions

diff --git a/lib/api/internal.rb b/lib/api/internal.rb
index 5b54c11ef62..6e6efece7c4 100644
--- a/lib/api/internal.rb
+++ b/lib/api/internal.rb
@@ -105,15 +105,19 @@ module API
       post '/two_factor_recovery_codes' do
         status 200
 
-        key = Key.find(params[:key_id])
-        user = key.user
+        key = Key.find_by(id: params[:key_id])
+
+        unless key
+          return { 'success' => false, 'message' => 'Could not find the given key' }
+        end
 
-        # Make sure this isn't a deploy key
-        unless key.type.nil?
+        if key.is_a?(DeployKey)
           return { success: false, message: 'Deploy keys cannot be used to retrieve recovery codes' }
         end
 
-        unless user.present?
+        user = key.user
+
+        unless user
           return { success: false, message: 'Could not find a user for the given key' }
         end