Improve network code for dependency proxy feature

Signed-off-by: Dmitriy Zaporozhets <dmitriy.zaporozhets@gmail.com>

ee/app/controllers/groups/dependency_proxy_for_containers_controller.rb

@@ -4,11 +4,12 @@ class Groups::DependencyProxyForContainersController < Groups::ApplicationContro
   include SendFileUpload
 
   before_action :ensure_feature_enabled!
+  before_action :ensure_token_granted!
 
   def manifest
-    output = DependencyProxy::PullManifestService.new(image, tag, token).execute
+    response = DependencyProxy::PullManifestService.new(image, tag, token).execute
 
-    render json: output
+    render status: response[:code], json: response[:body]
   end
 
   def blob
@@ -29,11 +30,7 @@ class Groups::DependencyProxyForContainersController < Groups::ApplicationContro
   end
 
   def token
-    @token ||= request_token
-  end
-
-  def request_token
-    DependencyProxy::RequestTokenService.new(image).execute
+    @token
   end
 
   def ensure_feature_enabled!
@@ -41,4 +38,14 @@ class Groups::DependencyProxyForContainersController < Groups::ApplicationContro
         group.feature_available?(:dependency_proxy) &&
         group.dependency_proxy_setting&.enabled
   end
+
+  def ensure_token_granted!
+    response = DependencyProxy::RequestTokenService.new(image).execute
+
+    if response[:code] == 200 and response[:body].present?
+      @token = response[:body]
+    else
+      render status: response[:code], json: response[:body]
+    end
+  end
 end

ee/app/services/dependency_proxy/base_service.rb

@@ -13,5 +13,12 @@ module DependencyProxy
         Authorization: "Bearer #{@token}"
       }
     end
+
+    def to_response(code, body)
+      {
+        code: code,
+        body: body
+      }
+    end
   end
 end

ee/app/services/dependency_proxy/pull_manifest_service.rb

@@ -10,7 +10,10 @@ module DependencyProxy
 
     def execute
       response = Gitlab::HTTP.get(manifest_url, headers: auth_headers)
-      response.body
+
+      to_response(response.code, response.body)
+    rescue Net::OpenTimeout, Net::ReadTimeout => exception
+      to_response(599, exception.message)
     end
 
     private

ee/app/services/dependency_proxy/request_token_service.rb

@@ -9,7 +9,15 @@ module DependencyProxy
     def execute
       response = Gitlab::HTTP.get(auth_url)
 
-      JSON.parse(response.body)['token']
+      if response.code == 200
+        to_response(200, JSON.parse(response.body)['token'])
+      else
+        to_response(response.code, 'Expected 200 response code for an access token')
+      end
+    rescue Net::OpenTimeout, Net::ReadTimeout => exception
+      to_response(599, exception.message)
+    rescue JSON::ParserError
+      to_response(500, 'Failed to parse a response body for an access token')
     end
 
     private

ee/spec/controllers/groups/dependency_proxy_for_containers_controller_spec.rb

@@ -4,30 +4,58 @@ require 'spec_helper'
 
 describe Groups::DependencyProxyForContainersController do
   let(:group) { create(:group) }
+  let(:token_response) { { code: 200, body: 'abcd1234' } }
 
   before do
     allow(Gitlab.config.dependency_proxy)
       .to receive(:enabled).and_return(true)
 
     allow_any_instance_of(DependencyProxy::RequestTokenService)
-      .to receive(:execute).and_return('abcd1234')
+      .to receive(:execute).and_return(token_response)
   end
 
   describe 'GET #manifest' do
     let(:manifest) { { foo: 'bar' }.to_json }
+    let(:pull_response) { { code: 200, body: manifest } }
 
     before do
       allow_any_instance_of(DependencyProxy::PullManifestService)
-        .to receive(:execute).and_return(manifest)
+        .to receive(:execute).and_return(pull_response)
     end
 
-    it 'returns 200 with manifest file' do
-      enable_dependency_proxy
+    context 'feature enabled' do
+      before do
+        enable_dependency_proxy
+      end
 
-      get_manifest
+      context 'remote token request fails' do
+        let(:token_response) { { code: 503, body: 'Service Unavailable' } }
+
+        it 'proxies status from the remote token request' do
+          get_manifest
+
+          expect(response).to have_gitlab_http_status(503)
+          expect(response.body).to eq('Service Unavailable')
+        end
+      end
+
+      context 'remote manifest request fails' do
+        let(:pull_response) { { code: 400, body: '' } }
 
-      expect(response).to have_gitlab_http_status(200)
-      expect(response.body).to eq(manifest)
+        it 'proxies status from the remote manifest request' do
+          get_manifest
+
+          expect(response).to have_gitlab_http_status(400)
+          expect(response.body).to be_empty
+        end
+      end
+
+      it 'returns 200 with manifest file' do
+        get_manifest
+
+        expect(response).to have_gitlab_http_status(200)
+        expect(response.body).to eq(manifest)
+      end
     end
 
     it 'returns 404 when feature is disabled' do

ee/spec/services/dependency_proxy/find_or_create_blob_service_spec.rb

@@ -39,7 +39,7 @@ describe DependencyProxy::FindOrCreateBlobService do
 
   context 'no such blob exists remotely' do
     before do
-      stub_blob_download_not_found(image, blob_sha)
+      stub_blob_download(image, blob_sha, 404)
     end
 
     it { is_expected.to be nil }

ee/spec/services/dependency_proxy/pull_manifest_service_spec.rb

@@ -9,13 +9,34 @@ describe DependencyProxy::PullManifestService do
   let(:token) { Digest::SHA256.hexdigest('123') }
   let(:manifest) { { foo: 'bar' }.to_json }
 
-  subject { described_class.new(image, tag, token) }
+  subject { described_class.new(image, tag, token).execute }
 
-  before do
-    stub_manifest_download(image, tag)
+  context 'remote request is successful' do
+    before do
+      stub_manifest_download(image, tag)
+    end
+
+    it { expect(subject[:code]).to eq(200) }
+    it { expect(subject[:body]).to eq(manifest) }
+  end
+
+  context 'remote request is not found' do
+    before do
+      stub_manifest_download(image, tag, 404, 'Not found')
+    end
+
+    it { expect(subject[:code]).to eq(404) }
+    it { expect(subject[:body]).to eq('Not found') }
   end
 
-  it 'downloads blob and writes it into the file' do
-    expect(subject.execute).to eq(manifest)
+  context 'net timeout exception' do
+    before do
+      manifest_link = DependencyProxy::Registry.manifest_url(image, tag)
+
+      stub_request(:get, manifest_link).to_timeout
+    end
+
+    it { expect(subject[:code]).to eq(599) }
+    it { expect(subject[:body]).to eq('execution expired') }
   end
 end

ee/spec/services/dependency_proxy/request_token_service_spec.rb

@@ -9,11 +9,41 @@ describe DependencyProxy::RequestTokenService do
 
   subject { described_class.new(image).execute }
 
-  before do
-    stub_registry_auth(image, token)
+  context 'remote request is successful' do
+    before do
+      stub_registry_auth(image, token)
+    end
+
+    it { expect(subject[:code]).to eq(200) }
+    it { expect(subject[:body]).to eq(token) }
   end
 
-  it 'requests an access token from auth service' do
-    is_expected.to eq(token)
+  context 'remote request is not found' do
+    before do
+      stub_registry_auth(image, token, 404)
+    end
+
+    it { expect(subject[:code]).to eq(404) }
+    it { expect(subject[:body]).to eq('Expected 200 response code for an access token') }
+  end
+
+  context 'failed to parse response body' do
+    before do
+      stub_registry_auth(image, token, 200, 'dasd1321: wow')
+    end
+
+    it { expect(subject[:code]).to eq(500) }
+    it { expect(subject[:body]).to eq('Failed to parse a response body for an access token') }
+  end
+
+  context 'net timeout exception' do
+    before do
+      auth_link = DependencyProxy::Registry.auth_url(image)
+
+      stub_request(:any, auth_link).to_timeout
+    end
+
+    it { expect(subject[:code]).to eq(599) }
+    it { expect(subject[:body]).to eq('execution expired') }
   end
 end

ee/spec/support/helpers/ee/dependency_proxy_helpers.rb

@@ -2,33 +2,26 @@
 
 module EE
   module DependencyProxyHelpers
-    def stub_registry_auth(image, token)
+    def stub_registry_auth(image, token, status = 200, body = nil)
       auth_body = { 'token' => token }.to_json
       auth_link = registry.auth_url(image)
 
       stub_request(:get, auth_link)
-        .to_return(status: 200, body: auth_body)
+        .to_return(status: status, body: body || auth_body)
     end
 
-    def stub_manifest_download(image, tag)
+    def stub_manifest_download(image, tag, status = 200, body = nil)
       manifest_url = registry.manifest_url(image, tag)
 
       stub_request(:get, manifest_url)
-        .to_return(status: 200, body: manifest)
+        .to_return(status: status, body: body || manifest)
     end
 
-    def stub_blob_download(image, blob_sha)
+    def stub_blob_download(image, blob_sha, status = 200, body = '123456')
       download_link = registry.blob_url(image, blob_sha)
 
       stub_request(:get, download_link)
-        .to_return(status: 200, body: '123456')
-    end
-
-    def stub_blob_download_not_found(image, blob_sha)
-      download_link = registry.blob_url(image, blob_sha)
-
-      stub_request(:get, download_link)
-        .to_return(status: 404)
+        .to_return(status: status, body: body)
     end
 
     private