Implement SSH authentication support in Ruby

2 files changed, 139 insertions, 19 deletions

diff --git a/lib/gitlab_projects.rb b/lib/gitlab_projects.rb
index 4c63a40..267c679 100644
--- a/lib/gitlab_projects.rb
+++ b/lib/gitlab_projects.rb
@@ -211,20 +211,23 @@ class GitlabProjects
     cmd = %W(git --git-dir=#{full_path} fetch #{@name} --prune --quiet)
     cmd << '--force' if forced
     cmd << tags_option
-    pid = Process.spawn(*cmd)
 
-    begin
-      Timeout.timeout(timeout) do
-        Process.wait(pid)
-      end
+    setup_ssh_auth do |env|
+      pid = Process.spawn(env, *cmd)
 
-      $?.exitstatus.zero?
-    rescue => exception
-      $logger.error "Fetching remote #{@name} for project #{@project_name} failed due to: #{exception.message}."
+      begin
+        _, status = Timeout.timeout(timeout) do
+          Process.wait2(pid)
+        end
 
-      Process.kill('KILL', pid)
-      Process.wait
-      false
+        status.success?
+      rescue => exception
+        $logger.error "Fetching remote #{@name} for project #{@project_name} failed due to: #{exception.message}."
+
+        Process.kill('KILL', pid)
+        Process.wait
+        false
+      end
     end
   end
 
@@ -406,6 +409,59 @@ class GitlabProjects
     false
   end
 
+  # Builds a small shell script that can be used to execute SSH with a set of
+  # custom options.
+  #
+  # Options are expanded as `'-oKey="Value"'`, so SSH will correctly interpret
+  # paths with spaces in them. We trust the user not to embed single or double
+  # quotes in the key or value.
+  def custom_ssh_script(options = {})
+    args = options.map { |k, v| "'-o#{k}=\"#{v}\"'" }.join(' ')
+
+    [
+      "#!/bin/sh",
+      "exec ssh #{args} \"$@\""
+    ].join("\n")
+  end
+
+  # Known hosts data and private keys can be passed to gitlab-shell in the
+  # environment. If present, this method puts them into temporary files, writes
+  # a script that can substitute as `ssh`, setting the options to respect those
+  # files, and yields: { "GIT_SSH" => "/tmp/myScript" }
+  def setup_ssh_auth
+    options = {}
+
+    if ENV.key?('GITLAB_SHELL_SSH_KEY')
+      key_file = Tempfile.new('gitlab-shell-key-file', mode: 0o400)
+      key_file.write(ENV['GITLAB_SHELL_SSH_KEY'])
+      key_file.close
+
+      options['IdentityFile'] = key_file.path
+      options['IdentitiesOnly'] = true
+    end
+
+    if ENV.key?('GITLAB_SHELL_KNOWN_HOSTS')
+      known_hosts_file = Tempfile.new('gitlab-shell-known-hosts', mode: 0o400)
+      known_hosts_file.write(ENV['GITLAB_SHELL_KNOWN_HOSTS'])
+      known_hosts_file.close
+
+      options['StrictHostKeyChecking'] = true
+      options['UserKnownHostsFile'] = known_hosts_file.path
+    end
+
+    return yield({}) if options.empty?
+
+    script = Tempfile.new('gitlab-shell-ssh-wrapper', mode: 0o755)
+    script.write(custom_ssh_script(options))
+    script.close
+
+    yield('GIT_SSH' => script.path)
+  ensure
+    key_file.close! unless key_file.nil?
+    known_hosts_file.close! unless known_hosts_file.nil?
+    script.close! unless script.nil?
+  end
+
   def gitlab_reference_counter
     @gitlab_reference_counter ||= begin
       # Defer loading because this pulls in gitlab_net, which takes 100-200 ms
diff --git a/spec/gitlab_projects_spec.rb b/spec/gitlab_projects_spec.rb
index 7f80cbd..626f933 100644
--- a/spec/gitlab_projects_spec.rb
+++ b/spec/gitlab_projects_spec.rb
@@ -1,5 +1,6 @@
 require_relative 'spec_helper'
 require_relative '../lib/gitlab_projects'
+require_relative '../lib/gitlab_reference_counter'
 
 describe GitlabProjects do
   before do
@@ -322,33 +323,55 @@ describe GitlabProjects do
     let(:pid) { 1234 }
     let(:branch_name) { 'master' }
 
+    def stub_spawn(*args, wait: true, success: true)
+      expect(Process).to receive(:spawn).with(*args).and_return(pid)
+      expect(Process).to receive(:wait2).with(pid).and_return([pid, double(success?: success)]) if wait
+    end
+
+    def stub_env(args = {})
+      original = ENV.to_h
+      args.each { |k, v| ENV[k] = v }
+      yield
+    ensure
+      ENV.replace(original)
+    end
+
+    def stub_tempfile(name, *args)
+      file = StringIO.new
+      allow(file).to receive(:close!)
+      allow(file).to receive(:path).and_return(name)
+
+      expect(Tempfile).to receive(:new).with(*args).and_return(file)
+
+      file
+    end
+
     describe 'with default args' do
       let(:gl_projects) { build_gitlab_projects('fetch-remote', repos_path, project_name, remote_name, '600') }
       let(:cmd) { %W(git --git-dir=#{full_path} fetch #{remote_name} --prune --quiet --tags) }
 
       it 'executes the command' do
-        expect(Process).to receive(:spawn).with(*cmd).and_return(pid)
-        expect(Process).to receive(:wait).with(pid)
+        stub_spawn({}, *cmd)
 
         expect(gl_projects.exec).to be true
       end
 
       it 'raises timeout' do
+        stub_spawn({}, *cmd, wait: false)
         expect(Timeout).to receive(:timeout).with(600).and_raise(Timeout::Error)
-        expect(Process).to receive(:spawn).with(*cmd).and_return(pid)
-        expect(Process).to receive(:wait)
         expect(Process).to receive(:kill).with('KILL', pid)
+
         expect(gl_projects.exec).to be false
       end
     end
 
     describe 'with --force' do
       let(:gl_projects) { build_gitlab_projects('fetch-remote', repos_path, project_name, remote_name, '600', '--force') }
+      let(:env) { {} }
       let(:cmd) { %W(git --git-dir=#{full_path} fetch #{remote_name} --prune --quiet --force --tags) }
 
       it 'executes the command with forced option' do
-        expect(Process).to receive(:spawn).with(*cmd).and_return(pid)
-        expect(Process).to receive(:wait).with(pid)
+        stub_spawn({}, *cmd)
 
         expect(gl_projects.exec).to be true
       end
@@ -359,12 +382,53 @@ describe GitlabProjects do
       let(:cmd) { %W(git --git-dir=#{full_path} fetch #{remote_name} --prune --quiet --no-tags) }
 
       it 'executes the command' do
-        expect(Process).to receive(:spawn).with(*cmd).and_return(pid)
-        expect(Process).to receive(:wait).with(pid)
+        stub_spawn({}, *cmd)
 
         expect(gl_projects.exec).to be true
       end
     end
+
+    describe 'with GITLAB_SHELL_SSH_KEY' do
+      let(:gl_projects) { build_gitlab_projects('fetch-remote', repos_path, project_name, remote_name, '600') }
+      let(:cmd) { %W(git --git-dir=#{full_path} fetch #{remote_name} --prune --quiet --tags) }
+
+      around(:each) do |example|
+        stub_env('GITLAB_SHELL_SSH_KEY' => 'SSH KEY') { example.run }
+      end
+
+      it 'sets GIT_SSH to a custom script' do
+        script = stub_tempfile('scriptFile', 'gitlab-shell-ssh-wrapper', mode: 0755)
+        key = stub_tempfile('/tmp files/keyFile', 'gitlab-shell-key-file', mode: 0400)
+
+        stub_spawn({ 'GIT_SSH' => 'scriptFile' }, *cmd)
+
+        expect(gl_projects.exec).to be true
+
+        expect(script.string).to eq("#!/bin/sh\nexec ssh '-oIdentityFile=\"/tmp files/keyFile\"' '-oIdentitiesOnly=\"true\"' \"$@\"")
+        expect(key.string).to eq('SSH KEY')
+      end
+    end
+
+    describe 'with GITLAB_SHELL_KNOWN_HOSTS' do
+      let(:gl_projects) { build_gitlab_projects('fetch-remote', repos_path, project_name, remote_name, '600') }
+      let(:cmd) { %W(git --git-dir=#{full_path} fetch #{remote_name} --prune --quiet --tags) }
+
+      around(:each) do |example|
+        stub_env('GITLAB_SHELL_KNOWN_HOSTS' => 'KNOWN HOSTS') { example.run }
+      end
+
+      it 'sets GIT_SSH to a custom script' do
+        script = stub_tempfile('scriptFile', 'gitlab-shell-ssh-wrapper', mode: 0755)
+        key = stub_tempfile('/tmp files/knownHosts', 'gitlab-shell-known-hosts', mode: 0400)
+
+        stub_spawn({ 'GIT_SSH' => 'scriptFile' }, *cmd)
+
+        expect(gl_projects.exec).to be true
+
+        expect(script.string).to eq("#!/bin/sh\nexec ssh '-oStrictHostKeyChecking=\"true\"' '-oUserKnownHostsFile=\"/tmp files/knownHosts\"' \"$@\"")
+        expect(key.string).to eq('KNOWN HOSTS')
+      end
+    end
   end
 
   describe :import_project do