Pass Gitaly token on Ruby gRPC requests

4 files changed, 44 insertions, 27 deletions

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