Merge branch 'rs-ruby-feature-flags' into 'master'

Pass Ruby-specific feature flags to the Ruby server

See merge request gitlab-org/gitaly!1818

6 files changed, 141 insertions, 0 deletions

diff --git a/changelogs/unreleased/rs-ruby-feature-flags.yml b/changelogs/unreleased/rs-ruby-feature-flags.yml
new file mode 100644
index 000000000..c044dcb53
--- /dev/null
+++ b/changelogs/unreleased/rs-ruby-feature-flags.yml
@@ -0,0 +1,5 @@
+---
+title: Pass Ruby-specific feature flags to the Ruby server
+merge_request: 1818
+author:
+type: added
diff --git a/internal/rubyserver/proxy.go b/internal/rubyserver/proxy.go
index 7535ab681..cbe0947ec 100644
--- a/internal/rubyserver/proxy.go
+++ b/internal/rubyserver/proxy.go
@@ -15,6 +15,9 @@ import (
 // forwarded as-is to gitaly-ruby.
 var ProxyHeaderWhitelist = []string{"gitaly-servers"}
 
+// Headers prefixed with this string get whitelisted automatically
+const rubyFeaturePrefix = "gitaly-feature-ruby-"
+
 const (
 	storagePathHeader  = "gitaly-storage-path"
 	repoPathHeader     = "gitaly-repo-path"
@@ -63,6 +66,13 @@ func setHeaders(ctx context.Context, repo *gitalypb.Repository, mustExist bool)
 	)
 
 	if inMD, ok := metadata.FromIncomingContext(ctx); ok {
+		// Automatically whitelist any Ruby-specific feature flag
+		for header := range inMD {
+			if strings.HasPrefix(header, rubyFeaturePrefix) {
+				ProxyHeaderWhitelist = append(ProxyHeaderWhitelist, header)
+			}
+		}
+
 		for _, header := range ProxyHeaderWhitelist {
 			for _, v := range inMD[header] {
 				md = metadata.Join(md, metadata.Pairs(header, v))
diff --git a/internal/rubyserver/proxy_test.go b/internal/rubyserver/proxy_test.go
index 5b804c061..03936f41c 100644
--- a/internal/rubyserver/proxy_test.go
+++ b/internal/rubyserver/proxy_test.go
@@ -43,3 +43,20 @@ func TestSetHeadersPreservesWhitelistedMetadata(t *testing.T) {
 
 	require.Equal(t, []string{value}, outMd[key], "outgoing MD should contain whitelisted key")
 }
+
+func TestRubyFeatureHeaders(t *testing.T) {
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	key := "gitaly-feature-ruby-test-feature"
+	value := "true"
+	inCtx := metadata.NewIncomingContext(ctx, metadata.Pairs(key, value))
+
+	outCtx, err := SetHeaders(inCtx, testRepo)
+	require.NoError(t, err)
+
+	outMd, ok := metadata.FromOutgoingContext(outCtx)
+	require.True(t, ok, "outgoing context should have metadata")
+
+	require.Equal(t, []string{value}, outMd[key], "outgoing MD should contain whitelisted feature key")
+}
diff --git a/ruby/lib/gitaly_server.rb b/ruby/lib/gitaly_server.rb
index a57f4dc24..acbfd2f96 100644
--- a/ruby/lib/gitaly_server.rb
+++ b/ruby/lib/gitaly_server.rb
@@ -13,6 +13,7 @@ require_relative 'gitaly_server/wiki_service.rb'
 require_relative 'gitaly_server/conflicts_service.rb'
 require_relative 'gitaly_server/remote_service.rb'
 require_relative 'gitaly_server/health_service.rb'
+require_relative 'gitaly_server/feature_flags.rb'
 
 module GitalyServer
   STORAGE_PATH_HEADER = 'gitaly-storage-path'.freeze
@@ -37,6 +38,10 @@ module GitalyServer
     call.metadata.fetch(REPO_ALT_DIRS_HEADER)
   end
 
+  def self.feature_flags(call)
+    FeatureFlags.new(call.metadata)
+  end
+
   def self.client(call)
     Client.new(call.metadata[GITALY_SERVERS_HEADER])
   end
diff --git a/ruby/lib/gitaly_server/feature_flags.rb b/ruby/lib/gitaly_server/feature_flags.rb
new file mode 100644
index 000000000..fcfa994f8
--- /dev/null
+++ b/ruby/lib/gitaly_server/feature_flags.rb
@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+module GitalyServer
+  # Interface to Ruby-specific feature flags passed to the Gitaly Ruby server
+  # via headers.
+  class FeatureFlags
+    # Only headers prefixed with this String will be made available
+    HEADER_PREFIX = 'gitaly-feature-ruby-'
+
+    def initialize(metadata)
+      @flags = metadata.select do |key, _|
+        key.start_with?(HEADER_PREFIX)
+      end
+    end
+
+    # Check if a given flag is enabled
+    #
+    # The `gitaly-feature-ruby-` prefix is optional, and underscores are
+    # translated to hyphens automatically.
+    #
+    # Examples
+    #
+    #   enabled?('gitaly-feature-ruby-my-flag')
+    #   => true
+    #
+    #   enabled?(:my_flag)
+    #   => true
+    #
+    #   enabled?('my-flag')
+    #   => true
+    #
+    #   enabled?(:unknown_flag)
+    #   => false
+    def enabled?(flag)
+      flag = normalize_flag(flag)
+
+      @flags.fetch(flag, false) == 'true'
+    end
+
+    def disabled?(flag)
+      !enabled?(flag)
+    end
+
+    def inspect
+      pairs = @flags.map { |name, value| "#{name}=#{value}" }
+      pairs.unshift(self.class.name)
+
+      "#<#{pairs.join(' ')}>"
+    end
+
+    private
+
+    def normalize_flag(flag)
+      flag = flag.to_s.delete_prefix(HEADER_PREFIX).tr('_', '-')
+
+      "#{HEADER_PREFIX}#{flag}"
+    end
+  end
+end
diff --git a/ruby/spec/lib/gitaly_server/feature_flags_spec.rb b/ruby/spec/lib/gitaly_server/feature_flags_spec.rb
new file mode 100644
index 000000000..56a911bcb
--- /dev/null
+++ b/ruby/spec/lib/gitaly_server/feature_flags_spec.rb
@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe GitalyServer::FeatureFlags do
+  describe '#enabled?' do
+    let(:metadata) do
+      {
+        "#{described_class::HEADER_PREFIX}some-feature" => 'true',
+        'gitaly-storage-path' => 'foo',
+        'gitaly-repo-path' => 'bar'
+      }
+    end
+
+    subject { described_class.new(metadata) }
+
+    it 'returns true for an enabled flag' do
+      expect(subject.enabled?(:some_feature)).to eq(true)
+    end
+
+    it 'returns false for an unknown flag' do
+      expect(subject.enabled?(:missing_feature)).to eq(false)
+    end
+
+    it 'removes the prefix if provided' do
+      expect(subject.enabled?(metadata.keys.first)).to eq(true)
+    end
+
+    it 'translates underscores' do
+      expect(subject.enabled?('some-feature')).to eq(true)
+    end
+  end
+
+  describe '#disabled?' do
+    it 'is the inverse of `enabled?`' do
+      instance = described_class.new({})
+
+      expect(instance).to receive(:enabled?)
+        .with(:some_feature)
+        .and_return(false)
+
+      expect(instance.disabled?(:some_feature)).to eq(true)
+    end
+  end
+end