Backport authorized_keys_enabled defaults to true'

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/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|

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

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

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