From 76bafc00e62b3a90252f0d229c7ce98c2691da30 Mon Sep 17 00:00:00 2001
From: Jacob Vosmaer <jacob@gitlab.com>
Date: Fri, 16 Jun 2017 18:42:41 +0200
Subject: Pass Gitaly token on Ruby gRPC requests

---
 config/gitlab.yml.example                          |  3 +++
 lib/gitlab/gitaly_client.rb                        | 22 ++++++++++++++++++++++
 lib/gitlab/gitaly_client/commit.rb                 | 16 +++++-----------
 lib/gitlab/gitaly_client/notifications.rb          | 12 +++++++-----
 lib/gitlab/gitaly_client/ref.rb                    | 21 ++++++++++-----------
 spec/lib/gitlab/gitaly_client/commit_spec.rb       |  8 ++++----
 .../lib/gitlab/gitaly_client/notifications_spec.rb |  2 +-
 spec/lib/gitlab/gitaly_client/ref_spec.rb          | 10 +++++-----
 8 files changed, 57 insertions(+), 37 deletions(-)

diff --git a/config/gitlab.yml.example b/config/gitlab.yml.example
index 0b33783869b..7c7e444af3a 100644
--- a/config/gitlab.yml.example
+++ b/config/gitlab.yml.example
@@ -454,6 +454,8 @@ production: &base
     # introduced in 9.0). Eventually Gitaly use will become mandatory and
     # this option will disappear.
     enabled: true
+    # Default Gitaly authentication token. Can be overriden per storage.
+    token: ""
 
   #
   # 4. Advanced settings
@@ -469,6 +471,7 @@ production: &base
       default:
         path: /home/git/repositories/
         gitaly_address: unix:/home/git/gitlab/tmp/sockets/private/gitaly.socket # TCP connections are supported too (e.g. tcp://host:port)
+        # gitaly_token: 'special token' # Optional: override global gitaly.token for this storage.
 
   ## Backup settings
   backup:
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index 2343446bf22..e30bcd7ae3a 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -1,3 +1,5 @@
+require 'base64'
+
 require 'gitaly'
 
 module Gitlab
@@ -48,6 +50,26 @@ module Gitlab
       address
     end
 
+    # All RPC calls should use GitalyClient.call. This method makes sure
+    # that per-request authentication headers are set.
+    def self.call(storage, service, rpc, request)
+      metadata = request_metadata(storage)
+      metadata = yield(metadata) if block_given?
+      stub(service, storage).send(rpc, request, metadata)
+    end
+  
+    def self.request_metadata(storage)
+      encoded_token = Base64.strict_encode64(token(storage).to_s)
+      { metadata: { 'authorization' => "Bearer #{encoded_token}" } }
+    end
+
+    def self.token(storage)
+      params = Gitlab.config.repositories.storages[storage]
+      raise "storage not found: #{storage.inspect}" if params.nil?
+
+      params['gitaly_token'].presence || Gitlab.config.gitaly['token']
+    end
+
     def self.enabled?
       Gitlab.config.gitaly.enabled
     end
diff --git a/lib/gitlab/gitaly_client/commit.rb b/lib/gitlab/gitaly_client/commit.rb
index ba3da781dad..73c1848c95f 100644
--- a/lib/gitlab/gitaly_client/commit.rb
+++ b/lib/gitlab/gitaly_client/commit.rb
@@ -11,28 +11,26 @@ module Gitlab
       end
 
       def is_ancestor(ancestor_id, child_id)
-        stub = GitalyClient.stub(:commit, @repository.storage)
         request = Gitaly::CommitIsAncestorRequest.new(
           repository: @gitaly_repo,
           ancestor_id: ancestor_id,
           child_id: child_id
         )
 
-        stub.commit_is_ancestor(request).value
+        GitalyClient.call(@repository.storage, :commit, :commit_is_ancestor, request).value
       end
 
       def diff_from_parent(commit, options = {})
         request_params = commit_diff_request_params(commit, options)
         request_params[:ignore_whitespace_change] = options.fetch(:ignore_whitespace_change, false)
-
-        response = diff_service_stub.commit_diff(Gitaly::CommitDiffRequest.new(request_params))
+        request = Gitaly::CommitDiffRequest.new(request_params)
+        response = GitalyClient.call(@repository.storage, :diff, :commit_diff, request)
         Gitlab::Git::DiffCollection.new(GitalyClient::DiffStitcher.new(response), options)
       end
 
       def commit_deltas(commit)
-        request_params = commit_diff_request_params(commit)
-
-        response = diff_service_stub.commit_delta(Gitaly::CommitDeltaRequest.new(request_params))
+        request = Gitaly::CommitDeltaRequest.new(commit_diff_request_params(commit))
+        response = GitalyClient.call(@repository.storage, :diff, :commit_delta, request)
         response.flat_map do |msg|
           msg.deltas.map { |d| Gitlab::Git::Diff.new(d) }
         end
@@ -50,10 +48,6 @@ module Gitlab
           paths: options.fetch(:paths, [])
         }
       end
-
-      def diff_service_stub
-        GitalyClient.stub(:diff, @repository.storage)
-      end
     end
   end
 end
diff --git a/lib/gitlab/gitaly_client/notifications.rb b/lib/gitlab/gitaly_client/notifications.rb
index 719554eac52..78ed433e6b8 100644
--- a/lib/gitlab/gitaly_client/notifications.rb
+++ b/lib/gitlab/gitaly_client/notifications.rb
@@ -1,17 +1,19 @@
 module Gitlab
   module GitalyClient
     class Notifications
-      attr_accessor :stub
-
       # 'repository' is a Gitlab::Git::Repository
       def initialize(repository)
         @gitaly_repo = repository.gitaly_repository
-        @stub = GitalyClient.stub(:notifications, repository.storage)
+        @storage = repository.storage
       end
 
       def post_receive
-        request = Gitaly::PostReceiveRequest.new(repository: @gitaly_repo)
-        @stub.post_receive(request)
+        GitalyClient.call(
+          @storage,
+          :notifications,
+          :post_receive,
+          Gitaly::PostReceiveRequest.new(repository: @gitaly_repo)
+        )
       end
     end
   end
diff --git a/lib/gitlab/gitaly_client/ref.rb b/lib/gitlab/gitaly_client/ref.rb
index 227fe45642e..6d5f54dd959 100644
--- a/lib/gitlab/gitaly_client/ref.rb
+++ b/lib/gitlab/gitaly_client/ref.rb
@@ -1,29 +1,28 @@
 module Gitlab
   module GitalyClient
     class Ref
-      attr_accessor :stub
-
       # 'repository' is a Gitlab::Git::Repository
       def initialize(repository)
         @gitaly_repo = repository.gitaly_repository
-        @stub = GitalyClient.stub(:ref, repository.storage)
+        @storage = repository.storage
       end
 
       def default_branch_name
         request = Gitaly::FindDefaultBranchNameRequest.new(repository: @gitaly_repo)
-        branch_name = stub.find_default_branch_name(request).name
-
-        Gitlab::Git.branch_name(branch_name)
+        response = GitalyClient.call(@storage, :ref, :find_default_branch_name, request)
+        Gitlab::Git.branch_name(response.name)
       end
 
       def branch_names
         request = Gitaly::FindAllBranchNamesRequest.new(repository: @gitaly_repo)
-        consume_refs_response(stub.find_all_branch_names(request), prefix: 'refs/heads/')
+        response = GitalyClient.call(@storage, :ref, :find_all_branch_names, request)
+        consume_refs_response(response, prefix: 'refs/heads/')
       end
 
       def tag_names
         request = Gitaly::FindAllTagNamesRequest.new(repository: @gitaly_repo)
-        consume_refs_response(stub.find_all_tag_names(request), prefix: 'refs/tags/')
+        response = GitalyClient.call(@storage, :ref, :find_all_tag_names, request)
+        consume_refs_response(response, prefix: 'refs/tags/')
       end
 
       def find_ref_name(commit_id, ref_prefix)
@@ -32,8 +31,7 @@ module Gitlab
           commit_id: commit_id,
           prefix: ref_prefix
         )
-
-        stub.find_ref_name(request).name
+        GitalyClient.call(@storage, :ref, :find_ref_name, request).name
       end
 
       def count_tag_names
@@ -47,7 +45,8 @@ module Gitlab
       def local_branches(sort_by: nil)
         request = Gitaly::FindLocalBranchesRequest.new(repository: @gitaly_repo)
         request.sort_by = sort_by_param(sort_by) if sort_by
-        consume_branches_response(stub.find_local_branches(request))
+        response = GitalyClient.call(@storage, :ref, :find_local_branches, request)
+        consume_branches_response(response)
       end
 
       private
diff --git a/spec/lib/gitlab/gitaly_client/commit_spec.rb b/spec/lib/gitlab/gitaly_client/commit_spec.rb
index cf1bc74779e..dff5b25c712 100644
--- a/spec/lib/gitlab/gitaly_client/commit_spec.rb
+++ b/spec/lib/gitlab/gitaly_client/commit_spec.rb
@@ -16,7 +16,7 @@ describe Gitlab::GitalyClient::Commit do
           right_commit_id: commit.id
         )
 
-        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_diff).with(request)
+        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_diff).with(request, kind_of(Hash))
 
         described_class.new(repository).diff_from_parent(commit)
       end
@@ -31,7 +31,7 @@ describe Gitlab::GitalyClient::Commit do
           right_commit_id: initial_commit.id
         )
 
-        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_diff).with(request)
+        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_diff).with(request, kind_of(Hash))
 
         described_class.new(repository).diff_from_parent(initial_commit)
       end
@@ -61,7 +61,7 @@ describe Gitlab::GitalyClient::Commit do
           right_commit_id: commit.id
         )
 
-        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_delta).with(request).and_return([])
+        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_delta).with(request, kind_of(Hash)).and_return([])
 
         described_class.new(repository).commit_deltas(commit)
       end
@@ -76,7 +76,7 @@ describe Gitlab::GitalyClient::Commit do
           right_commit_id: initial_commit.id
         )
 
-        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_delta).with(request).and_return([])
+        expect_any_instance_of(Gitaly::Diff::Stub).to receive(:commit_delta).with(request, kind_of(Hash)).and_return([])
 
         described_class.new(repository).commit_deltas(initial_commit)
       end
diff --git a/spec/lib/gitlab/gitaly_client/notifications_spec.rb b/spec/lib/gitlab/gitaly_client/notifications_spec.rb
index e5c9e06a15e..c2b8ca9f501 100644
--- a/spec/lib/gitlab/gitaly_client/notifications_spec.rb
+++ b/spec/lib/gitlab/gitaly_client/notifications_spec.rb
@@ -9,7 +9,7 @@ describe Gitlab::GitalyClient::Notifications do
 
     it 'sends a post_receive message' do
       expect_any_instance_of(Gitaly::Notifications::Stub).
-        to receive(:post_receive).with(gitaly_request_with_path(storage_name, relative_path))
+        to receive(:post_receive).with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash))
 
       subject.post_receive
     end
diff --git a/spec/lib/gitlab/gitaly_client/ref_spec.rb b/spec/lib/gitlab/gitaly_client/ref_spec.rb
index 2ea44ef74b0..3272333bb33 100644
--- a/spec/lib/gitlab/gitaly_client/ref_spec.rb
+++ b/spec/lib/gitlab/gitaly_client/ref_spec.rb
@@ -21,7 +21,7 @@ describe Gitlab::GitalyClient::Ref do
     it 'sends a find_all_branch_names message' do
       expect_any_instance_of(Gitaly::Ref::Stub).
         to receive(:find_all_branch_names).
-          with(gitaly_request_with_path(storage_name, relative_path)).
+          with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash)).
           and_return([])
 
       client.branch_names
@@ -32,7 +32,7 @@ describe Gitlab::GitalyClient::Ref do
     it 'sends a find_all_tag_names message' do
       expect_any_instance_of(Gitaly::Ref::Stub).
         to receive(:find_all_tag_names).
-          with(gitaly_request_with_path(storage_name, relative_path)).
+          with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash)).
           and_return([])
 
       client.tag_names
@@ -43,7 +43,7 @@ describe Gitlab::GitalyClient::Ref do
     it 'sends a find_default_branch_name message' do
       expect_any_instance_of(Gitaly::Ref::Stub).
         to receive(:find_default_branch_name).
-          with(gitaly_request_with_path(storage_name, relative_path)).
+          with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash)).
         and_return(double(name: 'foo'))
 
       client.default_branch_name
@@ -54,7 +54,7 @@ describe Gitlab::GitalyClient::Ref do
     it 'sends a find_local_branches message' do
       expect_any_instance_of(Gitaly::Ref::Stub).
         to receive(:find_local_branches).
-          with(gitaly_request_with_path(storage_name, relative_path)).
+          with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash)).
           and_return([])
 
       client.local_branches
@@ -63,7 +63,7 @@ describe Gitlab::GitalyClient::Ref do
     it 'parses and sends the sort parameter' do
       expect_any_instance_of(Gitaly::Ref::Stub).
         to receive(:find_local_branches).
-          with(gitaly_request_with_params(sort_by: :UPDATED_DESC)).
+          with(gitaly_request_with_params(sort_by: :UPDATED_DESC), kind_of(Hash)).
           and_return([])
 
       client.local_branches(sort_by: 'updated_desc')
-- 
cgit v1.2.1


From 7bda1030a59b0723eade42b03f72918a75c20e9d Mon Sep 17 00:00:00 2001
From: Jacob Vosmaer <jacob@gitlab.com>
Date: Tue, 20 Jun 2017 16:09:58 +0200
Subject: Send gitaly token to workhorse when needed

---
 lib/gitlab/workhorse.rb           | 11 ++++++++---
 spec/lib/gitlab/workhorse_spec.rb |  6 +++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/lib/gitlab/workhorse.rb b/lib/gitlab/workhorse.rb
index 7f27317775c..43305f63911 100644
--- a/lib/gitlab/workhorse.rb
+++ b/lib/gitlab/workhorse.rb
@@ -26,7 +26,10 @@ module Gitlab
         }
 
         if Gitlab.config.gitaly.enabled
-          address = Gitlab::GitalyClient.address(project.repository_storage)
+          server = {
+            address: Gitlab::GitalyClient.address(project.repository_storage),
+            token: Gitlab::GitalyClient.token(project.repository_storage),
+          }
           params[:Repository] = repository.gitaly_repository.to_h
 
           feature_enabled = case action.to_s
@@ -39,8 +42,10 @@ module Gitlab
                             else
                               raise "Unsupported action: #{action}"
                             end
-
-          params[:GitalyAddress] = address if feature_enabled
+          if feature_enabled
+            params[:GitalyAddress] = server[:address] # This field will be deprecated
+            params[:GitalyServer] = server
+          end
         end
 
         params
diff --git a/spec/lib/gitlab/workhorse_spec.rb b/spec/lib/gitlab/workhorse_spec.rb
index ad19998dff4..a3e8166cb70 100644
--- a/spec/lib/gitlab/workhorse_spec.rb
+++ b/spec/lib/gitlab/workhorse_spec.rb
@@ -202,7 +202,11 @@ describe Gitlab::Workhorse, lib: true do
     context 'when Gitaly is enabled' do
       let(:gitaly_params) do
         {
-          GitalyAddress: Gitlab::GitalyClient.address('default')
+          GitalyAddress: Gitlab::GitalyClient.address('default'),
+          GitalyServer: {
+            address: Gitlab::GitalyClient.address('default'),
+            token: Gitlab::GitalyClient.token('default')
+          }
         }
       end
 
-- 
cgit v1.2.1


From 4f8af50b619dd87c16035f56dce5cbde3531aef1 Mon Sep 17 00:00:00 2001
From: Jacob Vosmaer <jacob@gitlab.com>
Date: Tue, 20 Jun 2017 17:31:49 +0200
Subject: Rubocop and comment fixes

---
 config/gitlab.yml.example   | 6 ++++--
 lib/gitlab/gitaly_client.rb | 4 ++--
 lib/gitlab/workhorse.rb     | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/config/gitlab.yml.example b/config/gitlab.yml.example
index 7c7e444af3a..82dd53013f9 100644
--- a/config/gitlab.yml.example
+++ b/config/gitlab.yml.example
@@ -454,8 +454,10 @@ production: &base
     # introduced in 9.0). Eventually Gitaly use will become mandatory and
     # this option will disappear.
     enabled: true
-    # Default Gitaly authentication token. Can be overriden per storage.
-    token: ""
+    # Default Gitaly authentication token. Can be overriden per storage. Can
+    # be left blank when Gitaly is running locally on a Unix socket, which
+    # is the normal way to deploy Gitaly.
+    token:
 
   #
   # 4. Advanced settings
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index e30bcd7ae3a..f605c06dfc3 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -50,8 +50,8 @@ module Gitlab
       address
     end
 
-    # All RPC calls should use GitalyClient.call. This method makes sure
-    # that per-request authentication headers are set.
+    # All Gitaly RPC call sites should use GitalyClient.call. This method
+    # makes sure that per-request authentication headers are set.
     def self.call(storage, service, rpc, request)
       metadata = request_metadata(storage)
       metadata = yield(metadata) if block_given?
diff --git a/lib/gitlab/workhorse.rb b/lib/gitlab/workhorse.rb
index 43305f63911..f96ee69096d 100644
--- a/lib/gitlab/workhorse.rb
+++ b/lib/gitlab/workhorse.rb
@@ -28,7 +28,7 @@ module Gitlab
         if Gitlab.config.gitaly.enabled
           server = {
             address: Gitlab::GitalyClient.address(project.repository_storage),
-            token: Gitlab::GitalyClient.token(project.repository_storage),
+            token: Gitlab::GitalyClient.token(project.repository_storage)
           }
           params[:Repository] = repository.gitaly_repository.to_h
 
-- 
cgit v1.2.1


From 218da88e1aa7b0538d069ce9dc6e3a4348c610fc Mon Sep 17 00:00:00 2001
From: Jacob Vosmaer <jacob@gitlab.com>
Date: Tue, 20 Jun 2017 17:42:54 +0200
Subject: Enable gitaly token auth when testing

---
 config/gitlab.yml.example             | 1 +
 lib/tasks/gitlab/gitaly.rake          | 5 +++--
 spec/tasks/gitlab/gitaly_rake_spec.rb | 1 +
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/config/gitlab.yml.example b/config/gitlab.yml.example
index 82dd53013f9..43a8c0078ca 100644
--- a/config/gitlab.yml.example
+++ b/config/gitlab.yml.example
@@ -599,6 +599,7 @@ test:
         gitaly_address: unix:tmp/tests/gitaly/gitaly.socket
   gitaly:
     enabled: true
+    token: secret
   backup:
     path: tmp/tests/backups
   gitlab_shell:
diff --git a/lib/tasks/gitlab/gitaly.rake b/lib/tasks/gitlab/gitaly.rake
index e88111c3725..a8db5701d0b 100644
--- a/lib/tasks/gitlab/gitaly.rake
+++ b/lib/tasks/gitlab/gitaly.rake
@@ -58,8 +58,9 @@ namespace :gitlab do
 
         storages << { name: key, path: val['path'] }
       end
-
-      TOML.dump(socket_path: address.sub(%r{\Aunix:}, ''), storage: storages)
+      config = { socket_path: address.sub(%r{\Aunix:}, ''), storage: storages }
+      config[:auth] = { token: 'secret' } if Rails.env.test?
+      TOML.dump(config)
     end
 
     def create_gitaly_configuration
diff --git a/spec/tasks/gitlab/gitaly_rake_spec.rb b/spec/tasks/gitlab/gitaly_rake_spec.rb
index cfa6c9ca8ce..c9a0f1cb144 100644
--- a/spec/tasks/gitlab/gitaly_rake_spec.rb
+++ b/spec/tasks/gitlab/gitaly_rake_spec.rb
@@ -89,6 +89,7 @@ describe 'gitlab:gitaly namespace rake task' do
         }
       }
       allow(Gitlab.config.repositories).to receive(:storages).and_return(config)
+      allow(Rails.env).to receive(:test?).and_return(false)
 
       expected_output = ''
       Timecop.freeze do
-- 
cgit v1.2.1