From f46aaad39d37288a53924a6cfa6f456cdba9f6f9 Mon Sep 17 00:00:00 2001
From: Albert Yi <ayi@desk.com>
Date: Thu, 1 Dec 2016 12:22:34 -0800
Subject: [PATCH] additional fixes to support iqdbs

---
 app/logical/iqdb/download.rb                  | 22 +++++++++++++++++--
 app/views/iqdb_queries/create_by_url.html.erb |  2 +-
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/app/logical/iqdb/download.rb b/app/logical/iqdb/download.rb
index c3c49b351..bf5ab06c2 100644
--- a/app/logical/iqdb/download.rb
+++ b/app/logical/iqdb/download.rb
@@ -6,11 +6,24 @@ module Iqdb
       @source = source
     end
 
+    def get_referer(url)
+      headers = {}
+      datums = {}
+
+      Downloads::RewriteStrategies::Base.strategies.each do |strategy|
+        url, headers, datums = strategy.new(url).rewrite(url, headers, datums)
+      end
+
+      [url, headers["Referer"]]
+    end
+
     def find_similar
       if Danbooru.config.iqdbs_server
+        url, ref = get_referer(source)
         params = {
           "key" => Danbooru.config.iqdbs_auth_key,
-          "url" => source
+          "url" => url,
+          "ref" => ref
         }
         uri = URI.parse("#{Danbooru.config.iqdbs_server}/similar")
         uri.query = URI.encode_www_form(params)
@@ -18,7 +31,12 @@ module Iqdb
         Net::HTTP.start(uri.host, uri.port) do |http|
           resp = http.request_get(uri.request_uri)
           if resp.is_a?(Net::HTTPSuccess)
-            JSON.parse(resp.body)
+            json = JSON.parse(resp.body)
+            if json.is_a?(Array)
+              json
+            else
+              []
+            end
           else
             raise "HTTP error code: #{resp.code} #{resp.message}"
           end
diff --git a/app/views/iqdb_queries/create_by_url.html.erb b/app/views/iqdb_queries/create_by_url.html.erb
index 8a5aa5ae3..b33993d41 100644
--- a/app/views/iqdb_queries/create_by_url.html.erb
+++ b/app/views/iqdb_queries/create_by_url.html.erb
@@ -1,7 +1,7 @@
 <% if @results.any? %>
   <h3>Similar</h3>
   <% @results.each do |match| %>
-    <%= PostPresenter.preview(Post.find(match.post_id), :tags => "status:any", :size => true) %>
+    <%= PostPresenter.preview(Post.find(match["post_id"]), :tags => "status:any", :size => true) %>
   <% end %>
 <% else %>
   <h3>Similar</h3>