From fa37b1edcdc4dee8fd06f7d2c9555e5822c2bb3d Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Mon, 14 Oct 2019 21:16:04 -0500
Subject: [PATCH] iqdb: lower similarity cutoff, return more results (fix
 #4190).

* Change cutoffs on upload page to max 5 results, min. 20% similarity.
* Change cutoffs on standalone /iqdb_queries page to max 20 results, min. 0% similarity.
* /iqdb_queries.json: add `limit` and `similarity` params to change default cutoffs.
---
 app/controllers/iqdb_queries_controller.rb      | 15 ++++++---------
 app/javascript/src/javascripts/uploads.js.erb   |  2 +-
 app/logical/iqdb_proxy.rb                       | 15 +++++++++------
 test/functional/iqdb_queries_controller_test.rb | 14 +++++++++-----
 4 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/app/controllers/iqdb_queries_controller.rb b/app/controllers/iqdb_queries_controller.rb
index a8bf9439f..ab1952015 100644
--- a/app/controllers/iqdb_queries_controller.rb
+++ b/app/controllers/iqdb_queries_controller.rb
@@ -3,15 +3,12 @@ class IqdbQueriesController < ApplicationController
 
   def show
     if params[:url]
-      strategy = Sources::Strategies.find(params[:url])
-      @matches = IqdbProxy.query(strategy.image_url)
-    end
-
-    if params[:post_id]
-      @matches = IqdbProxy.query(Post.find(params[:post_id]).preview_file_url)
-    end
-
-    if params[:matches]
+      url = Sources::Strategies.find(params[:url]).image_url
+      @matches = IqdbProxy.query(url, params[:limit], params[:similarity])
+    elsif params[:post_id]
+      url = Post.find(params[:post_id]).preview_file_url
+      @matches = IqdbProxy.query(url, params[:limit], params[:similarity])
+    elsif params[:matches]
       @matches = IqdbProxy.decorate_posts(JSON.parse(params[:matches]))
     end
 
diff --git a/app/javascript/src/javascripts/uploads.js.erb b/app/javascript/src/javascripts/uploads.js.erb
index 83bdc2407..d89ebcd20 100644
--- a/app/javascript/src/javascripts/uploads.js.erb
+++ b/app/javascript/src/javascripts/uploads.js.erb
@@ -68,7 +68,7 @@ Upload.validate_upload = function (e) {
 
 Upload.initialize_iqdb_source = function() {
   if (/^https?:\/\//.test($("#upload_source").val())) {
-    $.get("/iqdb_queries.js", { url: $("#upload_source").val() });
+    $.get("/iqdb_queries.js", { url: $("#upload_source").val(), limit: 5, similarity: 20 });
   }
 }
 
diff --git a/app/logical/iqdb_proxy.rb b/app/logical/iqdb_proxy.rb
index 0d8080134..71fe5da3f 100644
--- a/app/logical/iqdb_proxy.rb
+++ b/app/logical/iqdb_proxy.rb
@@ -1,12 +1,15 @@
 class IqdbProxy
-  def self.query(image_url)
+  def self.query(url, limit, similarity)
     raise NotImplementedError unless Danbooru.config.iqdbs_server.present?
 
-    url = URI.parse(Danbooru.config.iqdbs_server)
-    url.path = "/similar"
-    url.query = {url: image_url}.to_query
-    json = HTTParty.get(url.to_s, Danbooru.config.httparty_options).parsed_response
-    decorate_posts(json)
+    limit ||= 20
+    similarity ||= 0.0
+    query = { url: url, limit: limit }
+    response = HTTParty.get("#{Danbooru.config.iqdbs_server}/similar", query: query, **Danbooru.config.httparty_options)
+
+    json = decorate_posts(response.parsed_response)
+    json = json.select { |result| result["score"] >= similarity.to_f }.take(limit.to_i)
+    json
   end
 
   def self.decorate_posts(json)
diff --git a/test/functional/iqdb_queries_controller_test.rb b/test/functional/iqdb_queries_controller_test.rb
index 04a547f16..17c784462 100644
--- a/test/functional/iqdb_queries_controller_test.rb
+++ b/test/functional/iqdb_queries_controller_test.rb
@@ -23,8 +23,10 @@ class IqdbQueriesControllerTest < ActionDispatch::IntegrationTest
         end
 
         should "render a response" do
-          IqdbProxy.expects(:query).with(@url).returns(@mocked_response)
-          get_auth iqdb_queries_path(variant: "xhr"), @user, params: @params
+          IqdbProxy.expects(:query).returns(@mocked_response)
+          get_auth iqdb_queries_path, @user, as: :javascript, params: @params
+
+          assert_response :success
           assert_select("#post_#{@posts[0].id}")
         end
       end
@@ -41,22 +43,24 @@ class IqdbQueriesControllerTest < ActionDispatch::IntegrationTest
         end
 
         should "redirect to iqdbs" do
-          IqdbProxy.expects(:query).with(@posts[0].preview_file_url).returns(@mocked_response)
+          IqdbProxy.expects(:query).returns(@mocked_response)
           get_auth iqdb_queries_path, @user, params: @params
+
+          assert_response :success
           assert_select("#post_#{@posts[0].id}")
         end
       end
 
       context "with matches" do
         setup do
-          json = @posts.map {|x| {"post_id" => x.id, "score" => 1}}.to_json          
+          json = @posts.map {|x| {"post_id" => x.id, "score" => 1}}.to_json
           @params = { matches: json }
         end
 
         should "render with matches" do
           get_auth iqdb_queries_path, @user, params: @params
           assert_response :success
-        end        
+        end
       end
     end
   end