Backport authorized_keys_enabled defaults to true'

Originally from branch 'fix-authorized-keys-enabled-default-2738' via merge request https://gitlab.com/gitlab-org/gitlab-ee/merge_requests/2240

Removed background migrations which were intended to fix state after using Gitlab
without a default having been set

Squashed commits:
Locally, if Spring was not restarted, `current_application_settings` was still cached, which prevented the migration from editing the file. This will also ensure that any app server somehow hitting old cache data will properly default this setting regardless.
Retroactively fix migration
  This allows us to identify customers who ran the broken migration. Their `authorized_keys_enabled` column does not have a default at this point.
  We will fix the column after we fix the `authorized_keys` file.
Fix authorized_keys file if needed
Add default to authorized_keys_enabled setting
  Reminder: The original migration was fixed retroactively a few commits ago, so people who did not ever run GitLab 9.3.0 already have a column that defaults to true and disallows nulls. I have tested on PostgreSQL and MySQL that it is safe to run this migration regardless.
  Affected customers who did run 9.3.0 are the ones who need this migration to fix the authorized_keys_enabled column.
  The reason for the retroactive fix plus this migration is that it allows us to run a migration in between to fix the authorized_keys file only for those who ran 9.3.0.
Tweaks to address feedback
Extract work into background migration
Move batch-add-logic to background migration
  Do the work synchronously to avoid multiple workers attempting to add batches of keys at the same time.
  Also, make the delete portion wait until after adding is done.
Do read and delete work in background migration
Fix Rubocop offenses
Add changelog entry
Inform the user of actions taken or not taken
Prevent unnecessary `select`s and `remove_key`s
Add logs for action taken
Fix optimization
Reuse `Gitlab::ShellAdapter`
Guarantee the earliest key
Fix migration spec for MySQL

5 files changed, 289 insertions, 7 deletions

diff --git a/db/migrate/20170531180233_add_authorized_keys_enabled_to_application_settings.rb b/db/migrate/20170531180233_add_authorized_keys_enabled_to_application_settings.rb
index fdae309946c..1d86a531eb3 100644
--- a/db/migrate/20170531180233_add_authorized_keys_enabled_to_application_settings.rb
+++ b/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
diff --git a/db/schema.rb b/db/schema.rb
index 11273d2a82e..840abffd9f4 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -154,7 +154,7 @@ ActiveRecord::Schema.define(version: 20171230123729) do
     t.integer "gitaly_timeout_default", default: 55, null: false
     t.integer "gitaly_timeout_medium", default: 30, null: false
     t.integer "gitaly_timeout_fast", default: 10, null: false
-    t.boolean "authorized_keys_enabled"
+    t.boolean "authorized_keys_enabled", default: true, null: false
   end
 
   create_table "audit_events", force: :cascade do |t|
diff --git a/lib/gitlab/shell.rb b/lib/gitlab/shell.rb
index 20df19c734d..65a9e58adf1 100644
--- a/lib/gitlab/shell.rb
+++ b/lib/gitlab/shell.rb
@@ -206,12 +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
 
@@ -226,6 +225,57 @@ module Gitlab
       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.
@@ -420,6 +470,10 @@ module Gitlab
     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
diff --git a/spec/lib/gitlab/shell_spec.rb b/spec/lib/gitlab/shell_spec.rb
index 6d50f537430..e452fbb757d 100644
--- a/spec/lib/gitlab/shell_spec.rb
+++ b/spec/lib/gitlab/shell_spec.rb
@@ -74,6 +74,21 @@ describe Gitlab::Shell do
         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
+
+      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
@@ -100,6 +115,20 @@ describe Gitlab::Shell do
         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
@@ -125,6 +154,32 @@ describe Gitlab::Shell do
         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
@@ -148,6 +203,155 @@ describe Gitlab::Shell do
         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
 
   describe Gitlab::Shell::KeyAdder do
@@ -484,4 +688,12 @@ describe Gitlab::Shell 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
diff --git a/spec/workers/gitlab_shell_worker_spec.rb b/spec/workers/gitlab_shell_worker_spec.rb
new file mode 100644
index 00000000000..6b222af454d
--- /dev/null
+++ b/spec/workers/gitlab_shell_worker_spec.rb
@@ -0,0 +1,12 @@
+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