From 26968b70845e7780f14d835488fd45489792ca59 Mon Sep 17 00:00:00 2001
From: Steve Abrams <sabrams@gitlab.com>
Date: Fri, 15 May 2020 10:14:06 +0000
Subject: [PATCH] Detect container registry vendor/version/features
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: João Pereira <jpereira@gitlab.com>
---
 lib/container_registry/client.rb           | 17 +++++++
 spec/lib/container_registry/client_spec.rb | 52 +++++++++++++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/lib/container_registry/client.rb b/lib/container_registry/client.rb
index fdd889f54162a..118eb8e2d7ca5 100644
--- a/lib/container_registry/client.rb
+++ b/lib/container_registry/client.rb
@@ -13,6 +13,8 @@ class Client
     DOCKER_DISTRIBUTION_MANIFEST_V2_TYPE = 'application/vnd.docker.distribution.manifest.v2+json'
     OCI_MANIFEST_V1_TYPE = 'application/vnd.oci.image.manifest.v1+json'
     CONTAINER_IMAGE_V1_TYPE = 'application/vnd.docker.container.image.v1+json'
+    REGISTRY_VERSION_HEADER = 'gitlab-container-registry-version'
+    REGISTRY_FEATURES_HEADER = 'gitlab-container-registry-features'
 
     ACCEPTED_TYPES = [DOCKER_DISTRIBUTION_MANIFEST_V2_TYPE, OCI_MANIFEST_V1_TYPE].freeze
 
@@ -24,6 +26,21 @@ def initialize(base_uri, options = {})
       @options = options
     end
 
+    def registry_info
+      response = faraday.get("/v2/")
+
+      return {} unless response&.success?
+
+      version = response.headers[REGISTRY_VERSION_HEADER]
+      features = response.headers.fetch(REGISTRY_FEATURES_HEADER, '')
+
+      {
+        version: version,
+        features: features.split(',').map(&:strip),
+        vendor: version ? 'gitlab' : 'other'
+      }
+    end
+
     def repository_tags(name)
       response_body faraday.get("/v2/#{name}/tags/list")
     end
diff --git a/spec/lib/container_registry/client_spec.rb b/spec/lib/container_registry/client_spec.rb
index 0aad6568793b3..18bcff65f4102 100644
--- a/spec/lib/container_registry/client_spec.rb
+++ b/spec/lib/container_registry/client_spec.rb
@@ -85,7 +85,7 @@
     it 'follows 307 redirect for GET /v2/:name/blobs/:digest' do
       stub_request(:get, "http://container-registry/v2/group/test/blobs/sha256:0123456789012345")
         .with(headers: blob_headers)
-        .to_return(status: 307, body: "", headers: { Location: 'http://redirected' })
+        .to_return(status: 307, body: '', headers: { Location: 'http://redirected' })
       # We should probably use hash_excluding here, but that requires an update to WebMock:
       # https://github.com/bblimke/webmock/blob/master/lib/webmock/matchers/hash_excluding_matcher.rb
       stub_request(:get, "http://redirected/")
@@ -238,4 +238,54 @@ def stub_upload(path, content, digest, status = 200)
       it { is_expected.to be_falsey }
     end
   end
+
+  def stub_registry_info(headers: {}, status: 200)
+    stub_request(:get, 'http://container-registry/v2/')
+      .to_return(status: status, body: "", headers: headers)
+  end
+
+  describe '#registry_info' do
+    subject { client.registry_info }
+
+    context 'when the check is successful' do
+      context 'when using the GitLab container registry' do
+        before do
+          stub_registry_info(headers: {
+            'GitLab-Container-Registry-Version' => '2.9.1-gitlab',
+            'GitLab-Container-Registry-Features' => 'a,b,c'
+          })
+        end
+
+        it 'identifies the vendor as "gitlab"' do
+          expect(subject).to include(vendor: 'gitlab')
+        end
+
+        it 'identifies version and features' do
+          expect(subject).to include(version: '2.9.1-gitlab', features: %w[a b c])
+        end
+      end
+
+      context 'when using a third-party container registry' do
+        before do
+          stub_registry_info
+        end
+
+        it 'identifies the vendor as "other"' do
+          expect(subject).to include(vendor: 'other')
+        end
+
+        it 'does not identify version or features' do
+          expect(subject).to include(version: nil, features: [])
+        end
+      end
+    end
+
+    context 'when the check is not successful' do
+      it 'does not identify vendor, version or features' do
+        stub_registry_info(status: 500)
+
+        expect(subject).to eq({})
+      end
+    end
+  end
 end
-- 
GitLab