diff --git a/app/logical/danbooru_image_resizer.rb b/app/logical/danbooru_image_resizer.rb index 34db067b3..9321ccc2d 100644 --- a/app/logical/danbooru_image_resizer.rb +++ b/app/logical/danbooru_image_resizer.rb @@ -35,19 +35,4 @@ 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 f13a785a3..e8939122b 100644 --- a/app/logical/upload_service/utils.rb +++ b/app/logical/upload_service/utils.rb @@ -2,8 +2,6 @@ class UploadService module Utils module_function - class CorruptFileError < RuntimeError; end - def file_header_to_file_ext(file) case File.read(file.path, 16) when /^\xff\xd8/n @@ -25,6 +23,14 @@ class UploadService end end + def corrupt?(filename) + image = Vips::Image.new_from_file(filename, fail: true) + image.stats + false + rescue Vips::Error + true + end + def calculate_ugoira_dimensions(source_path) folder = Zip::File.new(source_path) Tempfile.open("ugoira-dim-") do |tempfile| @@ -180,23 +186,8 @@ class UploadService return file if file.present? raise "No file or source URL provided" if upload.source_url.blank? - attempts = 0 - - begin - download = Downloads::File.new(upload.source_url, upload.referer_url) - file, strategy = download.download! - - if !DanbooruImageResizer.validate_shell(file) - raise CorruptFileError.new("File is corrupted") - end - rescue StandardError - if attempts == 3 - raise - end - - attempts += 1 - retry - end + download = Downloads::File.new(upload.source_url, upload.referer_url) + file, strategy = download.download! if download.data[:ugoira_frame_data].present? upload.context = { diff --git a/app/models/upload.rb b/app/models/upload.rb index d5df336c4..a9d4e5fd4 100644 --- a/app/models/upload.rb +++ b/app/models/upload.rb @@ -6,6 +6,7 @@ class Upload < ApplicationRecord class FileValidator < ActiveModel::Validator def validate(record) validate_file_ext(record) + validate_integrity(record) validate_md5_uniqueness(record) validate_video_duration(record) validate_resolution(record) @@ -17,6 +18,12 @@ class Upload < ApplicationRecord end end + def validate_integrity(record) + if record.file_ext.in?(["jpg", "gif", "png"]) && UploadService::Utils.corrupt?(record.file.path) + record.errors[:file] << "File is corrupted" + end + end + def validate_md5_uniqueness(record) if record.md5.nil? return diff --git a/test/files/test-corrupt.gif b/test/files/test-corrupt.gif new file mode 100644 index 000000000..e1ab46117 Binary files /dev/null and b/test/files/test-corrupt.gif differ diff --git a/test/files/test-corrupt.png b/test/files/test-corrupt.png new file mode 100644 index 000000000..1a81abef8 Binary files /dev/null and b/test/files/test-corrupt.png differ diff --git a/test/test_helpers/upload_test_helper.rb b/test/test_helpers/upload_test_helper.rb index 66b404278..4f2d18dc9 100644 --- a/test/test_helpers/upload_test_helper.rb +++ b/test/test_helpers/upload_test_helper.rb @@ -1,4 +1,8 @@ module UploadTestHelper + def upload_from_file(filepath) + UploadService.new(file: upload_file(filepath)).start! + end + def upload_file(path) file = Tempfile.new(binmode: true) IO.copy_stream("#{Rails.root}/#{path}", file.path) diff --git a/test/unit/upload_service_test.rb b/test/unit/upload_service_test.rb index f8382dd86..d99360175 100644 --- a/test/unit/upload_service_test.rb +++ b/test/unit/upload_service_test.rb @@ -34,28 +34,6 @@ class UploadServiceTest < ActiveSupport::TestCase end end - context "for a corrupt jpeg" do - setup do - @source = "https://cdn.donmai.us/original/93/f4/93f4dd66ef1eb11a89e56d31f9adc8d0.jpg" - @mock_upload = mock("upload") - @mock_upload.stubs(:source_url).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.get_file_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" @@ -1133,6 +1111,33 @@ class UploadServiceTest < ActiveSupport::TestCase assert_equal("http://www.example.com", @upload.source) end end + + context "for a corrupted image" do + should "fail for a corrupted jpeg" do + @bad_jpeg_path = "test/files/test-corrupt.jpg" + + upload = upload_from_file(@bad_jpeg_path) + assert_match(/corrupt/, upload.status) + assert_equal(true, UploadService::Utils.corrupt?(@bad_jpeg_path)) + end + + should "fail for a corrupted gif" do + @bad_gif_path = "test/files/test-corrupt.gif" + + upload = upload_from_file(@bad_gif_path) + assert_match(/corrupt/, upload.status) + assert_equal(true, UploadService::Utils.corrupt?(@bad_gif_path)) + end + + # https://schaik.com/pngsuite/pngsuite_xxx_png.html + should "fail for a corrupted png" do + @bad_png_path = "test/files/test-corrupt.png" + + upload = upload_from_file(@bad_png_path) + assert_match(/corrupt/, upload.status) + assert_equal(true, UploadService::Utils.corrupt?(@bad_png_path)) + end + end end context "#create_post_from_upload" do