diff --git a/app/logical/danbooru_image_resizer.rb b/app/logical/danbooru_image_resizer.rb index 2dd58f8e7..20450b768 100644 --- a/app/logical/danbooru_image_resizer.rb +++ b/app/logical/danbooru_image_resizer.rb @@ -1,4 +1,6 @@ module DanbooruImageResizer + extend self + # Taken from ArgyllCMS 2.0.0 (see also: https://ninedegreesbelow.com/photography/srgb-profile-comparison.html) SRGB_PROFILE = "#{Rails.root}/config/sRGB.icm" # http://jcupitt.github.io/libvips/API/current/libvips-resample.html#vips-thumbnail @@ -9,7 +11,7 @@ module DanbooruImageResizer # XXX libvips-8.4 on Debian doesn't support the `Vips::Image.thumbnail` method. # On 8.4 we have to shell out to vipsthumbnail instead. Remove when Debian supports 8.5. - def self.resize(file, width, height, quality = 90) + def resize(file, width, height, quality = 90) if Vips.at_least_libvips?(8, 5) resize_ruby(file, width, height, quality) else @@ -17,13 +19,13 @@ module DanbooruImageResizer end end - def self.crop(file, width, height, quality = 90) + def crop(file, width, height, quality = 90) crop_shell(file, width, height, quality) end # https://github.com/jcupitt/libvips/wiki/HOWTO----Image-shrinking # http://jcupitt.github.io/libvips/API/current/Using-vipsthumbnail.md.html - def self.resize_ruby(file, width, height, resize_quality) + def resize_ruby(file, width, height, resize_quality) output_file = Tempfile.new resized_image = Vips::Image.thumbnail(file.path, width, height: height, **THUMBNAIL_OPTIONS) resized_image.jpegsave(output_file.path, Q: resize_quality, **JPEG_OPTIONS) @@ -31,7 +33,7 @@ module DanbooruImageResizer output_file end - def self.crop_ruby(file, width, height, resize_quality) + def crop_ruby(file, width, height, resize_quality) return nil unless Danbooru.config.enable_image_cropping output_file = Tempfile.new @@ -41,7 +43,7 @@ module DanbooruImageResizer output_file end - def self.resize_shell(file, width, height, quality) + def resize_shell(file, width, height, quality) output_file = Tempfile.new(["resize", ".jpg"]) # --size=WxH will upscale if the image is smaller than the target size. @@ -63,7 +65,7 @@ module DanbooruImageResizer output_file end - def self.crop_shell(file, width, height, quality) + def crop_shell(file, width, height, quality) return nil unless Danbooru.config.enable_image_cropping output_file = Tempfile.new(["crop", ".jpg"]) @@ -85,4 +87,20 @@ module DanbooruImageResizer output_file end + + def validate_shell(file) + temp = Tempfile.new("validate") + output, status = Open3.capture2e("vips stats #{file.path} #{temp.path}.v") + + # png | jpeg | gif + if output =~ /Read Error|Premature end of JPEG file|Failed to read from given file/m + return false + end + + return true + + ensure + temp.close + temp.unlink + end end diff --git a/app/logical/upload_service/utils.rb b/app/logical/upload_service/utils.rb index 644c41445..8abe24f46 100644 --- a/app/logical/upload_service/utils.rb +++ b/app/logical/upload_service/utils.rb @@ -1,6 +1,7 @@ class UploadService module Utils extend self + class CorruptFileError < RuntimeError; end def file_header_to_file_ext(file) case File.read(file.path, 16) @@ -23,9 +24,9 @@ class UploadService end end - def delete_file(md5, file_ext, upload_id = nil) + def delete_file(md5, file_ext, upload_id = nil) if Post.where(md5: md5).exists? - if upload_id + if upload_id.present? && Upload.where(id: upload_id).exists? CurrentUser.as_system do Upload.find(upload_id).update(status: "completed") end @@ -34,7 +35,7 @@ class UploadService return end - if upload_id && Upload.where(id: upload_id).exists? + if upload_id.present? && Upload.where(id: upload_id).exists? CurrentUser.as_system do Upload.find(upload_id).update(status: "preprocessed + deleted") end @@ -206,8 +207,24 @@ class UploadService end def download_for_upload(upload) - download = Downloads::File.new(upload.source, upload.referer_url) - file, strategy = download.download! + attempts = 0 + + begin + download = Downloads::File.new(upload.source, upload.referer_url) + file, strategy = download.download! + + if !DanbooruImageResizer.validate_shell(file) + raise CorruptFileError.new("File is corrupted") + end + + rescue + if attempts == 3 + raise + end + + attempts += 1 + retry + end if download.data[:ugoira_frame_data] upload.context = { diff --git a/test/files/test-corrupt.jpg b/test/files/test-corrupt.jpg new file mode 100644 index 000000000..8703c5d59 Binary files /dev/null and b/test/files/test-corrupt.jpg differ diff --git a/test/models/upload_service_test.rb b/test/models/upload_service_test.rb index 1d5cc4932..4330dec31 100644 --- a/test/models/upload_service_test.rb +++ b/test/models/upload_service_test.rb @@ -42,6 +42,28 @@ class UploadServiceTest < ActiveSupport::TestCase end end + context "for a corrupt jpeg" do + setup do + @source = "https://raikou1.donmai.us/93/f4/93f4dd66ef1eb11a89e56d31f9adc8d0.jpg" + @mock_upload = mock("upload") + @mock_upload.stubs(:source).returns(@source) + @mock_upload.stubs(:referer_url).returns(nil) + @bad_file = File.open("#{Rails.root}/test/files/test-corrupt.jpg", "rb") + Downloads::File.any_instance.stubs(:download!).returns([@bad_file, nil]) + end + + teardown do + @bad_file.close + end + + should "retry three times" do + DanbooruImageResizer.expects(:validate_shell).times(4).returns(false) + assert_raise(UploadService::Utils::CorruptFileError) do + subject.download_for_upload(@mock_upload) + end + end + end + context "for a pixiv" do setup do @source = "http://www.pixiv.net/member_illust.php?mode=medium&illust_id=62247350" @@ -849,7 +871,8 @@ class UploadServiceTest < ActiveSupport::TestCase context "a post that is replaced to another file then replaced back to the original file" do should "not delete the original files" do begin - FileUtils.expects(:rm_f).never + # this is called thrice to delete the file for 62247364 + FileUtils.expects(:rm_f).times(3) as_user do @post.replace!(replacement_url: "https://www.pixiv.net/member_illust.php?mode=medium&illust_id=62247350")