Move encoding methods to the more general EncodingHelpergitaly-encoding-helpers

9 files changed, 86 insertions, 80 deletions

diff --git a/lib/gitlab/encoding_helper.rb b/lib/gitlab/encoding_helper.rb
index 582028493e9..6b53eb4533d 100644
--- a/lib/gitlab/encoding_helper.rb
+++ b/lib/gitlab/encoding_helper.rb
@@ -71,6 +71,16 @@ module Gitlab
       end
     end
 
+    def encode_binary(s)
+      return "" if s.nil?
+
+      s.dup.force_encoding(Encoding::ASCII_8BIT)
+    end
+
+    def binary_stringio(s)
+      StringIO.new(s || '').tap { |io| io.set_encoding(Encoding::ASCII_8BIT) }
+    end
+
     private
 
     def clean(message)
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index b753ac46291..4507ea923b4 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -330,22 +330,6 @@ module Gitlab
       Google::Protobuf::Timestamp.new(seconds: t.to_i)
     end
 
-    def self.encode(s)
-      return "" if s.nil?
-
-      s.dup.force_encoding(Encoding::ASCII_8BIT)
-    end
-
-    def self.binary_stringio(s)
-      io = StringIO.new(s || '')
-      io.set_encoding(Encoding::ASCII_8BIT)
-      io
-    end
-
-    def self.encode_repeated(a)
-      Google::Protobuf::RepeatedField.new(:bytes, a.map { |s| self.encode(s) } )
-    end
-
     # The default timeout on all Gitaly calls
     def self.default_timeout
       return 0 if Sidekiq.server?
diff --git a/lib/gitlab/gitaly_client/commit_service.rb b/lib/gitlab/gitaly_client/commit_service.rb
index fb3e27770b4..8a29e8ec5b6 100644
--- a/lib/gitlab/gitaly_client/commit_service.rb
+++ b/lib/gitlab/gitaly_client/commit_service.rb
@@ -1,6 +1,8 @@
 module Gitlab
   module GitalyClient
     class CommitService
+      include Gitlab::EncodingHelper
+
       # The ID of empty tree.
       # See http://stackoverflow.com/a/40884093/1856239 and https://github.com/git/git/blob/3ad8b5bf26362ac67c9020bf8c30eee54a84f56d/cache.h#L1011-L1012
       EMPTY_TREE_ID = '4b825dc642cb6eb9a060e54bf8d69288fbee4904'.freeze
@@ -13,7 +15,7 @@ module Gitlab
       def ls_files(revision)
         request = Gitaly::ListFilesRequest.new(
           repository: @gitaly_repo,
-          revision: GitalyClient.encode(revision)
+          revision: encode_binary(revision)
         )
 
         response = GitalyClient.call(@repository.storage, :commit_service, :list_files, request, timeout: GitalyClient.medium_timeout)
@@ -73,7 +75,7 @@ module Gitlab
         request = Gitaly::TreeEntryRequest.new(
           repository: @gitaly_repo,
           revision: ref,
-          path: GitalyClient.encode(path),
+          path: encode_binary(path),
           limit: limit.to_i
         )
 
@@ -98,8 +100,8 @@ module Gitlab
       def tree_entries(repository, revision, path)
         request = Gitaly::GetTreeEntriesRequest.new(
           repository: @gitaly_repo,
-          revision: GitalyClient.encode(revision),
-          path: path.present? ? GitalyClient.encode(path) : '.'
+          revision: encode_binary(revision),
+          path: path.present? ? encode_binary(path) : '.'
         )
 
         response = GitalyClient.call(@repository.storage, :commit_service, :get_tree_entries, request, timeout: GitalyClient.medium_timeout)
@@ -112,8 +114,8 @@ module Gitlab
               type: gitaly_tree_entry.type.downcase,
               mode: gitaly_tree_entry.mode.to_s(8),
               name: File.basename(gitaly_tree_entry.path),
-              path: GitalyClient.encode(gitaly_tree_entry.path),
-              flat_path: GitalyClient.encode(gitaly_tree_entry.flat_path),
+              path: encode_binary(gitaly_tree_entry.path),
+              flat_path: encode_binary(gitaly_tree_entry.flat_path),
               commit_id: gitaly_tree_entry.commit_oid
             )
           end
@@ -135,8 +137,8 @@ module Gitlab
       def last_commit_for_path(revision, path)
         request = Gitaly::LastCommitForPathRequest.new(
           repository: @gitaly_repo,
-          revision: GitalyClient.encode(revision),
-          path: GitalyClient.encode(path.to_s)
+          revision: encode_binary(revision),
+          path: encode_binary(path.to_s)
         )
 
         gitaly_commit = GitalyClient.call(@repository.storage, :commit_service, :last_commit_for_path, request, timeout: GitalyClient.fast_timeout).commit
@@ -202,8 +204,8 @@ module Gitlab
       def raw_blame(revision, path)
         request = Gitaly::RawBlameRequest.new(
           repository: @gitaly_repo,
-          revision: GitalyClient.encode(revision),
-          path: GitalyClient.encode(path)
+          revision: encode_binary(revision),
+          path: encode_binary(path)
         )
 
         response = GitalyClient.call(@repository.storage, :commit_service, :raw_blame, request, timeout: GitalyClient.medium_timeout)
@@ -213,7 +215,7 @@ module Gitlab
       def find_commit(revision)
         request = Gitaly::FindCommitRequest.new(
           repository: @gitaly_repo,
-          revision: GitalyClient.encode(revision)
+          revision: encode_binary(revision)
         )
 
         response = GitalyClient.call(@repository.storage, :commit_service, :find_commit, request, timeout: GitalyClient.medium_timeout)
@@ -224,7 +226,7 @@ module Gitlab
       def patch(revision)
         request = Gitaly::CommitPatchRequest.new(
           repository: @gitaly_repo,
-          revision: GitalyClient.encode(revision)
+          revision: encode_binary(revision)
         )
         response = GitalyClient.call(@repository.storage, :diff_service, :commit_patch, request, timeout: GitalyClient.medium_timeout)
 
@@ -234,7 +236,7 @@ module Gitlab
       def commit_stats(revision)
         request = Gitaly::CommitStatsRequest.new(
           repository: @gitaly_repo,
-          revision: GitalyClient.encode(revision)
+          revision: encode_binary(revision)
         )
         GitalyClient.call(@repository.storage, :commit_service, :commit_stats, request, timeout: GitalyClient.medium_timeout)
       end
@@ -250,9 +252,9 @@ module Gitlab
         )
         request.after    = GitalyClient.timestamp(options[:after]) if options[:after]
         request.before   = GitalyClient.timestamp(options[:before]) if options[:before]
-        request.revision = GitalyClient.encode(options[:ref]) if options[:ref]
+        request.revision = encode_binary(options[:ref]) if options[:ref]
 
-        request.paths = GitalyClient.encode_repeated(Array(options[:path])) if options[:path].present?
+        request.paths = encode_repeated(Array(options[:path])) if options[:path].present?
 
         response = GitalyClient.call(@repository.storage, :commit_service, :find_commits, request, timeout: GitalyClient.medium_timeout)
 
@@ -264,7 +266,7 @@ module Gitlab
 
         enum = Enumerator.new do |y|
           shas.each_slice(20) do |revs|
-            request.shas = GitalyClient.encode_repeated(revs)
+            request.shas = encode_repeated(revs)
 
             y.yield request
 
@@ -303,7 +305,7 @@ module Gitlab
           repository: @gitaly_repo,
           left_commit_id: from_id,
           right_commit_id: to_id,
-          paths: options.fetch(:paths, []).compact.map { |path| GitalyClient.encode(path) }
+          paths: options.fetch(:paths, []).compact.map { |path| encode_binary(path) }
         }
       end
 
@@ -314,6 +316,10 @@ module Gitlab
           end
         end
       end
+
+      def encode_repeated(a)
+        Google::Protobuf::RepeatedField.new(:bytes, a.map { |s| encode_binary(s) } )
+      end
     end
   end
 end
diff --git a/lib/gitlab/gitaly_client/operation_service.rb b/lib/gitlab/gitaly_client/operation_service.rb
index 400a4af363b..c7732764880 100644
--- a/lib/gitlab/gitaly_client/operation_service.rb
+++ b/lib/gitlab/gitaly_client/operation_service.rb
@@ -1,6 +1,8 @@
 module Gitlab
   module GitalyClient
     class OperationService
+      include Gitlab::EncodingHelper
+
       def initialize(repository)
         @gitaly_repo = repository.gitaly_repository
         @repository = repository
@@ -9,7 +11,7 @@ module Gitlab
       def rm_tag(tag_name, user)
         request = Gitaly::UserDeleteTagRequest.new(
           repository: @gitaly_repo,
-          tag_name: GitalyClient.encode(tag_name),
+          tag_name: encode_binary(tag_name),
           user: Gitlab::Git::User.from_gitlab(user).to_gitaly
         )
 
@@ -24,9 +26,9 @@ module Gitlab
         request = Gitaly::UserCreateTagRequest.new(
           repository: @gitaly_repo,
           user: Gitlab::Git::User.from_gitlab(user).to_gitaly,
-          tag_name: GitalyClient.encode(tag_name),
-          target_revision: GitalyClient.encode(target),
-          message: GitalyClient.encode(message.to_s)
+          tag_name: encode_binary(tag_name),
+          target_revision: encode_binary(target),
+          message: encode_binary(message.to_s)
         )
 
         response = GitalyClient.call(@repository.storage, :operation_service, :user_create_tag, request)
@@ -44,9 +46,9 @@ module Gitlab
       def user_create_branch(branch_name, user, start_point)
         request = Gitaly::UserCreateBranchRequest.new(
           repository: @gitaly_repo,
-          branch_name: GitalyClient.encode(branch_name),
+          branch_name: encode_binary(branch_name),
           user: Gitlab::Git::User.from_gitlab(user).to_gitaly,
-          start_point: GitalyClient.encode(start_point)
+          start_point: encode_binary(start_point)
         )
         response = GitalyClient.call(@repository.storage, :operation_service,
           :user_create_branch, request)
@@ -64,7 +66,7 @@ module Gitlab
       def user_delete_branch(branch_name, user)
         request = Gitaly::UserDeleteBranchRequest.new(
           repository: @gitaly_repo,
-          branch_name: GitalyClient.encode(branch_name),
+          branch_name: encode_binary(branch_name),
           user: Gitlab::Git::User.from_gitlab(user).to_gitaly
         )
 
@@ -89,8 +91,8 @@ module Gitlab
             repository: @gitaly_repo,
             user: Gitlab::Git::User.from_gitlab(user).to_gitaly,
             commit_id: source_sha,
-            branch: GitalyClient.encode(target_branch),
-            message: GitalyClient.encode(message)
+            branch: encode_binary(target_branch),
+            message: encode_binary(message)
           )
         )
 
@@ -111,7 +113,7 @@ module Gitlab
           repository: @gitaly_repo,
           user: Gitlab::Git::User.from_gitlab(user).to_gitaly,
           commit_id: source_sha,
-          branch: GitalyClient.encode(target_branch)
+          branch: encode_binary(target_branch)
         )
 
         branch_update = GitalyClient.call(
@@ -152,9 +154,9 @@ module Gitlab
           repository: @gitaly_repo,
           user: Gitlab::Git::User.from_gitlab(user).to_gitaly,
           commit: commit.to_gitaly_commit,
-          branch_name: GitalyClient.encode(branch_name),
-          message: GitalyClient.encode(message),
-          start_branch_name: GitalyClient.encode(start_branch_name.to_s),
+          branch_name: encode_binary(branch_name),
+          message: encode_binary(message),
+          start_branch_name: encode_binary(start_branch_name.to_s),
           start_repository: start_repository.gitaly_repository
         )
 
diff --git a/lib/gitlab/gitaly_client/ref_service.rb b/lib/gitlab/gitaly_client/ref_service.rb
index 066e4e183c0..5bce1009878 100644
--- a/lib/gitlab/gitaly_client/ref_service.rb
+++ b/lib/gitlab/gitaly_client/ref_service.rb
@@ -72,7 +72,7 @@ module Gitlab
       end
 
       def ref_exists?(ref_name)
-        request = Gitaly::RefExistsRequest.new(repository: @gitaly_repo, ref: GitalyClient.encode(ref_name))
+        request = Gitaly::RefExistsRequest.new(repository: @gitaly_repo, ref: encode_binary(ref_name))
         response = GitalyClient.call(@storage, :ref_service, :ref_exists, request)
         response.value
       rescue GRPC::InvalidArgument => e
@@ -82,7 +82,7 @@ module Gitlab
       def find_branch(branch_name)
         request = Gitaly::FindBranchRequest.new(
           repository: @gitaly_repo,
-          name: GitalyClient.encode(branch_name)
+          name: encode_binary(branch_name)
         )
 
         response = GitalyClient.call(@repository.storage, :ref_service, :find_branch, request)
@@ -96,8 +96,8 @@ module Gitlab
       def create_branch(ref, start_point)
         request = Gitaly::CreateBranchRequest.new(
           repository: @gitaly_repo,
-          name: GitalyClient.encode(ref),
-          start_point: GitalyClient.encode(start_point)
+          name: encode_binary(ref),
+          start_point: encode_binary(start_point)
         )
 
         response = GitalyClient.call(@repository.storage, :ref_service, :create_branch, request)
@@ -121,7 +121,7 @@ module Gitlab
       def delete_branch(branch_name)
         request = Gitaly::DeleteBranchRequest.new(
           repository: @gitaly_repo,
-          name: GitalyClient.encode(branch_name)
+          name: encode_binary(branch_name)
         )
 
         GitalyClient.call(@repository.storage, :ref_service, :delete_branch, request)
diff --git a/lib/gitlab/gitaly_client/repository_service.rb b/lib/gitlab/gitaly_client/repository_service.rb
index c1f95396878..2115f388d5b 100644
--- a/lib/gitlab/gitaly_client/repository_service.rb
+++ b/lib/gitlab/gitaly_client/repository_service.rb
@@ -1,6 +1,8 @@
 module Gitlab
   module GitalyClient
     class RepositoryService
+      include Gitlab::EncodingHelper
+
       def initialize(repository)
         @repository = repository
         @gitaly_repo = repository.gitaly_repository
@@ -72,7 +74,7 @@ module Gitlab
       def find_merge_base(*revisions)
         request = Gitaly::FindMergeBaseRequest.new(
           repository: @gitaly_repo,
-          revisions: revisions.map { |r| GitalyClient.encode(r) }
+          revisions: revisions.map { |r| encode_binary(r) }
         )
 
         response = GitalyClient.call(@storage, :repository_service, :find_merge_base, request)
diff --git a/lib/gitlab/gitaly_client/wiki_service.rb b/lib/gitlab/gitaly_client/wiki_service.rb
index 337d225d081..5c5b170a3e0 100644
--- a/lib/gitlab/gitaly_client/wiki_service.rb
+++ b/lib/gitlab/gitaly_client/wiki_service.rb
@@ -3,6 +3,8 @@ require 'stringio'
 module Gitlab
   module GitalyClient
     class WikiService
+      include Gitlab::EncodingHelper
+
       MAX_MSG_SIZE = 128.kilobytes.freeze
 
       def initialize(repository)
@@ -13,12 +15,12 @@ module Gitlab
       def write_page(name, format, content, commit_details)
         request = Gitaly::WikiWritePageRequest.new(
           repository: @gitaly_repo,
-          name: GitalyClient.encode(name),
+          name: encode_binary(name),
           format: format.to_s,
           commit_details: gitaly_commit_details(commit_details)
         )
 
-        strio = GitalyClient.binary_stringio(content)
+        strio = binary_stringio(content)
 
         enum = Enumerator.new do |y|
           until strio.eof?
@@ -39,13 +41,13 @@ module Gitlab
       def update_page(page_path, title, format, content, commit_details)
         request = Gitaly::WikiUpdatePageRequest.new(
           repository: @gitaly_repo,
-          page_path: GitalyClient.encode(page_path),
-          title: GitalyClient.encode(title),
+          page_path: encode_binary(page_path),
+          title: encode_binary(title),
           format: format.to_s,
           commit_details: gitaly_commit_details(commit_details)
         )
 
-        strio = GitalyClient.binary_stringio(content)
+        strio = binary_stringio(content)
 
         enum = Enumerator.new do |y|
           until strio.eof?
@@ -63,7 +65,7 @@ module Gitlab
       def delete_page(page_path, commit_details)
         request = Gitaly::WikiDeletePageRequest.new(
           repository: @gitaly_repo,
-          page_path: GitalyClient.encode(page_path),
+          page_path: encode_binary(page_path),
           commit_details: gitaly_commit_details(commit_details)
         )
 
@@ -73,9 +75,9 @@ module Gitlab
       def find_page(title:, version: nil, dir: nil)
         request = Gitaly::WikiFindPageRequest.new(
           repository: @gitaly_repo,
-          title: GitalyClient.encode(title),
-          revision: GitalyClient.encode(version),
-          directory: GitalyClient.encode(dir)
+          title: encode_binary(title),
+          revision: encode_binary(version),
+          directory: encode_binary(dir)
         )
 
         response = GitalyClient.call(@repository.storage, :wiki_service, :wiki_find_page, request)
@@ -102,8 +104,8 @@ module Gitlab
       def find_file(name, revision)
         request = Gitaly::WikiFindFileRequest.new(
           repository: @gitaly_repo,
-          name: GitalyClient.encode(name),
-          revision: GitalyClient.encode(revision)
+          name: encode_binary(name),
+          revision: encode_binary(revision)
         )
 
         response = GitalyClient.call(@repository.storage, :wiki_service, :wiki_find_file, request)
@@ -158,9 +160,9 @@ module Gitlab
 
       def gitaly_commit_details(commit_details)
         Gitaly::WikiCommitDetails.new(
-          name: GitalyClient.encode(commit_details.name),
-          email: GitalyClient.encode(commit_details.email),
-          message: GitalyClient.encode(commit_details.message)
+          name: encode_binary(commit_details.name),
+          email: encode_binary(commit_details.email),
+          message: encode_binary(commit_details.message)
         )
       end
     end
diff --git a/spec/lib/gitlab/encoding_helper_spec.rb b/spec/lib/gitlab/encoding_helper_spec.rb
index f6e5c55240f..87ec2698fc1 100644
--- a/spec/lib/gitlab/encoding_helper_spec.rb
+++ b/spec/lib/gitlab/encoding_helper_spec.rb
@@ -145,4 +145,18 @@ describe Gitlab::EncodingHelper do
       end
     end
   end
+
+  describe 'encode_binary' do
+    [
+      [nil, ""],
+      ["", ""],
+      ["  ", "  "],
+      %w(a1 a1),
+      ["编码", "\xE7\xBC\x96\xE7\xA0\x81".b]
+    ].each do |input, result|
+      it "encodes #{input.inspect} to #{result.inspect}" do
+        expect(ext_class.encode_binary(input)).to eq(result)
+      end
+    end
+  end
 end
diff --git a/spec/lib/gitlab/gitaly_client_spec.rb b/spec/lib/gitlab/gitaly_client_spec.rb
index a871ed0df0e..309b7338ef0 100644
--- a/spec/lib/gitlab/gitaly_client_spec.rb
+++ b/spec/lib/gitlab/gitaly_client_spec.rb
@@ -38,20 +38,6 @@ describe Gitlab::GitalyClient, skip_gitaly_mock: true do
     end
   end
 
-  describe 'encode' do
-    [
-      [nil, ""],
-      ["", ""],
-      ["  ", "  "],
-      %w(a1 a1),
-      ["编码", "\xE7\xBC\x96\xE7\xA0\x81".b]
-    ].each do |input, result|
-      it "encodes #{input.inspect} to #{result.inspect}" do
-        expect(described_class.encode(input)).to eq result
-      end
-    end
-  end
-
   describe 'allow_n_plus_1_calls' do
     context 'when RequestStore is enabled', :request_store do
       it 'returns the result of the allow_n_plus_1_calls block' do