Add support for avatar upload transfer

app/services/geo/file_upload_service.rb

@@ -17,6 +17,8 @@ module Geo
 
       response =
         case params[:type]
+        when 'avatar'
+          handle_avatar_geo_request(params[:id], data)
         when 'lfs'
           handle_lfs_geo_request(params[:id], data)
         else
@@ -26,6 +28,27 @@ module Geo
       response
     end
 
+    def handle_avatar_geo_request(id, message)
+      status = { code: :not_found, message: 'Avatar file not found' }
+      upload = Upload.find(id)
+
+      return status unless upload.present?
+
+      if message[:id] != upload.model_id || message[:type] != upload.model_type || message[:checksum] != upload.checksum
+        return status
+      end
+
+      unless upload.model.avatar&.exists?
+        status[:message] = "#{upload.model_type} does not have a avatar"
+        return status
+      end
+
+      status[:code] = :ok
+      status[:message] = 'Success'
+      status[:file] = upload.model.avatar
+      status
+    end
+
     def handle_lfs_geo_request(id, message)
       status = { code: :not_found, message: 'LFS object not found' }
       lfs_object = LfsObject.find(id)

lib/gitlab/geo/avatar_transfer.rb

+module Gitlab
+  module Geo
+    class AvatarTransfer < Transfer
+      def initialize(upload)
+        @file_type = :avatar
+        @file_id = upload.id
+        @filename = upload.path
+        @request_data = build_request_data(upload)
+      end
+
+      private
+
+      def build_request_data(upload)
+        {
+          id: upload.model_id,
+          type: upload.model_type,
+          checksum: upload.checksum
+        }
+      end
+    end
+  end
+end

spec/requests/api/geo_spec.rb

@@ -121,6 +121,42 @@ describe API::Geo, api: true do
     end
   end
 
+  describe 'GET /geo/transfers/avatar/1' do
+    let!(:secondary_node) { create(:geo_node) }
+    let(:user) { create(:user, avatar: fixture_file_upload(Rails.root + 'spec/fixtures/dk.png', 'image/png')) }
+    let(:upload) { Upload.find_by(model: user, uploader: 'AvatarUploader') }
+    let(:transfer) { Gitlab::Geo::AvatarTransfer.new(upload) }
+    let(:req_header) { Gitlab::Geo::TransferRequest.new(transfer.request_data).headers }
+
+    before do
+      allow_any_instance_of(Gitlab::Geo::TransferRequest).to receive(:requesting_node).and_return(secondary_node)
+    end
+
+    it 'responds with 401 with invalid auth header' do
+      get api("/geo/transfers/avatar/#{upload.id}"), nil, Authorization: 'Test'
+
+      expect(response.status).to eq 401
+    end
+
+    context 'avatar file exists' do
+      it 'responds with 200 with X-Sendfile' do
+        get api("/geo/transfers/avatar/#{upload.id}"), nil, req_header
+
+        expect(response.status).to eq 200
+        expect(response.headers['Content-Type']).to eq('application/octet-stream')
+        expect(response.headers['X-Sendfile']).to eq(user.avatar.path)
+      end
+    end
+
+    context 'avatar does not exist' do
+      it 'responds with 404' do
+        get api("/geo/transfers/avatar/100000"), nil, req_header
+
+        expect(response.status).to eq 404
+      end
+    end
+  end
+
   describe 'GET /geo/transfers/lfs/1' do
     let!(:secondary_node) { create(:geo_node) }
     let(:lfs_object) { create(:lfs_object, :with_file) }

spec/services/geo/file_upload_service_spec.rb

 require 'spec_helper'
 
 describe Geo::FileUploadService, services: true do
-  let(:lfs_object) { create(:lfs_object, :with_file) }
-  let(:params) { { id: lfs_object.id, type: 'lfs' } }
-  let(:lfs_transfer) { Gitlab::Geo::LfsTransfer.new(lfs_object) }
-  let(:transfer_request) { Gitlab::Geo::TransferRequest.new(lfs_transfer.request_data) }
-  let(:req_header) { transfer_request.headers['Authorization'] }
-
-  before do
-    create(:geo_node, :current)
-  end
+  let!(:node) { create(:geo_node, :current) }
 
   describe '#execute' do
-    it 'sends LFS file' do
-      service = described_class.new(params, req_header)
+    context 'user avatar' do
+      let(:user) { create(:user, avatar: fixture_file_upload(Rails.root + 'spec/fixtures/dk.png', 'image/png')) }
+      let(:upload) { Upload.find_by(model: user, uploader: 'AvatarUploader') }
+      let(:params) { { id: upload.id, type: 'avatar' } }
+      let(:avatar_transfer) { Gitlab::Geo::AvatarTransfer.new(upload) }
+      let(:transfer_request) { Gitlab::Geo::TransferRequest.new(avatar_transfer.request_data) }
+      let(:req_header) { transfer_request.headers['Authorization'] }
+
+      it 'sends avatar file' do
+        service = described_class.new(params, req_header)
+
+        response = service.execute
+
+        expect(response[:code]).to eq(:ok)
+        expect(response[:file].file.path).to eq(user.avatar.path)
+      end
+
+      it 'returns nil if no authorization' do
+        service = described_class.new(params, nil)
+
+        expect(service.execute).to be_nil
+      end
+    end
+
+    context 'group avatar' do
+      let(:group) { create(:group, avatar: fixture_file_upload(Rails.root + 'spec/fixtures/dk.png', 'image/png')) }
+      let(:upload) { Upload.find_by(model: group, uploader: 'AvatarUploader') }
+      let(:params) { { id: upload.id, type: 'avatar' } }
+      let(:avatar_transfer) { Gitlab::Geo::AvatarTransfer.new(upload) }
+      let(:transfer_request) { Gitlab::Geo::TransferRequest.new(avatar_transfer.request_data) }
+      let(:req_header) { transfer_request.headers['Authorization'] }
+
+      it 'sends avatar file' do
+        service = described_class.new(params, req_header)
+
+        response = service.execute
+
+        expect(response[:code]).to eq(:ok)
+        expect(response[:file].file.path).to eq(group.avatar.path)
+      end
 
-      response = service.execute
+      it 'returns nil if no authorization' do
+        service = described_class.new(params, nil)
 
-      expect(response[:code]).to eq(:ok)
-      expect(response[:file].file.path).to eq(lfs_object.file.path)
+        expect(service.execute).to be_nil
+      end
     end
 
-    it 'returns nil if no authorization' do
-      service = described_class.new(params, nil)
+    context 'project avatar' do
+      let(:project) { create(:empty_project, avatar: fixture_file_upload(Rails.root + 'spec/fixtures/dk.png', 'image/png')) }
+      let(:upload) { Upload.find_by(model: project, uploader: 'AvatarUploader') }
+      let(:params) { { id: upload.id, type: 'avatar' } }
+      let(:avatar_transfer) { Gitlab::Geo::AvatarTransfer.new(upload) }
+      let(:transfer_request) { Gitlab::Geo::TransferRequest.new(avatar_transfer.request_data) }
+      let(:req_header) { transfer_request.headers['Authorization'] }
+
+      it 'sends avatar file' do
+        service = described_class.new(params, req_header)
+
+        response = service.execute
+
+        expect(response[:code]).to eq(:ok)
+        expect(response[:file].file.path).to eq(project.avatar.path)
+      end
+
+      it 'returns nil if no authorization' do
+        service = described_class.new(params, nil)
+
+        expect(service.execute).to be_nil
+      end
+    end
+
+    context 'LFS Object' do
+      let(:lfs_object) { create(:lfs_object, :with_file) }
+      let(:params) { { id: lfs_object.id, type: 'lfs' } }
+      let(:lfs_transfer) { Gitlab::Geo::LfsTransfer.new(lfs_object) }
+      let(:transfer_request) { Gitlab::Geo::TransferRequest.new(lfs_transfer.request_data) }
+      let(:req_header) { transfer_request.headers['Authorization'] }
+
+      it 'sends LFS file' do
+        service = described_class.new(params, req_header)
+
+        response = service.execute
+
+        expect(response[:code]).to eq(:ok)
+        expect(response[:file].file.path).to eq(lfs_object.file.path)
+      end
+
+      it 'returns nil if no authorization' do
+        service = described_class.new(params, nil)
 
-      expect(service.execute).to be_nil
+        expect(service.execute).to be_nil
+      end
     end
   end
 end