1 files changed, 68 insertions, 2 deletions

diff --git a/lib/gitlab/shell.rb b/lib/gitlab/shell.rb
index a8a4ec996c4..392f66c99d3 100644
--- a/lib/gitlab/shell.rb
+++ b/lib/gitlab/shell.rb
@@ -183,6 +183,8 @@ module Gitlab
     #   add_key("key-42", "sha-rsa ...")
     #
     def add_key(key_id, key_content)
+      return unless self.authorized_keys_enabled?
+
       gitlab_shell_fast_execute([gitlab_shell_keys_path,
                                  'add-key', key_id, self.class.strip_key(key_content)])
     end
@@ -192,6 +194,8 @@ module Gitlab
     # Ex.
     #   batch_add_keys { |adder| adder.add_key("key-42", "sha-rsa ...") }
     def batch_add_keys(&block)
+      return unless self.authorized_keys_enabled?
+
       IO.popen(%W(#{gitlab_shell_path}/bin/gitlab-keys batch-add-keys), 'w') do |io|
         yield(KeyAdder.new(io))
       end
@@ -202,10 +206,11 @@ module Gitlab
     # Ex.
     #   remove_key("key-342", "sha-rsa ...")
     #
-    def remove_key(key_id, key_content)
+    def remove_key(key_id, key_content = nil)
+      return unless self.authorized_keys_enabled?
+
       args = [gitlab_shell_keys_path, 'rm-key', key_id]
       args << key_content if key_content
-
       gitlab_shell_fast_execute(args)
     end
 
@@ -215,9 +220,62 @@ module Gitlab
     #   remove_all_keys
     #
     def remove_all_keys
+      return unless self.authorized_keys_enabled?
+
       gitlab_shell_fast_execute([gitlab_shell_keys_path, 'clear'])
     end
 
+    # Remove ssh keys from gitlab shell that are not in the DB
+    #
+    # Ex.
+    #   remove_keys_not_found_in_db
+    #
+    def remove_keys_not_found_in_db
+      return unless self.authorized_keys_enabled?
+
+      Rails.logger.info("Removing keys not found in DB")
+
+      batch_read_key_ids do |ids_in_file|
+        ids_in_file.uniq!
+        keys_in_db = Key.where(id: ids_in_file)
+
+        next unless ids_in_file.size > keys_in_db.count # optimization
+
+        ids_to_remove = ids_in_file - keys_in_db.pluck(:id)
+        ids_to_remove.each do |id|
+          Rails.logger.info("Removing key-#{id} not found in DB")
+          remove_key("key-#{id}")
+        end
+      end
+    end
+
+    # Iterate over all ssh key IDs from gitlab shell, in batches
+    #
+    # Ex.
+    #   batch_read_key_ids { |batch| keys = Key.where(id: batch) }
+    #
+    def batch_read_key_ids(batch_size: 100, &block)
+      return unless self.authorized_keys_enabled?
+
+      list_key_ids do |key_id_stream|
+        key_id_stream.lazy.each_slice(batch_size) do |lines|
+          key_ids = lines.map { |l| l.chomp.to_i }
+          yield(key_ids)
+        end
+      end
+    end
+
+    # Stream all ssh key IDs from gitlab shell, separated by newlines
+    #
+    # Ex.
+    #   list_key_ids
+    #
+    def list_key_ids(&block)
+      return unless self.authorized_keys_enabled?
+
+      IO.popen(%W(#{gitlab_shell_path}/bin/gitlab-keys list-key-ids), &block)
+    end
+
     # Add empty directory for storing repositories
     #
     # Ex.
@@ -333,6 +391,14 @@ module Gitlab
       File.join(gitlab_shell_path, 'bin', 'gitlab-keys')
     end
 
+    def authorized_keys_enabled?
+      # Return true if nil to ensure the authorized_keys methods work while
+      # fixing the authorized_keys file during migration.
+      return true if Gitlab::CurrentSettings.current_application_settings.authorized_keys_enabled.nil?
+
+      Gitlab::CurrentSettings.current_application_settings.authorized_keys_enabled
+    end
+
     private
 
     def gitlab_projects(shard_path, disk_path)