Integrate Gitlab::Keys with Gitlab::Shell

In this commit, some methods that aren't being used
are removed from `Gitlab::Shell`. They are the ff:
- `#remove_keys_not_found_in_db`
- `#batch_read_key_ids`
- `#list_key_ids`

The corresponding methods in `Gitlab::Keys` have been
removed as well.

changelogs/unreleased/56089-merge-gitlab-keys.yml

+---
+title: Merge the gitlab-shell "gitlab-keys" functionality into GitLab CE
+merge_request: 25598
+author:
+type: other

config/gitlab.yml.example

@@ -697,6 +697,7 @@ production: &base
   ## GitLab Shell settings
   gitlab_shell:
     path: /home/git/gitlab-shell/
+    authorized_keys_file: /home/git/.ssh/authorized_keys
 
     # File that contains the secret key for verifying access for gitlab-shell.
     # Default is '.gitlab_shell_secret' relative to Rails.root (i.e. root of the GitLab app).
@@ -854,6 +855,7 @@ test:
     path: tmp/tests/backups
   gitlab_shell:
     path: tmp/tests/gitlab-shell/
+    authorized_keys_file: tmp/tests/authorized_keys
   issues_tracker:
     redmine:
       title: "Redmine"

config/initializers/1_settings.rb

@@ -356,6 +356,7 @@ Settings['sidekiq']['log_format'] ||= 'default'
 Settings['gitlab_shell'] ||= Settingslogic.new({})
 Settings.gitlab_shell['path']           = Settings.absolute(Settings.gitlab_shell['path'] || Settings.gitlab['user_home'] + '/gitlab-shell/')
 Settings.gitlab_shell['hooks_path']     = :deprecated_use_gitlab_shell_path_instead
+Settings.gitlab_shell['authorized_keys_file'] ||= nil
 Settings.gitlab_shell['secret_file'] ||= Rails.root.join('.gitlab_shell_secret')
 Settings.gitlab_shell['receive_pack']   = true if Settings.gitlab_shell['receive_pack'].nil?
 Settings.gitlab_shell['upload_pack']    = true if Settings.gitlab_shell['upload_pack'].nil?

lib/gitlab/authorized_keys.rb

+# frozen_string_literal: true
+
+module Gitlab
+  class AuthorizedKeys
+    KeyError = Class.new(StandardError)
+
+    attr_reader :logger
+
+    # Initializes the class
+    #
+    # @param [Gitlab::Logger] logger
+    def initialize(logger = Gitlab::AppLogger)
+      @logger = logger
+    end
+
+    # Add id and its key to the authorized_keys file
+    #
+    # @param [String] id identifier of key prefixed by `key-`
+    # @param [String] key public key to be added
+    # @return [Boolean]
+    def add_key(id, key)
+      lock do
+        public_key = strip(key)
+        logger.info("Adding key (#{id}): #{public_key}")
+        open_authorized_keys_file('a') { |file| file.puts(key_line(id, public_key)) }
+      end
+
+      true
+    end
+
+    # Atomically add all the keys to the authorized_keys file
+    #
+    # @param [Array<::Key>] keys list of Key objects to be added
+    # @return [Boolean]
+    def batch_add_keys(keys)
+      lock(300) do # Allow 300 seconds (5 minutes) for batch_add_keys
+        open_authorized_keys_file('a') do |file|
+          keys.each do |key|
+            public_key = strip(key.key)
+            logger.info("Adding key (#{key.shell_id}): #{public_key}")
+            file.puts(key_line(key.shell_id, public_key))
+          end
+        end
+      end
+
+      true
+    rescue Gitlab::AuthorizedKeys::KeyError
+      false
+    end
+
+    # Remove key by ID from the authorized_keys file
+    #
+    # @param [String] id identifier of the key to be removed prefixed by `key-`
+    # @return [Boolean]
+    def rm_key(id)
+      lock do
+        logger.info("Removing key (#{id})")
+        open_authorized_keys_file('r+') do |f|
+          while line = f.gets
+            next unless line.start_with?("command=\"#{command(id)}\"")
+
+            f.seek(-line.length, IO::SEEK_CUR)
+            # Overwrite the line with #'s. Because the 'line' variable contains
+            # a terminating '\n', we write line.length - 1 '#' characters.
+            f.write('#' * (line.length - 1))
+          end
+        end
+      end
+
+      true
+    end
+
+    # Clear the authorized_keys file
+    #
+    # @return [Boolean]
+    def clear
+      open_authorized_keys_file('w') { |file| file.puts '# Managed by gitlab-rails' }
+
+      true
+    end
+
+    # Read the authorized_keys file and return IDs of each key
+    #
+    # @return [Array<Integer>]
+    def list_key_ids
+      logger.info('Listing all key IDs')
+
+      [].tap do |a|
+        open_authorized_keys_file('r') do |f|
+          f.each_line do |line|
+            key_id = line.match(/key-(\d+)/)
+
+            next unless key_id
+
+            a << key_id[1].chomp.to_i
+          end
+        end
+      end
+    end
+
+    private
+
+    def lock(timeout = 10)
+      File.open("#{authorized_keys_file}.lock", "w+") do |f|
+        f.flock File::LOCK_EX
+        Timeout.timeout(timeout) { yield }
+      ensure
+        f.flock File::LOCK_UN
+      end
+    end
+
+    def open_authorized_keys_file(mode)
+      File.open(authorized_keys_file, mode, 0o600) do |file|
+        file.chmod(0o600)
+        yield file
+      end
+    end
+
+    def key_line(id, key)
+      key = key.chomp
+
+      if key.include?("\n") || key.include?("\t")
+        raise KeyError, "Invalid public_key: #{key.inspect}"
+      end
+
+      %Q(command="#{command(id)}",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty #{strip(key)})
+    end
+
+    def command(id)
+      unless /\A[a-z0-9-]+\z/ =~ id
+        raise KeyError, "Invalid ID: #{id.inspect}"
+      end
+
+      "#{File.join(Gitlab.config.gitlab_shell.path, 'bin', 'gitlab-shell')} #{id}"
+    end
+
+    def strip(key)
+      key.split(/[ ]+/)[0, 2].join(' ')
+    end
+
+    def authorized_keys_file
+      Gitlab.config.gitlab_shell.authorized_keys_file
+    end
+  end
+end

lib/gitlab/shell.rb

@@ -10,18 +10,6 @@ module Gitlab
 
     Error = Class.new(StandardError)
 
-    KeyAdder = Struct.new(:io) do
-      def add_key(id, key)
-        key = Gitlab::Shell.strip_key(key)
-        # Newline and tab are part of the 'protocol' used to transmit id+key to the other end
-        if key.include?("\t") || key.include?("\n")
-          raise Error.new("Invalid key: #{key.inspect}")
-        end
-
-        io.puts("#{id}\t#{key}")
-      end
-    end
-
     class << self
       def secret_token
         @secret_token ||= begin
@@ -40,10 +28,6 @@ module Gitlab
                                         .join('GITLAB_SHELL_VERSION')).strip
       end
 
-      def strip_key(key)
-        key.split(/[ ]+/)[0, 2].join(' ')
-      end
-
       private
 
       # Create (if necessary) and link the secret token file
@@ -173,7 +157,7 @@ module Gitlab
       false
     end
 
-    # Add new key to gitlab-shell
+    # Add new key to authorized_keys
     #
     # Ex.
     #   add_key("key-42", "sha-rsa ...")
@@ -181,33 +165,53 @@ module Gitlab
     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)])
+      if shell_out_for_gitlab_keys?
+        gitlab_shell_fast_execute([
+          gitlab_shell_keys_path,
+          'add-key',
+          key_id,
+          strip_key(key_content)
+        ])
+      else
+        gitlab_authorized_keys.add_key(key_id, key_content)
+      end
     end
 
     # Batch-add keys to authorized_keys
     #
     # Ex.
-    #   batch_add_keys { |adder| adder.add_key("key-42", "sha-rsa ...") }
-    def batch_add_keys(&block)
+    #   batch_add_keys(Key.all)
+    def batch_add_keys(keys)
       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))
+      if shell_out_for_gitlab_keys?
+        begin
+          IO.popen("#{gitlab_shell_keys_path} batch-add-keys", 'w') do |io|
+            add_keys_to_io(keys, io)
+          end
+
+          $?.success?
+        rescue Error
+          false
+        end
+      else
+        gitlab_authorized_keys.batch_add_keys(keys)
       end
     end
 
-    # Remove ssh key from gitlab shell
+    # Remove ssh key from authorized_keys
     #
     # Ex.
-    #   remove_key("key-342", "sha-rsa ...")
+    #   remove_key("key-342")
     #
-    def remove_key(key_id, key_content = nil)
+    def remove_key(id, _ = 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)
+      if shell_out_for_gitlab_keys?
+        gitlab_shell_fast_execute([gitlab_shell_keys_path, 'rm-key', id])
+      else
+        gitlab_authorized_keys.rm_key(id)
+      end
     end
 
     # Remove all ssh keys from gitlab shell
@@ -218,7 +222,11 @@ module Gitlab
     def remove_all_keys
       return unless self.authorized_keys_enabled?
 
-      gitlab_shell_fast_execute([gitlab_shell_keys_path, 'clear'])
+      if shell_out_for_gitlab_keys?
+        gitlab_shell_fast_execute([gitlab_shell_keys_path, 'clear'])
+      else
+        gitlab_authorized_keys.clear
+      end
     end
 
     # Remove ssh keys from gitlab shell that are not in the DB
@@ -247,33 +255,6 @@ module Gitlab
     end
     # rubocop: enable CodeReuse/ActiveRecord
 
-    # 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.
@@ -378,6 +359,10 @@ module Gitlab
 
     private
 
+    def shell_out_for_gitlab_keys?
+      Gitlab.config.gitlab_shell.authorized_keys_file.blank?
+    end
+
     def gitlab_shell_fast_execute(cmd)
       output, status = gitlab_shell_fast_execute_helper(cmd)
 
@@ -415,6 +400,40 @@ module Gitlab
       raise Error, e
     end
 
+    def gitlab_authorized_keys
+      @gitlab_authorized_keys ||= Gitlab::AuthorizedKeys.new
+    end
+
+    def batch_read_key_ids(batch_size: 100, &block)
+      return unless self.authorized_keys_enabled?
+
+      if shell_out_for_gitlab_keys?
+        IO.popen("#{gitlab_shell_keys_path} list-key-ids") do |key_id_stream|
+          key_id_stream.lazy.each_slice(batch_size) do |lines|
+            yield(lines.map { |l| l.chomp.to_i })
+          end
+        end
+      else
+        gitlab_authorized_keys.list_key_ids.lazy.each_slice(batch_size) do |key_ids|
+          yield(key_ids)
+        end
+      end
+    end
+
+    def strip_key(key)
+      key.split(/[ ]+/)[0, 2].join(' ')
+    end
+
+    def add_keys_to_io(keys, io)
+      keys.each do |k|
+        key = strip_key(k.key)
+
+        raise Error.new("Invalid key: #{key.inspect}") if key.include?("\t") || key.include?("\n")
+
+        io.puts("#{k.shell_id}\t#{key}")
+      end
+    end
+
     class GitalyGitlabProjects
       attr_reader :shard_name, :repository_relative_path, :output, :gl_project_path
 

lib/tasks/gitlab/shell.rake

@@ -103,19 +103,12 @@ namespace :gitlab do
 
     Gitlab::Shell.new.remove_all_keys
 
-    Gitlab::Shell.new.batch_add_keys do |adder|
-      Key.find_each(batch_size: 1000) do |key|
-        adder.add_key(key.shell_id, key.key)
-        print '.'
+    Key.find_in_batches(batch_size: 1000) do |keys|
+      unless Gitlab::Shell.new.batch_add_keys(keys)
+        puts "Failed to add keys...".color(:red)
+        exit 1
       end
     end
-    puts ""
-
-    unless $?.success?
-      puts "Failed to add keys...".color(:red)
-      exit 1
-    end
-
   rescue Gitlab::TaskAbortedByUserError
     puts "Quitting...".color(:red)
     exit 1

spec/lib/gitlab/authorized_keys_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::AuthorizedKeys do
+  let(:logger) { double('logger').as_null_object }
+
+  subject { described_class.new(logger) }
+
+  describe '#add_key' do
+    it "adds a line at the end of the file and strips trailing garbage" do
+      create_authorized_keys_fixture
+      auth_line = "command=\"#{Gitlab.config.gitlab_shell.path}/bin/gitlab-shell key-741\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty ssh-rsa AAAAB3NzaDAxx2E"
+
+      expect(logger).to receive(:info).with('Adding key (key-741): ssh-rsa AAAAB3NzaDAxx2E')
+      expect(subject.add_key('key-741', 'ssh-rsa AAAAB3NzaDAxx2E trailing garbage'))
+        .to be_truthy
+      expect(File.read(tmp_authorized_keys_path)).to eq("existing content\n#{auth_line}\n")
+    end
+  end
+
+  describe '#batch_add_keys' do
+    let(:keys) do
+      [
+        double(shell_id: 'key-12', key: 'ssh-dsa ASDFASGADG trailing garbage'),
+        double(shell_id: 'key-123', key: 'ssh-rsa GFDGDFSGSDFG')
+      ]
+    end
+
+    before do
+      create_authorized_keys_fixture
+    end
+
+    it "adds lines at the end of the file" do
+      auth_line1 = "command=\"#{Gitlab.config.gitlab_shell.path}/bin/gitlab-shell key-12\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty ssh-dsa ASDFASGADG"
+      auth_line2 = "command=\"#{Gitlab.config.gitlab_shell.path}/bin/gitlab-shell key-123\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty ssh-rsa GFDGDFSGSDFG"
+
+      expect(logger).to receive(:info).with('Adding key (key-12): ssh-dsa ASDFASGADG')
+      expect(logger).to receive(:info).with('Adding key (key-123): ssh-rsa GFDGDFSGSDFG')
+      expect(subject.batch_add_keys(keys)).to be_truthy
+      expect(File.read(tmp_authorized_keys_path)).to eq("existing content\n#{auth_line1}\n#{auth_line2}\n")
+    end
+
+    context "invalid key" do
+      let(:keys) { [double(shell_id: 'key-123', key: "ssh-rsa A\tSDFA\nSGADG")] }
+
+      it "doesn't add keys" do
+        expect(subject.batch_add_keys(keys)).to be_falsey
+        expect(File.read(tmp_authorized_keys_path)).to eq("existing content\n")
+      end
+    end
+  end
+
+  describe '#rm_key' do
+    it "removes the right line" do
+      create_authorized_keys_fixture
+      other_line = "command=\"#{Gitlab.config.gitlab_shell.path}/bin/gitlab-shell key-742\",options ssh-rsa AAAAB3NzaDAxx2E"
+      delete_line = "command=\"#{Gitlab.config.gitlab_shell.path}/bin/gitlab-shell key-741\",options ssh-rsa AAAAB3NzaDAxx2E"
+      erased_line = delete_line.gsub(/./, '#')
+      File.open(tmp_authorized_keys_path, 'a') do |auth_file|
+        auth_file.puts delete_line
+        auth_file.puts other_line
+      end
+
+      expect(logger).to receive(:info).with('Removing key (key-741)')
+      expect(subject.rm_key('key-741')).to be_truthy
+      expect(File.read(tmp_authorized_keys_path)).to eq("existing content\n#{erased_line}\n#{other_line}\n")
+    end
+  end
+
+  describe '#clear' do
+    it "should return true" do
+      expect(subject.clear).to be_truthy
+    end
+  end
+
+  describe '#list_key_ids' do
+    before do
+      create_authorized_keys_fixture(
+        existing_content:
+          "key-1\tssh-dsa AAA\nkey-2\tssh-rsa BBB\nkey-3\tssh-rsa CCC\nkey-9000\tssh-rsa DDD\n"
+      )
+    end
+
+    it 'returns array of key IDs' do
+      expect(subject.list_key_ids).to eq([1, 2, 3, 9000])
+    end
+  end
+
+  def create_authorized_keys_fixture(existing_content: 'existing content')
+    FileUtils.mkdir_p(File.dirname(tmp_authorized_keys_path))
+    File.open(tmp_authorized_keys_path, 'w') { |file| file.puts(existing_content) }
+  end
+
+  def tmp_authorized_keys_path
+    Gitlab.config.gitlab_shell.authorized_keys_file
+  end
+end

spec/lib/gitlab/shell_spec.rb

@@ -8,6 +8,7 @@ describe Gitlab::Shell do
   let(:gitlab_shell) { described_class.new }
   let(:popen_vars) { { 'GIT_TERMINAL_PROMPT' => ENV['GIT_TERMINAL_PROMPT'] } }
   let(:timeout) { Gitlab.config.gitlab_shell.git_timeout }
+  let(:gitlab_authorized_keys) { double }
 
   before do
     allow(Project).to receive(:find).and_return(project)
@@ -49,13 +50,38 @@ describe Gitlab::Shell do
 
   describe '#add_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_shell).to receive(:gitlab_shell_fast_execute).with(
-          [:gitlab_shell_keys_path, 'add-key', 'key-123', 'ssh-rsa foobar']
-        )
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          allow(gitlab_shell)
+            .to receive(:gitlab_shell_keys_path)
+            .and_return(:gitlab_shell_keys_path)
+        end
+
+        it 'calls #gitlab_shell_fast_execute with add-key command' do
+          expect(gitlab_shell)
+            .to receive(:gitlab_shell_fast_execute)
+            .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 'authorized_keys_file set' do
+        it 'calls Gitlab::AuthorizedKeys#add_key with id and key' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
 
-        gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+          expect(gitlab_authorized_keys)
+            .to receive(:add_key)
+            .with('key-123', 'ssh-rsa foobar')
+
+          gitlab_shell.add_key('key-123', 'ssh-rsa foobar')
+        end
       end
     end
 
@@ -64,10 +90,24 @@ describe Gitlab::Shell do
         stub_application_setting(authorized_keys_enabled: false)
       end
 
-      it 'does nothing' do
-        expect(gitlab_shell).not_to receive(:gitlab_shell_fast_execute)
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+        end
+
+        it 'does nothing' do
+          expect(gitlab_shell).not_to receive(:gitlab_shell_fast_execute)
+
+          gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+        end
+      end
+
+      context 'authorized_keys_file set' do
+        it 'does nothing' do
+          expect(Gitlab::AuthorizedKeys).not_to receive(:new)
 
-        gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+          gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+        end
       end
     end
 
@@ -76,24 +116,89 @@ describe Gitlab::Shell 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_shell).to receive(:gitlab_shell_fast_execute).with(
-          [:gitlab_shell_keys_path, 'add-key', 'key-123', 'ssh-rsa foobar']
-        )
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          allow(gitlab_shell)
+            .to receive(:gitlab_shell_keys_path)
+            .and_return(:gitlab_shell_keys_path)
+        end
+
+        it 'calls #gitlab_shell_fast_execute with add-key command' do
+          expect(gitlab_shell)
+            .to receive(:gitlab_shell_fast_execute)
+            .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 'authorized_keys_file set' do
+        it 'calls Gitlab::AuthorizedKeys#add_key with id and key' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
+
+          expect(gitlab_authorized_keys)
+            .to receive(:add_key)
+            .with('key-123', 'ssh-rsa foobar')
 
-        gitlab_shell.add_key('key-123', 'ssh-rsa foobar trailing garbage')
+          gitlab_shell.add_key('key-123', 'ssh-rsa foobar')
+        end
       end
     end
   end
 
   describe '#batch_add_keys' do
+    let(:keys) { [double(shell_id: 'key-123', key: 'ssh-rsa foobar')] }
+
     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')
+      context 'authorized_keys_file not set' do
+        let(:io) { double }
+
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+        end
+
+        context 'valid keys' do
+          before do
+            allow(gitlab_shell)
+              .to receive(:gitlab_shell_keys_path)
+              .and_return(:gitlab_shell_keys_path)
+          end
+
+          it 'calls gitlab-keys with batch-add-keys command' do
+            expect(IO)
+              .to receive(:popen)
+              .with("gitlab_shell_keys_path batch-add-keys", 'w')
+              .and_yield(io)
+
+            expect(io).to receive(:puts).with("key-123\tssh-rsa foobar")
+            expect(gitlab_shell.batch_add_keys(keys)).to be_truthy
+          end
+        end
+
+        context 'invalid keys' do
+          let(:keys) { [double(shell_id: 'key-123', key: "ssh-rsa A\tSDFA\nSGADG")] }
+
+          it 'catches failure and returns false' do
+            expect(gitlab_shell.batch_add_keys(keys)).to be_falsey
+          end
+        end
+      end
 
-        gitlab_shell.batch_add_keys do |adder|
-          adder.add_key('key-123', 'ssh-rsa foobar')
+      context 'authorized_keys_file set' do
+        it 'calls Gitlab::AuthorizedKeys#batch_add_keys with keys to be added' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
+
+          expect(gitlab_authorized_keys)
+            .to receive(:batch_add_keys)
+            .with(keys)
+
+          gitlab_shell.batch_add_keys(keys)
         end
       end
     end
@@ -103,11 +208,23 @@ describe Gitlab::Shell 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)
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+        end
+
+        it 'does nothing' do
+          expect(IO).not_to receive(:popen)
+
+          gitlab_shell.batch_add_keys(keys)
+        end
+      end
+
+      context 'authorized_keys_file set' do
+        it 'does nothing' do
+          expect(Gitlab::AuthorizedKeys).not_to receive(:new)
 
-        gitlab_shell.batch_add_keys do |adder|
-          adder.add_key('key-123', 'ssh-rsa foobar')
+          gitlab_shell.batch_add_keys(keys)
         end
       end
     end
@@ -117,11 +234,37 @@ describe Gitlab::Shell 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')
+      context 'authorized_keys_file not set' do
+        let(:io) { double }
 
-        gitlab_shell.batch_add_keys do |adder|
-          adder.add_key('key-123', 'ssh-rsa foobar')
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          allow(gitlab_shell)
+            .to receive(:gitlab_shell_keys_path)
+            .and_return(:gitlab_shell_keys_path)
+        end
+
+        it 'calls gitlab-keys with batch-add-keys command' do
+          expect(IO)
+            .to receive(:popen)
+            .with("gitlab_shell_keys_path batch-add-keys", 'w')
+            .and_yield(io)
+
+          expect(io).to receive(:puts).with("key-123\tssh-rsa foobar")
+
+          gitlab_shell.batch_add_keys(keys)
+        end
+      end
+
+      context 'authorized_keys_file set' do
+        it 'calls Gitlab::AuthorizedKeys#batch_add_keys with keys to be added' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
+
+          expect(gitlab_authorized_keys)
+            .to receive(:batch_add_keys)
+            .with(keys)
+
+          gitlab_shell.batch_add_keys(keys)
         end
       end
     end
@@ -129,13 +272,34 @@ describe Gitlab::Shell do
 
   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_shell).to receive(:gitlab_shell_fast_execute).with(
-          [:gitlab_shell_keys_path, 'rm-key', 'key-123', 'ssh-rsa foobar']
-        )
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          allow(gitlab_shell)
+            .to receive(:gitlab_shell_keys_path)
+            .and_return(:gitlab_shell_keys_path)
+        end
 
-        gitlab_shell.remove_key('key-123', 'ssh-rsa foobar')
+        it 'calls #gitlab_shell_fast_execute with rm-key command' do
+          expect(gitlab_shell)
+            .to receive(:gitlab_shell_fast_execute)
+            .with([
+              :gitlab_shell_keys_path,
+              'rm-key',
+              'key-123'
+            ])
+
+          gitlab_shell.remove_key('key-123')
+        end
+      end
+
+      context 'authorized_keys_file not set' do
+        it 'calls Gitlab::AuthorizedKeys#rm_key with the key to be removed' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
+          expect(gitlab_authorized_keys).to receive(:rm_key).with('key-123')
+
+          gitlab_shell.remove_key('key-123')
+        end
       end
     end
 
@@ -144,10 +308,24 @@ describe Gitlab::Shell do
         stub_application_setting(authorized_keys_enabled: false)
       end
 
-      it 'does nothing' do
-        expect(gitlab_shell).not_to receive(:gitlab_shell_fast_execute)
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+        end
+
+        it 'does nothing' do
+          expect(gitlab_shell).not_to receive(:gitlab_shell_fast_execute)
 
-        gitlab_shell.remove_key('key-123', 'ssh-rsa foobar')
+          gitlab_shell.remove_key('key-123')
+        end
+      end
+
+      context 'authorized_keys_file set' do
+        it 'does nothing' do
+          expect(Gitlab::AuthorizedKeys).not_to receive(:new)
+
+          gitlab_shell.remove_key('key-123')
+        end
       end
     end
 
@@ -156,232 +334,256 @@ describe Gitlab::Shell 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_shell).to receive(:gitlab_shell_fast_execute).with(
-          [:gitlab_shell_keys_path, 'rm-key', 'key-123', 'ssh-rsa foobar']
-        )
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          allow(gitlab_shell)
+            .to receive(:gitlab_shell_keys_path)
+            .and_return(:gitlab_shell_keys_path)
+        end
+
+        it 'calls #gitlab_shell_fast_execute with rm-key command' do
+          expect(gitlab_shell)
+            .to receive(:gitlab_shell_fast_execute)
+            .with([
+              :gitlab_shell_keys_path,
+              'rm-key',
+              'key-123'
+            ])
 
-        gitlab_shell.remove_key('key-123', 'ssh-rsa foobar')
+          gitlab_shell.remove_key('key-123')
+        end
       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_shell).to receive(:gitlab_shell_fast_execute).with(
-          [:gitlab_shell_keys_path, 'rm-key', 'key-123']
-        )
+      context 'authorized_keys_file not set' do
+        it 'calls Gitlab::AuthorizedKeys#rm_key with the key to be removed' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
+          expect(gitlab_authorized_keys).to receive(:rm_key).with('key-123')
 
-        gitlab_shell.remove_key('key-123')
+          gitlab_shell.remove_key('key-123')
+        end
       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_shell).to receive(:gitlab_shell_fast_execute).with([:gitlab_shell_keys_path, 'clear'])
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          allow(gitlab_shell)
+            .to receive(:gitlab_shell_keys_path)
+            .and_return(:gitlab_shell_keys_path)
+        end
 
-        gitlab_shell.remove_all_keys
-      end
-    end
+        it 'calls #gitlab_shell_fast_execute with clear command' do
+          expect(gitlab_shell)
+            .to receive(:gitlab_shell_fast_execute)
+            .with([:gitlab_shell_keys_path, 'clear'])
 
-    context 'when authorized_keys_enabled is false' do
-      before do
-        stub_application_setting(authorized_keys_enabled: false)
+          gitlab_shell.remove_all_keys
+        end
       end
 
-      it 'does nothing' do
-        expect(gitlab_shell).not_to receive(:gitlab_shell_fast_execute)
+      context 'authorized_keys_file set' do
+        it 'calls Gitlab::AuthorizedKeys#clear' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
+          expect(gitlab_authorized_keys).to receive(:clear)
 
-        gitlab_shell.remove_all_keys
+          gitlab_shell.remove_all_keys
+        end
       end
     end
 
-    context 'when authorized_keys_enabled is nil' do
+    context 'when authorized_keys_enabled is false' do
       before do
-        stub_application_setting(authorized_keys_enabled: nil)
+        stub_application_setting(authorized_keys_enabled: false)
       end
 
-      it 'removes trailing garbage' do
-        allow(gitlab_shell).to receive(:gitlab_shell_keys_path).and_return(:gitlab_shell_keys_path)
-        expect(gitlab_shell).to receive(:gitlab_shell_fast_execute).with(
-          [:gitlab_shell_keys_path, 'clear']
-        )
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+        end
 
-        gitlab_shell.remove_all_keys
-      end
-    end
-  end
+        it 'does nothing' do
+          expect(gitlab_shell).not_to receive(:gitlab_shell_fast_execute)
 
-  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
+          gitlab_shell.remove_all_keys
+        end
       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
+      context 'authorized_keys_file set' do
+        it 'does nothing' do
+          expect(Gitlab::AuthorizedKeys).not_to receive(:new)
+
+          gitlab_shell.remove_all_keys
+        end
       end
     end
 
-    context 'when keys there are duplicate keys in the file that are not in the DB' do
+    context 'when authorized_keys_enabled is nil' 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')
+        stub_application_setting(authorized_keys_enabled: nil)
       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
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          allow(gitlab_shell)
+            .to receive(:gitlab_shell_keys_path)
+            .and_return(:gitlab_shell_keys_path)
+        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
+        it 'calls #gitlab_shell_fast_execute with clear command' do
+          expect(gitlab_shell)
+            .to receive(:gitlab_shell_fast_execute)
+            .with([:gitlab_shell_keys_path, 'clear'])
 
-    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)
+          gitlab_shell.remove_all_keys
+        end
       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
+      context 'authorized_keys_file set' do
+        it 'calls Gitlab::AuthorizedKeys#clear' do
+          expect(Gitlab::AuthorizedKeys).to receive(:new).and_return(gitlab_authorized_keys)
+          expect(gitlab_authorized_keys).to receive(:clear)
 
-      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
+          gitlab_shell.remove_all_keys
+        end
       end
     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
+  describe '#remove_keys_not_found_in_db' do
+    context 'when keys are in the file that are not in the DB' do
+      context 'authorized_keys_file not set' do
         before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
           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')
+          gitlab_shell.add_key('key-9876', 'ssh-rsa ASDFASDF')
+          @another_key = create(:key) # this one IS in the DB
         end
 
-        it 'removes the keys not in the DB' do
-          expect(find_in_authorized_keys_file(1234)).to be_truthy
+        it 'removes the keys' do
+          expect(gitlab_shell).to receive(:remove_key).with('key-1234')
+          expect(gitlab_shell).to receive(:remove_key).with('key-9876')
+          expect(gitlab_shell).not_to receive(:remove_key).with("key-#{@another_key.id}")
+
           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}")
+      context 'authorized_keys_file set' 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
-      end
 
-      it 'iterates over the key IDs in the file, in batches' do
-        loop_count = 0
-        first_batch = [1, 2]
-        second_batch = [3, 4]
+        it 'removes the keys' do
+          expect(gitlab_shell).to receive(:remove_key).with('key-1234')
+          expect(gitlab_shell).to receive(:remove_key).with('key-9876')
+          expect(gitlab_shell).not_to receive(:remove_key).with("key-#{@another_key.id}")
 
-        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
+          gitlab_shell.remove_keys_not_found_in_db
         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}")
+    context 'when keys there are duplicate keys in the file that are not in the DB' do
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          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(gitlab_shell).to receive(:remove_key).with('key-1234')
+
+          gitlab_shell.remove_keys_not_found_in_db
         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
+      context 'authorized_keys_file set' 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
 
-        expect(ids).to eq(%W{1\n 2\n 3\n 4\n})
-      end
-    end
+        it 'removes the keys' do
+          expect(gitlab_shell).to receive(:remove_key).with('key-1234')
 
-    context 'when there are no keys in the authorized_keys file' do
-      before do
-        gitlab_shell.remove_all_keys
+          gitlab_shell.remove_keys_not_found_in_db
+        end
       end
+    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
+    context 'when keys there are duplicate keys in the file that ARE in the DB' do
+      context 'authorized_keys_file not set' do
+        before do
+          stub_gitlab_shell_setting(authorized_keys_file: nil)
+          gitlab_shell.remove_all_keys
+          @key = create(:key)
+          gitlab_shell.add_key(@key.shell_id, @key.key)
         end
 
-        expect(ids).to eq([])
+        it 'does not remove the key' do
+          expect(gitlab_shell).not_to receive(:remove_key).with("key-#{@key.id}")
+
+          gitlab_shell.remove_keys_not_found_in_db
+        end
       end
-    end
-  end
 
-  describe Gitlab::Shell::KeyAdder do
-    describe '#add_key' do
-      it 'removes trailing garbage' do
-        io = spy(:io)
-        adder = described_class.new(io)
+      context 'authorized_keys_file set' do
+        before do
+          gitlab_shell.remove_all_keys
+          @key = create(:key)
+          gitlab_shell.add_key(@key.shell_id, @key.key)
+        end
 
-        adder.add_key('key-42', "ssh-rsa foo bar\tbaz")
+        it 'does not remove the key' do
+          expect(gitlab_shell).not_to receive(:remove_key).with("key-#{@key.id}")
 
-        expect(io).to have_received(:puts).with("key-42\tssh-rsa foo")
+          gitlab_shell.remove_keys_not_found_in_db
+        end
       end
+    end
 
-      it 'handles multiple spaces in the key' do
-        io = spy(:io)
-        adder = described_class.new(io)
+    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
+        context 'authorized_keys_file not set' do
+          before do
+            stub_gitlab_shell_setting(authorized_keys_file: nil)
+            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
 
-        adder.add_key('key-42', "ssh-rsa  foo")
+          it 'removes the keys not in the DB' do
+            expect(gitlab_shell).to receive(:remove_key).with('key-1234')
 
-        expect(io).to have_received(:puts).with("key-42\tssh-rsa foo")
-      end
+            gitlab_shell.remove_keys_not_found_in_db
+          end
+        end
 
-      it 'raises an exception if the key contains a tab' do
-        expect do
-          described_class.new(StringIO.new).add_key('key-42', "ssh-rsa\tfoobar")
-        end.to raise_error(Gitlab::Shell::Error)
-      end
+        context 'authorized_keys_file set' 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(gitlab_shell).to receive(:remove_key).with('key-1234')
 
-      it 'raises an exception if the key contains a newline' do
-        expect do
-          described_class.new(StringIO.new).add_key('key-42', "ssh-rsa foobar\nssh-rsa pawned")
-        end.to raise_error(Gitlab::Shell::Error)
+            gitlab_shell.remove_keys_not_found_in_db
+          end
+        end
       end
     end
   end
@@ -566,12 +768,4 @@ 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) # rubocop:disable Cop/AvoidReturnFromBlocks
-    end
-
-    false
-  end
 end

spec/support/helpers/stub_configuration.rb

@@ -84,6 +84,10 @@ module StubConfiguration
     allow(Gitlab.config.kerberos).to receive_messages(to_settings(messages))
   end
 
+  def stub_gitlab_shell_setting(messages)
+    allow(Gitlab.config.gitlab_shell).to receive_messages(to_settings(messages))
+  end
+
   private
 
   # Modifies stubbed messages to also stub possible predicate versions