Merge branch 'fix-authorized-keys-enabled-default-2738' into 'master'

Default `authorized_keys_enabled` setting to true

Closes #2738

See merge request !2240

GITLAB_SHELL_VERSION

-5.0.6
+5.1.0

app/controllers/admin/application_settings_controller.rb

@@ -180,7 +180,8 @@ class Admin::ApplicationSettingsController < Admin::ApplicationController
       :check_namespace_plan,
       :mirror_max_delay,
       :mirror_max_capacity,
-      :mirror_capacity_threshold
+      :mirror_capacity_threshold,
+      :authorized_keys_enabled
     ]
   end
 end

app/views/admin/application_settings/_form.html.haml

@@ -671,6 +671,22 @@
           installations. Set to 0 to completely disable polling.
           = link_to icon('question-circle'), help_page_path('administration/polling')
 
+  %fieldset
+    %legend Performance optimization
+    .form-group
+      .col-sm-offset-2.col-sm-10
+        .checkbox
+          = f.label :authorized_keys_enabled do
+            = f.check_box :authorized_keys_enabled
+            Write to "authorized_keys" file
+          .help-block
+            By default, we write to the "authorized_keys" file to support Git
+            over SSH without additional configuration. GitLab can be optimized
+            to authenticate SSH keys via the database file. Only uncheck this
+            if you have configured your OpenSSH server to use the
+            AuthorizedKeysCommand. Click on the help icon for more details.
+            = link_to icon('question-circle'), help_page_path('administration/operations/speed_up_ssh', anchor: 'the-solution')
+
   - if Gitlab::Geo.license_allows?
     %fieldset
       %legend GitLab Geo

changelogs/unreleased-ee/fix-authorized-keys-enabled-default-2738.yml

+---
+title: Fix locked and stale SSH keys file from 9.3.0 upgrade
+merge_request: 2240
+author:

db/migrate/20170531180233_add_authorized_keys_enabled_to_application_settings.rb

@@ -7,9 +7,13 @@ class AddAuthorizedKeysEnabledToApplicationSettings < ActiveRecord::Migration
   # Set this constant to true if this migration requires downtime.
   DOWNTIME = false
 
-  def change
-    # allow_null: true because we want to set the default based on if the
-    # instance is configured to use AuthorizedKeysCommand
-    add_column :application_settings, :authorized_keys_enabled, :boolean, allow_null: true
+  disable_ddl_transaction!
+
+  def up
+    add_column_with_default :application_settings, :authorized_keys_enabled, :boolean, default: true, allow_null: false
+  end
+
+  def down
+    remove_column :application_settings, :authorized_keys_enabled
   end
 end

db/migrate/20170626202753_update_authorized_keys_file.rb

+# See http://doc.gitlab.com/ce/development/migration_style_guide.html
+# for more information on how to write migrations for GitLab.
+
+class UpdateAuthorizedKeysFile < ActiveRecord::Migration
+  include Gitlab::Database::MigrationHelpers
+
+  class ApplicationSetting < ActiveRecord::Base
+    self.table_name = 'application_settings'
+  end
+
+  # Set this constant to true if this migration requires downtime.
+  DOWNTIME = false
+
+  # We started deploying 9.3.0 to GitLab.com at 2017-06-22 17:24 UTC, and the
+  # package was pushed some time before that. Some buffer room is included here.
+  DATETIME_9_3_0_RELEASED = DateTime.parse('2017-06-22T00:00:00+00:00')
+
+  def up
+    if authorized_keys_file_in_use_and_stale?
+      say 'The authorized_keys file is in use, and may be stale. Now bringing it up-to-date in the background...'
+
+      # Update nil authorized_keys_enabled to true to ensure that Gitlab::Shell
+      # key methods work properly for workers running 9.3.0 during the
+      # migration. If the setting remained nil, the workers would not edit the
+      # file.
+      update_nil_setting_to_true
+
+      update_authorized_keys_file_since(DATETIME_9_3_0_RELEASED)
+    else
+      say 'The authorized_keys file does not need to be updated. Skipping...'
+    end
+  end
+
+  def down
+    # Do nothing
+  end
+
+  def authorized_keys_file_in_use_and_stale?
+    return false unless ran_broken_migration?
+
+    @uncached_application_setting = ApplicationSetting.last
+
+    # If there is no ApplicationSetting record in the DB, then the instance was
+    # never in a state where `authorized_keys_enabled` field was `nil`. So the
+    # file is not stale.
+    return false unless @uncached_application_setting
+
+    if @uncached_application_setting.authorized_keys_enabled == false # not falsey!
+      # If authorized_keys_enabled is explicitly false, then the file is not in
+      # use, so it doesn't need to be fixed. I.e. GitLab.com.
+      #
+      # Unfortunately it is possible some users may have saved Application
+      # Settings without realizing the new option existed, and since it
+      # mistakenly defaulted to unchecked, now it is explicitly false. These
+      # users need this warning.
+      say false_negative_warning
+      return false
+    end
+
+    # If authorized_keys_enabled is true or nil, then we need to rebuild the
+    # file in case it is stale.
+    true
+  end
+
+  def ran_broken_migration?
+    # If the column is already fixed, then the migration wasn't run before now.
+    default_value = Gitlab::Database.postgresql? ? 'true' : '1'
+
+    column_has_no_default = !column_exists?(:application_settings, :authorized_keys_enabled, :boolean, default: default_value, null: false)
+    say "This GitLab installation was #{'never ' unless column_has_no_default}upgraded to exactly version 9.3.0."
+
+    column_has_no_default
+  end
+
+  def false_negative_warning
+    <<-MSG.strip_heredoc
+      WARNING
+
+      If you did not intentionally disable the "Write to authorized_keys file"
+      option in Application Settings as outlined in the Speed up SSH
+      documentation,
+
+      https://docs.gitlab.com/ee/administration/operations/speed_up_ssh.html
+
+      then the authorized_keys file may be out-of-date, affecting SSH
+      operations.
+
+      If you are affected, please check the "Write to authorized_keys file"
+      checkbox, and Save. Then rebuild the authorized_keys file as shown here:
+
+      https://docs.gitlab.com/ee/administration/raketasks/maintenance.html#rebuild-authorized_keys-file
+
+      For more information, see the issue:
+
+      https://gitlab.com/gitlab-org/gitlab-ee/issues/2738
+    MSG
+  end
+
+  def update_nil_setting_to_true
+    @uncached_application_setting.update_attribute(:authorized_keys_enabled, true)
+  end
+
+  def update_authorized_keys_file_since(cutoff_datetime)
+    job = ['UpdateAuthorizedKeysFileSince', [cutoff_datetime]]
+    BackgroundMigrationWorker.perform_bulk(job)
+  end
+end

db/migrate/20170627211700_add_default_to_authorized_keys_enabled_application_setting.rb

+# See http://doc.gitlab.com/ce/development/migration_style_guide.html
+# for more information on how to write migrations for GitLab.
+
+class AddDefaultToAuthorizedKeysEnabledApplicationSetting < ActiveRecord::Migration
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+
+  disable_ddl_transaction!
+
+  def up
+    change_column :application_settings, :authorized_keys_enabled, :boolean, default: true
+    change_column_null :application_settings, :authorized_keys_enabled, false, true
+  end
+
+  def down
+    change_column_null :application_settings, :authorized_keys_enabled, true
+    change_column :application_settings, :authorized_keys_enabled, :boolean, default: nil
+  end
+end

db/schema.rb

@@ -11,7 +11,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 20170622162730) do
+ActiveRecord::Schema.define(version: 20170627211700) do
 
   # These are extensions that must be enabled in order to support this database
   enable_extension "plpgsql"
@@ -141,7 +141,7 @@ ActiveRecord::Schema.define(version: 20170622162730) do
     t.integer "mirror_max_delay", default: 5, null: false
     t.integer "mirror_max_capacity", default: 100, null: false
     t.integer "mirror_capacity_threshold", default: 50, null: false
-    t.boolean "authorized_keys_enabled"
+    t.boolean "authorized_keys_enabled", default: true, null: false
     t.boolean "help_page_hide_commercial_content", default: false
     t.string "help_page_support_url"
   end

doc/administration/operations.md

@@ -6,3 +6,4 @@
 - [Cleaning up Redis sessions](operations/cleaning_up_redis_sessions.md)
 - [Understanding Unicorn and unicorn-worker-killer](operations/unicorn.md)
 - [Moving repositories to a new location](operations/moving_repositories.md)
+- [Speed up SSH operations](operations/speed_up_ssh.md)

doc/administration/operations/speed_up_ssh.md

@@ -51,6 +51,18 @@ sudo service sshd reload
 
 Confirm that SSH is working by removing your user's SSH key in the UI, adding a new one, and attempting to pull a repo.
 
+> **Warning:** Do not disable writes until SSH is confirmed to be working perfectly because the file will quickly become out-of-date.
+
+In the case of lookup failures (which are not uncommon), the `authorized_keys` file will still be scanned. So git SSH performance will still be slow for many users as long as a large file exists.
+
+You can disable any more writes to the `authorized_keys` file by unchecking `Write to "authorized_keys" file` in the Application Settings of your GitLab installation.
+
+![Write to authorized keys setting](img/write_to_authorized_keys_setting.png)
+
+Again, confirm that SSH is working by removing your user's SSH key in the UI, adding a new one, and attempting to pull a repo.
+
+Then you can backup and delete your `authorized_keys` file for best performance.
+
 ## How to go back to using the `authorized_keys` file
 
 This is a brief overview. Please refer to the above instructions for more context.

lib/gitlab/background_migration/update_authorized_keys_file_since.rb

+module Gitlab
+  module BackgroundMigration
+    class UpdateAuthorizedKeysFileSince
+      include Gitlab::ShellAdapter
+
+      class Key < ActiveRecord::Base
+        self.table_name = 'keys'
+
+        def shell_id
+          "key-#{id}"
+        end
+      end
+
+      delegate :remove_keys_not_found_in_db, to: :gitlab_shell
+
+      def perform(cutoff_datetime)
+        add_keys_since(cutoff_datetime)
+
+        remove_keys_not_found_in_db
+      end
+
+      def add_keys_since(cutoff_datetime)
+        start_key = Key.select(:id).where("created_at >= ?", cutoff_datetime).order('id ASC').take
+        if start_key
+          batch_add_keys_in_db_starting_from(start_key.id)
+        end
+      end
+
+      # Not added to Gitlab::Shell because I don't expect this to be used again
+      def batch_add_keys_in_db_starting_from(start_id)
+        Rails.logger.info("Adding all keys starting from ID: #{start_id}")
+
+        gitlab_shell.batch_add_keys do |adder|
+          Key.find_each(start: start_id, batch_size: 1000) do |key|
+            adder.add_key(key.shell_id, key.key)
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/shell.rb

@@ -197,6 +197,8 @@ module Gitlab
     #   add_key("key-42", "sha-rsa ...")
     #
     def add_key(key_id, key_content)
+      return unless self.authorized_keys_enabled?
+
       Gitlab::Utils.system_silent([gitlab_shell_keys_path,
                                    'add-key', key_id, self.class.strip_key(key_content)])
     end
@@ -206,6 +208,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
@@ -216,9 +220,12 @@ module Gitlab
     # Ex.
     #   remove_key("key-342", "sha-rsa ...")
     #
-    def remove_key(key_id, key_content)
-      Gitlab::Utils.system_silent([gitlab_shell_keys_path,
-                                   'rm-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::Utils.system_silent(args)
     end
 
     # Remove all ssh keys from gitlab shell
@@ -227,9 +234,62 @@ module Gitlab
     #   remove_all_keys
     #
     def remove_all_keys
+      return unless self.authorized_keys_enabled?
+
       Gitlab::Utils.system_silent([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.
@@ -356,5 +416,13 @@ module Gitlab
     def gitlab_shell_keys_path
       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 current_application_settings.authorized_keys_enabled.nil?
+
+      current_application_settings.authorized_keys_enabled
+    end
   end
 end

spec/lib/gitlab/background_migration/update_authorized_keys_file_since_spec.rb

+require 'spec_helper'
+
+describe Gitlab::BackgroundMigration::UpdateAuthorizedKeysFileSince do
+  let(:background_migration) { described_class.new }
+
+  describe '#perform' do
+    let!(:cutoff_datetime) { DateTime.now }
+
+    subject { background_migration.perform(cutoff_datetime) }
+
+    context 'when an SSH key was created after the cutoff datetime' do
+      before do
+        Timecop.freeze
+      end
+      after do
+        Timecop.return
+      end
+
+      before do
+        Timecop.travel 1.day.from_now
+        @key = create(:key)
+      end
+
+      it 'calls batch_add_keys_in_db_starting_from with the start key ID' do
+        expect(background_migration).to receive(:batch_add_keys_in_db_starting_from).with(@key.id)
+        subject
+      end
+    end
+
+    it 'calls remove_keys_not_found_in_db on Gitlab::Shell' do
+      expect_any_instance_of(Gitlab::Shell).to receive(:remove_keys_not_found_in_db)
+      subject
+    end
+  end
+
+  describe '#add_keys_since' do
+    let!(:cutoff_datetime) { DateTime.now }
+
+    subject { background_migration.add_keys_since(cutoff_datetime) }
+
+    before do
+      Timecop.freeze
+    end
+    after do
+      Timecop.return
+    end
+
+    context 'when an SSH key was created after the cutoff datetime' do
+      before do
+        Timecop.travel 1.day.from_now
+        @key = create(:key)
+        create(:key) # other key
+      end
+
+      it 'calls batch_add_keys_in_db_starting_from with the start key ID' do
+        expect(background_migration).to receive(:batch_add_keys_in_db_starting_from).with(@key.id)
+        subject
+      end
+    end
+
+    context 'when an SSH key was not created after the cutoff datetime' do
+      it 'does not call batch_add_keys_in_db_starting_from' do
+        expect(background_migration).not_to receive(:batch_add_keys_in_db_starting_from)
+        subject
+      end
+    end
+  end
+
+  describe '#remove_keys_not_found_in_db' do
+    it 'calls remove_keys_not_found_in_db on Gitlab::Shell' do
+      expect_any_instance_of(Gitlab::Shell).to receive(:remove_keys_not_found_in_db)
+      background_migration.remove_keys_not_found_in_db
+    end
+  end
+
+  describe '#batch_add_keys_in_db_starting_from' do
+    context 'when there are many keys in the DB' do
+      before do
+        @keys = []
+        10.times do
+          @keys << create(:key)
+        end
+      end
+
+      it 'adds all the keys in the DB, starting from the given ID, to the authorized_keys file' do
+        Gitlab::Shell.new.remove_all_keys
+
+        background_migration.batch_add_keys_in_db_starting_from(@keys[3].id)
+
+        file = File.read(Rails.root.join('tmp/tests/.ssh/authorized_keys'))
+        expect(file.scan(/ssh-rsa/).count).to eq(7)
+
+        expect(file).not_to include(Gitlab::Shell.strip_key(@keys[0].key))
+        expect(file).not_to include(Gitlab::Shell.strip_key(@keys[2].key))
+        expect(file).to include(Gitlab::Shell.strip_key(@keys[3].key))
+        expect(file).to include(Gitlab::Shell.strip_key(@keys[9].key))
+      end
+    end
+  end
+end

spec/lib/gitlab/shell_spec.rb

@@ -104,13 +104,305 @@ describe Gitlab::Shell, lib: true do
   end
 
   describe '#add_key' do
-    it 'removes trailing garbage' do
-      allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
-      expect(Gitlab::Utils).to receive(:system_silent).with(
-        [:gitlab_shell_keys_path, 'add-key', 'key-123', 'ssh-rsa foobar']
-      )
+    context 'when authorized_keys_enabled is true' do
+      it 'removes trailing garbage' do
+        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
+        expect(Gitlab::Utils).to receive(:system_silent).with(
+          [:gitlab_shell_keys_path, 'add-key', 'key-123', 'ssh-rsa foobar']
+        )
+
+        gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+      end
+    end
+
+    context 'when authorized_keys_enabled is false' do
+      before do
+        stub_application_setting(authorized_keys_enabled: false)
+      end
+
+      it 'does nothing' do
+        expect(Gitlab::Utils).not_to receive(:system_silent)
+
+        gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+      end
+    end
+
+    context 'when authorized_keys_enabled is nil' do
+      before do
+        stub_application_setting(authorized_keys_enabled: nil)
+      end
 
-      gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+      it 'removes trailing garbage' do
+        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
+        expect(Gitlab::Utils).to receive(:system_silent).with(
+          [:gitlab_shell_keys_path, 'add-key', 'key-123', 'ssh-rsa foobar']
+        )
+
+        gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+      end
+    end
+  end
+
+  describe '#batch_add_keys' do
+    context 'when authorized_keys_enabled is true' do
+      it 'instantiates KeyAdder' do
+        expect_any_instance_of(Gitlab::Shell::KeyAdder).to receive(:add_key).with('key-123', 'ssh-rsa foobar')
+
+        gitlab_shell.batch_add_keys do |adder|
+          adder.add_key('key-123', 'ssh-rsa foobar')
+        end
+      end
+    end
+
+    context 'when authorized_keys_enabled is false' do
+      before do
+        stub_application_setting(authorized_keys_enabled: false)
+      end
+
+      it 'does nothing' do
+        expect_any_instance_of(Gitlab::Shell::KeyAdder).not_to receive(:add_key)
+
+        gitlab_shell.batch_add_keys do |adder|
+          adder.add_key('key-123', 'ssh-rsa foobar')
+        end
+      end
+    end
+
+    context 'when authorized_keys_enabled is nil' do
+      before do
+        stub_application_setting(authorized_keys_enabled: nil)
+      end
+
+      it 'instantiates KeyAdder' do
+        expect_any_instance_of(Gitlab::Shell::KeyAdder).to receive(:add_key).with('key-123', 'ssh-rsa foobar')
+
+        gitlab_shell.batch_add_keys do |adder|
+          adder.add_key('key-123', 'ssh-rsa foobar')
+        end
+      end
+    end
+  end
+
+  describe '#remove_key' do
+    context 'when authorized_keys_enabled is true' do
+      it 'removes trailing garbage' do
+        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
+        expect(Gitlab::Utils).to receive(:system_silent).with(
+          [:gitlab_shell_keys_path, 'rm-key', 'key-123', 'ssh-rsa foobar']
+        )
+
+        gitlab_shell.remove_key('key-123', 'ssh-rsa foobar')
+      end
+    end
+
+    context 'when authorized_keys_enabled is false' do
+      before do
+        stub_application_setting(authorized_keys_enabled: false)
+      end
+
+      it 'does nothing' do
+        expect(Gitlab::Utils).not_to receive(:system_silent)
+
+        gitlab_shell.remove_key('key-123', 'ssh-rsa foobar')
+      end
+    end
+
+    context 'when authorized_keys_enabled is nil' do
+      before do
+        stub_application_setting(authorized_keys_enabled: nil)
+      end
+
+      it 'removes trailing garbage' do
+        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
+        expect(Gitlab::Utils).to receive(:system_silent).with(
+          [:gitlab_shell_keys_path, 'rm-key', 'key-123', 'ssh-rsa foobar']
+        )
+
+        gitlab_shell.remove_key('key-123', 'ssh-rsa foobar')
+      end
+    end
+
+    context 'when key content is not given' do
+      it 'calls rm-key with only one argument' do
+        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
+        expect(Gitlab::Utils).to receive(:system_silent).with(
+          [:gitlab_shell_keys_path, 'rm-key', 'key-123']
+        )
+
+        gitlab_shell.remove_key('key-123')
+      end
+    end
+  end
+
+  describe '#remove_all_keys' do
+    context 'when authorized_keys_enabled is true' do
+      it 'removes trailing garbage' do
+        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
+        expect(Gitlab::Utils).to receive(:system_silent).with([:gitlab_shell_keys_path, 'clear'])
+
+        gitlab_shell.remove_all_keys
+      end
+    end
+
+    context 'when authorized_keys_enabled is false' do
+      before do
+        stub_application_setting(authorized_keys_enabled: false)
+      end
+
+      it 'does nothing' do
+        expect(Gitlab::Utils).not_to receive(:system_silent)
+
+        gitlab_shell.remove_all_keys
+      end
+    end
+
+    context 'when authorized_keys_enabled is nil' do
+      before do
+        stub_application_setting(authorized_keys_enabled: nil)
+      end
+
+      it 'removes trailing garbage' do
+        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
+        expect(Gitlab::Utils).to receive(:system_silent).with([:gitlab_shell_keys_path, 'clear'])
+
+        gitlab_shell.remove_all_keys
+      end
+    end
+  end
+
+  describe '#remove_keys_not_found_in_db' do
+    context 'when keys are in the file that are not in the DB' do
+      before do
+        gitlab_shell.remove_all_keys
+        gitlab_shell.add_key('key-1234', 'ssh-rsa ASDFASDF')
+        gitlab_shell.add_key('key-9876', 'ssh-rsa ASDFASDF')
+        @another_key = create(:key) # this one IS in the DB
+      end
+
+      it 'removes the keys' do
+        expect(find_in_authorized_keys_file(1234)).to be_truthy
+        expect(find_in_authorized_keys_file(9876)).to be_truthy
+        expect(find_in_authorized_keys_file(@another_key.id)).to be_truthy
+        gitlab_shell.remove_keys_not_found_in_db
+        expect(find_in_authorized_keys_file(1234)).to be_falsey
+        expect(find_in_authorized_keys_file(9876)).to be_falsey
+        expect(find_in_authorized_keys_file(@another_key.id)).to be_truthy
+      end
+    end
+
+    context 'when keys there are duplicate keys in the file that are not in the DB' do
+      before do
+        gitlab_shell.remove_all_keys
+        gitlab_shell.add_key('key-1234', 'ssh-rsa ASDFASDF')
+        gitlab_shell.add_key('key-1234', 'ssh-rsa ASDFASDF')
+      end
+
+      it 'removes the keys' do
+        expect(find_in_authorized_keys_file(1234)).to be_truthy
+        gitlab_shell.remove_keys_not_found_in_db
+        expect(find_in_authorized_keys_file(1234)).to be_falsey
+      end
+
+      it 'does not run remove more than once per key (in a batch)' do
+        expect(gitlab_shell).to receive(:remove_key).with('key-1234').once
+        gitlab_shell.remove_keys_not_found_in_db
+      end
+    end
+
+    context 'when keys there are duplicate keys in the file that ARE in the DB' do
+      before do
+        gitlab_shell.remove_all_keys
+        @key = create(:key)
+        gitlab_shell.add_key(@key.shell_id, @key.key)
+      end
+
+      it 'does not remove the key' do
+        gitlab_shell.remove_keys_not_found_in_db
+        expect(find_in_authorized_keys_file(@key.id)).to be_truthy
+      end
+
+      it 'does not need to run a SELECT query for that batch, on account of that key' do
+        expect_any_instance_of(ActiveRecord::Relation).not_to receive(:pluck)
+        gitlab_shell.remove_keys_not_found_in_db
+      end
+    end
+
+    unless ENV['CI'] # Skip in CI, it takes 1 minute
+      context 'when the first batch can be skipped, but the next batch has keys that are not in the DB' do
+        before do
+          gitlab_shell.remove_all_keys
+          100.times { |i| create(:key) } # first batch is all in the DB
+          gitlab_shell.add_key('key-1234', 'ssh-rsa ASDFASDF')
+        end
+
+        it 'removes the keys not in the DB' do
+          expect(find_in_authorized_keys_file(1234)).to be_truthy
+          gitlab_shell.remove_keys_not_found_in_db
+          expect(find_in_authorized_keys_file(1234)).to be_falsey
+        end
+      end
+    end
+  end
+
+  describe '#batch_read_key_ids' do
+    context 'when there are keys in the authorized_keys file' do
+      before do
+        gitlab_shell.remove_all_keys
+        (1..4).each do |i|
+          gitlab_shell.add_key("key-#{i}", "ssh-rsa ASDFASDF#{i}")
+        end
+      end
+
+      it 'iterates over the key IDs in the file, in batches' do
+        loop_count = 0
+        first_batch = [1, 2]
+        second_batch = [3, 4]
+
+        gitlab_shell.batch_read_key_ids(batch_size: 2) do |batch|
+          expected = (loop_count == 0 ? first_batch : second_batch)
+          expect(batch).to eq(expected)
+          loop_count += 1
+        end
+      end
+    end
+  end
+
+  describe '#list_key_ids' do
+    context 'when there are keys in the authorized_keys file' do
+      before do
+        gitlab_shell.remove_all_keys
+        (1..4).each do |i|
+          gitlab_shell.add_key("key-#{i}", "ssh-rsa ASDFASDF#{i}")
+        end
+      end
+
+      it 'outputs the key IDs in the file, separated by newlines' do
+        ids = []
+        gitlab_shell.list_key_ids do |io|
+          io.each do |line|
+            ids << line
+          end
+        end
+
+        expect(ids).to eq(%W{1\n 2\n 3\n 4\n})
+      end
+    end
+
+    context 'when there are no keys in the authorized_keys file' do
+      before do
+        gitlab_shell.remove_all_keys
+      end
+
+      it 'outputs nothing, not even an empty string' do
+        ids = []
+        gitlab_shell.list_key_ids do |io|
+          io.each do |line|
+            ids << line
+          end
+        end
+
+        expect(ids).to eq([])
+      end
     end
   end
 
@@ -188,4 +480,12 @@ describe Gitlab::Shell, lib: true do
       end
     end
   end
+
+  def find_in_authorized_keys_file(key_id)
+    gitlab_shell.batch_read_key_ids do |ids|
+      return true if ids.include?(key_id)
+    end
+
+    false
+  end
 end

spec/migrations/update_authorized_keys_file_spec.rb

+require 'spec_helper'
+require Rails.root.join('db', 'migrate', '20170626202753_update_authorized_keys_file.rb')
+
+describe UpdateAuthorizedKeysFile, :migration do
+  let(:migration) { described_class.new }
+
+  describe '#up' do
+    context 'when authorized_keys_enabled is nil' do
+      before do
+        # Ensure the column can be null for the test
+        ActiveRecord::Base.connection.change_column_null :application_settings, :authorized_keys_enabled, true
+        ActiveRecord::Base.connection.change_column :application_settings, :authorized_keys_enabled, :boolean, default: nil
+
+        ApplicationSetting.create!(authorized_keys_enabled: nil)
+      end
+
+      it 'sets authorized_keys_enabled to true' do
+        migration.up
+
+        expect(ApplicationSetting.last.authorized_keys_enabled).to be_truthy
+      end
+
+      context 'there are keys created before and after the cutoff datetime' do
+        before do
+          Timecop.freeze
+        end
+        after do
+          Timecop.return
+        end
+
+        before do
+          @cutoff_datetime = UpdateAuthorizedKeysFile::DATETIME_9_3_0_RELEASED
+          @keys = []
+          Timecop.travel(@cutoff_datetime - 1.day)
+          2.times { @keys << create(:key) } # 2 keys before cutoff
+          Timecop.travel(@cutoff_datetime + 1.day)
+          2.times { @keys << create(:key) } # 2 keys after cutoff
+        end
+
+        it 'adds the keys created after the cutoff datetime to the authorized_keys file' do
+          Gitlab::Shell.new.remove_all_keys
+
+          migration.up
+
+          file = File.read(Rails.root.join('tmp/tests/.ssh/authorized_keys'))
+          expect(file.scan(/ssh-rsa/).count).to eq(2)
+
+          expect(file).not_to include(Gitlab::Shell.strip_key(@keys[0].key))
+          expect(file).not_to include(Gitlab::Shell.strip_key(@keys[1].key))
+          expect(file).to include(Gitlab::Shell.strip_key(@keys[2].key))
+          expect(file).to include(Gitlab::Shell.strip_key(@keys[3].key))
+        end
+      end
+
+      context 'when an SSH key exists in authorized_keys but not in the DB' do
+        before do
+          @key_to_stay = create(:key)
+          @key_to_delete = create(:key)
+          @key_to_delete.delete
+        end
+
+        it 'deletes the SSH key from authorized_keys' do
+          file = File.read(Rails.root.join('tmp/tests/.ssh/authorized_keys'))
+          expect(file).to include(Gitlab::Shell.strip_key(@key_to_stay.key))
+          expect(file).to include(Gitlab::Shell.strip_key(@key_to_delete.key))
+
+          migration.up
+
+          file = File.read(Rails.root.join('tmp/tests/.ssh/authorized_keys'))
+          expect(file).to include(Gitlab::Shell.strip_key(@key_to_stay.key))
+          expect(file).not_to include(Gitlab::Shell.strip_key(@key_to_delete.key))
+        end
+      end
+    end
+  end
+
+  describe '#authorized_keys_file_in_use_and_stale?' do
+    subject { migration.authorized_keys_file_in_use_and_stale? }
+
+    context 'when the customer ran the broken migration' do
+      before do
+        allow(migration).to receive(:ran_broken_migration?).and_return(true)
+      end
+
+      context 'when is a record in application_settings table' do
+        before do
+          ApplicationSetting.create!(authorized_keys_enabled: true)
+        end
+
+        context 'when authorized_keys_enabled is true' do
+          it { is_expected.to be_truthy }
+        end
+
+        context 'when authorized_keys_enabled is nil' do
+          before do
+            # Ensure the column can be null for the test
+            ActiveRecord::Base.connection.change_column_null :application_settings, :authorized_keys_enabled, true
+            ActiveRecord::Base.connection.change_column :application_settings, :authorized_keys_enabled, :boolean, default: nil
+
+            ApplicationSetting.first.update(authorized_keys_enabled: nil)
+          end
+
+          it { is_expected.to be_truthy }
+        end
+
+        context 'when authorized_keys_enabled is explicitly false' do
+          before do
+            ApplicationSetting.first.update!(authorized_keys_enabled: false)
+          end
+
+          it { is_expected.to be_falsey }
+
+          it 'outputs a warning message for users who unintentionally Saved the setting unchecked' do
+            expect{ subject }.to output(/warning.*intentionally/mi).to_stdout
+          end
+        end
+      end
+
+      context 'when there is no record in application_settings table' do
+        before do
+          expect(ApplicationSetting.count).to eq(0)
+        end
+
+        it { is_expected.to be_falsey }
+      end
+    end
+
+    context 'when the customer did not run the broken migration' do
+      before do
+        allow(migration).to receive(:ran_broken_migration?).and_return(false)
+      end
+
+      it { is_expected.to be_falsey }
+    end
+  end
+
+  describe '#ran_broken_migration?' do
+    subject { migration.ran_broken_migration? }
+
+    context 'for unaffected customers: the authorized_keys_enabled column has a default (so the fixed migration ran)' do
+      before do
+        ActiveRecord::Base.connection.change_column :application_settings, :authorized_keys_enabled, :boolean, default: true
+        ActiveRecord::Base.connection.change_column_null :application_settings, :authorized_keys_enabled, false, true
+      end
+
+      it 'returns false' do
+        expect(subject).to be_falsey
+      end
+    end
+
+    context 'for affected customers: the authorized_keys_enabled column does not have a default (so the broken migration ran)' do
+      before do
+        ActiveRecord::Base.connection.change_column_null :application_settings, :authorized_keys_enabled, true
+        ActiveRecord::Base.connection.change_column :application_settings, :authorized_keys_enabled, :boolean, default: nil
+      end
+
+      it 'returns true' do
+        expect(subject).to be_truthy
+      end
+    end
+  end
+end

spec/workers/gitlab_shell_worker_spec.rb

+require 'spec_helper'
+
+describe GitlabShellWorker do
+  let(:worker) { described_class.new }
+
+  describe '#perform with add_key' do
+    it 'calls add_key on Gitlab::Shell' do
+      expect_any_instance_of(Gitlab::Shell).to receive(:add_key).with('foo', 'bar')
+      worker.perform(:add_key, 'foo', 'bar')
+    end
+  end
+end