Save Geo files to a temporary file and rename after success

Closes #3618

changelogs/unreleased-ee/sh-geo-fix-file-transfer-empty-files.yml

+---
+title: Save Geo files to a temporary file and rename after success
+merge_request:
+author:
+type: fixed

lib/gitlab/geo/log_helpers.rb

+module Gitlab
+  module Geo
+    module LogHelpers
+      def log_info(message, details = {})
+        data = base_log_data(message)
+        data.merge!(details) if details
+        Gitlab::Geo::Logger.info(data)
+      end
+
+      def log_error(message, error)
+        data = base_log_data(message)
+        data[:error] = error.to_s
+        Gitlab::Geo::Logger.error(data)
+      end
+
+      protected
+
+      def base_log_data(message)
+        {
+          class: self.class.name,
+          message: message
+        }
+      end
+    end
+  end
+end

lib/gitlab/geo/project_log_helpers.rb

 module Gitlab
   module Geo
     module ProjectLogHelpers
-      def log_info(message, details = {})
-        data = base_log_data(message)
-        data.merge!(details) if details
-        Gitlab::Geo::Logger.info(data)
-      end
-
-      def log_error(message, error)
-        data = base_log_data(message)
-        data[:error] = error.to_s
-        Gitlab::Geo::Logger.error(data)
-      end
-
-      private
+      include LogHelpers
 
       def base_log_data(message)
         {

lib/gitlab/geo/transfer.rb

 module Gitlab
   module Geo
     class Transfer
+      include LogHelpers
+
       attr_reader :file_type, :file_id, :filename, :request_data
 
+      TEMP_PREFIX = 'tmp_'.freeze
+
       def initialize(file_type, file_id, filename, request_data)
         @file_type = file_type
         @file_id = file_id
@@ -50,34 +54,61 @@ module Gitlab
         true
       end
 
-      def log_transfer_error(message)
-        Rails.logger.error("#{self.class.name}: #{message}")
-      end
-
       # Use HTTParty for now but switch to curb if performance becomes
       # an issue
       def download_file(url, req_headers)
         file_size = -1
+        temp_file = open_temp_file(filename)
+
+        return unless temp_file
 
         begin
-          File.open(filename, "wb") do |file|
-            response = HTTParty.get(url, headers: req_headers, stream_body: true) do |fragment|
-              file.write(fragment)
-            end
-
-            if response.success?
-              file_size = File.stat(filename).size
-              Rails.logger.info("GitLab Geo: Successfully downloaded #{filename} (#{file_size} bytes)")
-            else
-              log_transfer_error("Unsuccessful download: #{response.code} #{response.msg}")
-            end
+          response = HTTParty.get(url, headers: req_headers, stream_body: true) do |fragment|
+            temp_file.write(fragment)
+          end
+
+          temp_file.flush
+
+          unless response.success?
+            log_error("Unsuccessful download", response_code: response.code, response_msg: response.msg)
+            return file_size
+          end
+
+          if File.directory?(filename)
+            log_error("Destination file is a directory", filename: filename)
+            return file_size
           end
+
+          FileUtils.mv(temp_file.path, filename)
+
+          file_size = File.stat(filename).size
+          log_info("Successful downloaded", filename: filename, file_size_bytes: file_size)
         rescue StandardError, HTTParty::Error => e
-          log_transfer_error("Error downloading file: #{e}")
+          log_error("Error downloading file", error: e)
+        ensure
+          temp_file.close
+          temp_file.unlink
         end
 
         file_size
       end
+
+      def default_permissions
+        0666 - File.umask
+      end
+
+      def open_temp_file(target_filename)
+        begin
+          # Make sure the file is in the same directory to prevent moves across filesystems
+          pathname = Pathname.new(target_filename)
+          temp = Tempfile.new(TEMP_PREFIX + pathname.basename.to_s, pathname.dirname.to_s)
+          temp.chmod(default_permissions)
+          temp.binmode
+          temp
+        rescue StandardError => e
+          log_error("Error creating temporary file", error: e)
+        end
+      end
     end
   end
 end

spec/lib/gitlab/geo/transfer_spec.rb

@@ -7,13 +7,9 @@ describe Gitlab::Geo::Transfer do
   set(:secondary_node) { create(:geo_node) }
   set(:lfs_object) { create(:lfs_object, :with_file) }
   let(:url) { primary_node.geo_transfers_url(:lfs, lfs_object.id.to_s) }
-  let(:content) { StringIO.new("1\n2\n3") }
+  let(:content) { SecureRandom.random_bytes(10) }
   let(:size) { File.stat(lfs_object.file.path).size }
 
-  before do
-    allow(File).to receive(:open).with(lfs_object.file.path, "wb").and_yield(content)
-  end
-
   subject do
     described_class.new(:lfs,
                         lfs_object.id,
@@ -21,12 +17,33 @@ describe Gitlab::Geo::Transfer do
                         { sha256: lfs_object.oid })
   end
 
-  it '#download_from_primary' do
-    stub_current_geo_node(secondary_node)
+  context '#download_from_primary' do
+    before do
+      stub_current_geo_node(secondary_node)
+    end
+
+    it 'when the destination filename is a directory' do
+      transfer = described_class.new(:lfs, lfs_object.id, '/tmp', { sha256: lfs_object.id })
+
+      expect(transfer.download_from_primary).to eq(nil)
+    end
+
+    it 'when the HTTP response is successful' do
+      expect(FileUtils).to receive(:mv).with(anything, lfs_object.file.path).and_call_original
+      response = double(success?: true)
+      expect(HTTParty).to receive(:get).and_yield(content.to_s).and_return(response)
+
+      expect(subject.download_from_primary).to eq(size)
+      expect(File.stat(lfs_object.file.path).size).to eq(size)
+      expect(File.binread(lfs_object.file.path)).to eq(content)
+    end
 
-    response = double(success?: true)
-    expect(HTTParty).to receive(:get).and_return(response)
+    it 'when the HTTP response is unsuccessful' do
+      expect(FileUtils).not_to receive(:mv).with(anything, lfs_object.file.path).and_call_original
+      response = double(success?: false, code: 404, msg: 'No such file')
+      expect(HTTParty).to receive(:get).and_return(response)
 
-    expect(subject.download_from_primary).to eq(size)
+      expect(subject.download_from_primary).to eq(-1)
+    end
   end
 end