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