Refactor sources

app/controllers/sources_controller.rb

@@ -1,20 +1,12 @@
 class SourcesController < ApplicationController
   respond_to :json, :xml
-  rescue_from Sources::Site::NoStrategyError, :with => :no_strategy
 
   def show
-    @source = Sources::Site.new(params[:url], :referer_url => params[:ref])
-    @source.get
+    @source = Sources::Strategies.find(params[:url], params[:ref])
 
     respond_with(@source.to_h) do |format|
       format.xml { render xml: @source.to_h.to_xml(root: "source") }
       format.json { render json: @source.to_h.to_json }
     end
   end
-
-protected
-
-  def no_strategy
-    render json: {message: "Unsupported site"}.to_json, status: 400
-  end
 end

app/controllers/uploads_controller.rb

@@ -5,7 +5,7 @@ class UploadsController < ApplicationController
 
   def new
     @upload_notice_wiki = WikiPage.titled(Danbooru.config.upload_notice_wiki_page).first
-    @upload, @post, @source, @normalized_url, @remote_size = UploadService::ControllerHelper.prepare(
+    @upload, @post, @source, @remote_size = UploadService::ControllerHelper.prepare(
       url: params[:url], ref: params[:ref]
     )
     respond_with(@upload)
@@ -43,7 +43,7 @@ def show
   end
 
   def preprocess
-    @upload, @post, @source, @normalized_url, @remote_size = UploadService::ControllerHelper.prepare(
+    @upload, @post, @source, @remote_size = UploadService::ControllerHelper.prepare(
       url: params[:url], file: params[:file], ref: params[:ref]
     )
     render body: nil

app/javascript/src/javascripts/related_tag.js.erb

@@ -289,7 +289,7 @@ RelatedTag.toggle_tag = function(e) {
 RelatedTag.find_artist = function(e) {
   $("#artist-tags").html("<em>Loading...</em>");
   var url = $("#upload_source,#post_source");
-  var referer_url = $("#upload_referer_url");
+  var referer_url = $("#ref");
   $.get("/artists/finder.json", {"url": url.val(), "referer_url": referer_url.val()}, RelatedTag.process_artist);
   if (e) {
     e.preventDefault();

app/javascript/src/javascripts/uploads.js

@@ -87,7 +87,7 @@ Upload.initialize_info_bookmarklet = function() {
 Upload.initialize_info_manual = function() {
   $("#fetch-data-manual").click(function(e) {
     var source = $("#upload_source,#post_source").val();
-    var referer = $("#upload_referer_url").val();
+    var referer = $("#ref").val();
 
     if (/^https?:\/\//.test(source)) {
       $("#source-info span#loading-data").show();

app/logical/downloads/file.rb

@@ -3,16 +3,33 @@ class File
     class Error < Exception ; end
 
     attr_reader :data, :options
-    attr_accessor :source, :original_source, :downloaded_source
+    attr_accessor :source, :referer
 
-    def initialize(source, options = {})
+    # Prevent Cloudflare from potentially mangling the image. See issue #3528.
+    def self.uncached_url(url, headers = {})
+      url = Addressable::URI.parse(url)
+
+      if is_cloudflare?(url, headers)
+        url.query_values = (url.query_values || {}).merge(danbooru_no_cache: SecureRandom.uuid)
+      end
+
+      url
+    end
+
+    def self.is_cloudflare?(url, headers = {})
+      Cache.get("is_cloudflare:#{url.origin}", 4.hours) do
+        res = HTTParty.head(url, { headers: headers }.deep_merge(Danbooru.config.httparty_options))
+        raise Error.new("HTTP error code: #{res.code} #{res.message}") unless res.success?
+
+        res.key?("CF-Ray")
+      end
+    end
+
+    def initialize(source, referer=nil, options = {})
       # source can potentially get rewritten in the course
       # of downloading a file, so check it again
       @source = source
-      @original_source = source
-
-      # the URL actually downloaded after rewriting the original source.
-      @downloaded_source = nil
+      @referer = referer
 
       # we sometimes need to capture data from the source page
       @data = {}
@@ -22,48 +39,31 @@ def initialize(source, options = {})
       @data[:get_thumbnail] = options[:get_thumbnail]
     end
 
-    def rewrite_url
-      url, _, _ = before_download(@source, @data)
-      return url
-    end
-
     def size
-      url, headers, _ = before_download(@source, @data)
-      options = { timeout: 3, headers: headers }.deep_merge(Danbooru.config.httparty_options)
-      res = HTTParty.head(url, options)
-      res.content_length
-    end
+      strategy = Sources::Strategies.find(source, referer)
+      options = { timeout: 3, headers: strategy.headers }.deep_merge(Danbooru.config.httparty_options)
 
-    def download!
-      url, headers, @data = before_download(@source, @data)
+      res = HTTParty.head(strategy.file_url, options)
 
-      output_file = Tempfile.new(binmode: true)
-      http_get_streaming(uncached_url(url, headers), output_file, headers)
-
-      @downloaded_source = url
-      @source = after_download(url)
-
-      output_file
+      if res.success?
+        res.content_length
+      else
+        raise HTTParty::ResponseError.new(res)
+      end
     end
 
-    def before_download(url, datums)
-      original_url = url
-      headers = Danbooru.config.http_headers
-
-      RewriteStrategies::Base.strategies.each do |strategy|
-        url, headers, datums = strategy.new(url).rewrite(url, headers, datums)
-        url = original_url if url.nil?
-      end
+    def download!
+      strategy = Sources::Strategies.find(source, referer)
+      output_file = Tempfile.new(binmode: true)
+      @data = strategy.data
 
-      return [url, headers, datums]
-    end
+      http_get_streaming(
+        self.class.uncached_url(strategy.file_url, strategy.headers), 
+        output_file, 
+        strategy.headers
+      )
 
-    def after_download(src)
-      src = fix_twitter_sources(src)
-      if options[:referer_url].present?
-        src = set_source_to_referer(src, options[:referer_url])
-      end
-      src
+      [output_file, strategy]
     end
 
     def validate_local_hosts(url)
@@ -111,50 +111,5 @@ def http_get_streaming(src, file, headers = {}, max_size: Danbooru.config.max_fi
         end
       end # while
     end # def
-
-    def fix_twitter_sources(src)
-      if src =~ %r!^https?://(?:video|pbs)\.twimg\.com/! && original_source =~ %r!^https?://twitter\.com/!
-        original_source
-      elsif src =~ %r!^https?://img\.pawoo\.net/! && original_source =~ %r!^https?://pawoo\.net/!
-        original_source
-      else
-        src
-      end
-    end
-
-    def set_source_to_referer(src, referer)
-      if Sources::Strategies::Nijie.url_match?(src) ||
-         Sources::Strategies::Twitter.url_match?(src) || Sources::Strategies::Twitter.url_match?(referer) ||
-         Sources::Strategies::Pawoo.url_match?(src) ||
-         Sources::Strategies::Tumblr.url_match?(src) || Sources::Strategies::Tumblr.url_match?(referer) ||
-         Sources::Strategies::ArtStation.url_match?(src) || Sources::Strategies::ArtStation.url_match?(referer)
-        strategy = Sources::Site.new(src, :referer_url => referer)
-        strategy.referer_url
-      else
-        src
-      end
-    end
-
-    private
-
-    # Prevent Cloudflare from potentially mangling the image. See issue #3528.
-    def uncached_url(url, headers = {})
-      url = Addressable::URI.parse(url)
-
-      if is_cloudflare?(url, headers)
-        url.query_values = (url.query_values || {}).merge(danbooru_no_cache: SecureRandom.uuid)
-      end
-
-      url
-    end
-
-    def is_cloudflare?(url, headers = {})
-      Cache.get("is_cloudflare:#{url.origin}", 4.hours) do
-        res = HTTParty.head(url, { headers: headers }.deep_merge(Danbooru.config.httparty_options))
-        raise Error.new("HTTP error code: #{res.code} #{res.message}") unless res.success?
-
-        res.key?("CF-Ray")
-      end
-    end
   end
 end