From a6efaa54a165f986470f3e2afb1d9abae4577f34 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Sat, 26 Oct 2019 14:04:19 -0500
Subject: [PATCH] Downloads::File: fix following 302 redirects during download.

Fixes downloading yande.re preview images for iqdb. Yande.re previews
return a 302 redirect to the actual file. Before the html in the body of
the 302 response would get prepended to the image file.

https://files.yande.re/data/preview/12/8f/128fb480d8589be26d1dce7e1d841bcb.jpg
=> https://assets.yande.re/data/preview/12/8f/128fb480d8589be26d1dce7e1d841bcb.jpg
---
 app/logical/downloads/file.rb    | 10 ++++++++--
 test/unit/downloads/file_test.rb | 17 ++++++++++++++---
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/app/logical/downloads/file.rb b/app/logical/downloads/file.rb
index e4fa2d128..d5db778be 100644
--- a/app/logical/downloads/file.rb
+++ b/app/logical/downloads/file.rb
@@ -26,9 +26,9 @@ module Downloads
       end
     end
 
-    def download!(tries: 3, **options)
+    def download!(url: uncached_url, tries: 3, **options)
       Retriable.retriable(on: RETRIABLE_ERRORS, tries: tries, base_interval: 0) do
-        file = http_get_streaming(uncached_url, headers: strategy.headers, **options)
+        file = http_get_streaming(url, headers: strategy.headers, **options)
         return [file, strategy]
       end
     end
@@ -43,6 +43,8 @@ module Downloads
       size = 0
 
       res = HTTParty.get(url, httparty_options) do |chunk|
+        next if chunk.code == 302
+
         size += chunk.size
         raise Error.new("File is too large (max size: #{max_size})") if size > max_size && max_size > 0
 
@@ -66,6 +68,10 @@ module Downloads
       url
     end
 
+    def preview_url
+      @preview_url ||= Addressable::URI.parse(strategy.preview_url)
+    end
+
     def file_url
       @file_url ||= Addressable::URI.parse(strategy.image_url)
     end
diff --git a/test/unit/downloads/file_test.rb b/test/unit/downloads/file_test.rb
index 45d33e0e5..9c8db4d25 100644
--- a/test/unit/downloads/file_test.rb
+++ b/test/unit/downloads/file_test.rb
@@ -49,14 +49,19 @@ module Downloads
 
         should "return an uncorrupted file on the second try" do
           bomb = stub("bomb")
-          bomb.expects(:size).raises(IOError)
+          bomb.stubs(:code).raises(IOError)
           resp = stub("resp", success?: true)
 
-          HTTParty.expects(:get).twice.multiple_yields("a", bomb, "b", "c").then.multiple_yields("a", "b", "c").returns(resp)
+          chunk = stub("a")
+          chunk.stubs(:code).returns(200)
+          chunk.stubs(:size).returns(1)
+          chunk.stubs(:to_s).returns("a")
+
+          HTTParty.expects(:get).twice.multiple_yields(chunk, bomb).then.multiple_yields(chunk, chunk).returns(resp)
           @download.stubs(:is_cloudflare?).returns(false)
           tempfile, _ = @download.download!
 
-          assert_equal("abc", tempfile.read)
+          assert_equal("aa", tempfile.read)
         end
       end
 
@@ -70,6 +75,12 @@ module Downloads
         tempfile, strategy = @download.download!
         assert_operator(tempfile.size, :>, 0, "should have data")
       end
+
+      should "correctly save the file when following 302 redirects" do
+        download = Downloads::File.new("https://yande.re/post/show/578014")
+        file, strategy  = download.download!(url: download.preview_url)
+        assert_equal(19134, file.size)
+      end
     end
   end
 end