diff options
Diffstat (limited to 'lib/gitlab/import_export/command_line_util.rb')
-rw-r--r-- | lib/gitlab/import_export/command_line_util.rb | 29 |
1 files changed, 23 insertions, 6 deletions
diff --git a/lib/gitlab/import_export/command_line_util.rb b/lib/gitlab/import_export/command_line_util.rb index e520cade517..2b0467d8779 100644 --- a/lib/gitlab/import_export/command_line_util.rb +++ b/lib/gitlab/import_export/command_line_util.rb @@ -6,6 +6,8 @@ module Gitlab UNTAR_MASK = 'u+rwX,go+rX,go-w' DEFAULT_DIR_MODE = 0700 + FileOversizedError = Class.new(StandardError) + def tar_czf(archive:, dir:) tar_with_options(archive: archive, dir: dir, options: 'czf') end @@ -51,19 +53,34 @@ module Gitlab private - def download_or_copy_upload(uploader, upload_path) + def download_or_copy_upload(uploader, upload_path, size_limit: nil) if uploader.upload.local? copy_files(uploader.path, upload_path) else - download(uploader.url, upload_path) + download(uploader.url, upload_path, size_limit: size_limit) end end - def download(url, upload_path) - File.open(upload_path, 'w') do |file| - # Download (stream) file from the uploader's location - IO.copy_stream(URI.parse(url).open, file) + def download(url, upload_path, size_limit: nil) + File.open(upload_path, 'wb') do |file| + current_size = 0 + + Gitlab::HTTP.get(url, stream_body: true, allow_object_storage: true) do |fragment| + if [301, 302, 307].include?(fragment.code) + Gitlab::Import::Logger.warn(message: "received redirect fragment", fragment_code: fragment.code) + elsif fragment.code == 200 + current_size += fragment.bytesize + + raise FileOversizedError if size_limit.present? && current_size > size_limit + + file.write(fragment) + else + raise Gitlab::ImportExport::Error, "unsupported response downloading fragment #{fragment.code}" + end + end end + rescue FileOversizedError + nil end def tar_with_options(archive:, dir:, options:) |