From 21dcf53dcb111579d38135ca84444a3b22d3d805 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Fri, 28 Jan 2022 21:07:06 -0600
Subject: [PATCH] uploads: show similar images for disk uploads.

Fix the upload page so that it shows similar images (IQDB matches) for
files uploaded from your computer. Before this only worked for files
uploaded from a source.
---
 app/controllers/iqdb_queries_controller.rb |  2 +-
 app/javascript/src/javascripts/uploads.js  | 20 +++++++++-----------
 app/logical/iqdb_client.rb                 |  4 +++-
 app/models/upload.rb                       |  9 +++++++++
 app/policies/post_policy.rb                |  6 +++---
 app/views/uploads/_related_posts.html.erb  | 10 +---------
 app/views/uploads/show.html.erb            | 10 ++++++----
 7 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/app/controllers/iqdb_queries_controller.rb b/app/controllers/iqdb_queries_controller.rb
index 2e89db0d4..9e77438cf 100644
--- a/app/controllers/iqdb_queries_controller.rb
+++ b/app/controllers/iqdb_queries_controller.rb
@@ -5,7 +5,7 @@ class IqdbQueriesController < ApplicationController
 
   def show
     # XXX allow bare search params for backwards compatibility.
-    search_params.merge!(params.slice(:url, :hash, :image_url, :file_url, :post_id, :limit, :similarity, :high_similarity).permit!)
+    search_params.merge!(params.slice(:url, :hash, :image_url, :file_url, :post_id, :media_asset_id, :limit, :similarity, :high_similarity).permit!)
 
     iqdb_params = search_params.to_h.symbolize_keys
     @high_similarity_matches, @low_similarity_matches, @matches = IqdbClient.new.search(**iqdb_params)
diff --git a/app/javascript/src/javascripts/uploads.js b/app/javascript/src/javascripts/uploads.js
index 5d34a8e1c..87de7422e 100644
--- a/app/javascript/src/javascripts/uploads.js
+++ b/app/javascript/src/javascripts/uploads.js
@@ -26,18 +26,16 @@ Upload.initialize_all = function() {
 }
 
 Upload.initialize_similar = function() {
-  let source = $("#post_source").val();
+  let media_asset_id = $("input[name='post[media_asset_id]']").val();
 
-  if (/^https?:\/\//.test(source)) {
-    $.get("/iqdb_queries.js", {
-      limit: Upload.IQDB_LIMIT,
-      search: {
-        url: source,
-        similarity: Upload.IQDB_MIN_SIMILARITY,
-        high_similarity: Upload.IQDB_HIGH_SIMILARITY
-      }
-    });
-  }
+  $.get("/iqdb_queries.js", {
+    limit: Upload.IQDB_LIMIT,
+    search: {
+      media_asset_id: media_asset_id,
+      similarity: Upload.IQDB_MIN_SIMILARITY,
+      high_similarity: Upload.IQDB_HIGH_SIMILARITY
+    }
+  });
 }
 
 Upload.initialize_image = function() {
diff --git a/app/logical/iqdb_client.rb b/app/logical/iqdb_client.rb
index dc11d3d92..b4e8a45c7 100644
--- a/app/logical/iqdb_client.rb
+++ b/app/logical/iqdb_client.rb
@@ -18,7 +18,7 @@ class IqdbClient
 
   concerning :QueryMethods do
     # Search for an image by file, URL, hash, or post ID.
-    def search(post_id: nil, file: nil, hash: nil, url: nil, image_url: nil, file_url: nil, similarity: 0.0, high_similarity: 65.0, limit: 20)
+    def search(post_id: nil, media_asset_id: nil, file: nil, hash: nil, url: nil, image_url: nil, file_url: nil, similarity: 0.0, high_similarity: 65.0, limit: 20)
       limit = limit.to_i.clamp(1, 1000)
       similarity = similarity.to_f.clamp(0.0, 100.0)
       high_similarity = high_similarity.to_f.clamp(0.0, 100.0)
@@ -33,6 +33,8 @@ class IqdbClient
         file = download(file_url, :image_url)
       elsif post_id.present?
         file = Post.find(post_id).file(:preview)
+      elsif media_asset_id.present?
+        file = MediaAsset.find(media_asset_id).variant("360x360").open_file
       end
 
       if hash.present?
diff --git a/app/models/upload.rb b/app/models/upload.rb
index 654f37ac5..bb1d5ffa1 100644
--- a/app/models/upload.rb
+++ b/app/models/upload.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 class Upload < ApplicationRecord
+  extend Memoist
+
   MAX_VIDEO_DURATION = 140
 
   attr_accessor :file
@@ -76,7 +78,14 @@ class Upload < ApplicationRecord
     raise
   end
 
+  def source_strategy
+    return nil if source.blank?
+    Sources::Strategies.find(source, referer_url)
+  end
+
   def self.available_includes
     [:uploader, :upload_media_assets, :media_assets]
   end
+
+  memoize :source_strategy
 end
diff --git a/app/policies/post_policy.rb b/app/policies/post_policy.rb
index 84f596771..aa498d77c 100644
--- a/app/policies/post_policy.rb
+++ b/app/policies/post_policy.rb
@@ -71,9 +71,9 @@ class PostPolicy < ApplicationPolicy
   end
 
   def permitted_attributes_for_create
-    %i[upload_media_asset_id tag_string rating parent_id source is_pending
-    artist_commentary_desc artist_commentary_title translated_commentary_desc
-    translated_commentary_title]
+    %i[upload_id media_asset_id upload_media_asset_id tag_string rating
+    parent_id source is_pending artist_commentary_desc artist_commentary_title
+    translated_commentary_desc translated_commentary_title]
   end
 
   def permitted_attributes_for_update
diff --git a/app/views/uploads/_related_posts.html.erb b/app/views/uploads/_related_posts.html.erb
index 4a9c3e3e8..b62377377 100644
--- a/app/views/uploads/_related_posts.html.erb
+++ b/app/views/uploads/_related_posts.html.erb
@@ -16,12 +16,4 @@
   </div>
 <% end %>
 
-<% if Danbooru.config.iqdb_url %>
-  <% if params[:url] %>
-    <div class="input" id="iqdb-similar">
-      <p><em>Loading similar...</em></p>
-    </div>
-  <% else %>
-    <div class="input" id="iqdb-similar" style="display: none;"></div>
-  <% end %>
-<% end %>
+<div class="input" id="iqdb-similar" style="display: none;"></div>
diff --git a/app/views/uploads/show.html.erb b/app/views/uploads/show.html.erb
index 9608d3813..02612d334 100644
--- a/app/views/uploads/show.html.erb
+++ b/app/views/uploads/show.html.erb
@@ -45,13 +45,15 @@
         <%= tag.img src: @media_asset.variant("original").file_url, title: "Preview", id: "image", class: "fit-width fit-height", "data-shortcut": "z"  %>
       </div>
 
-      <% if @upload.source.present? %>
-        <% @source = ::Sources::Strategies.find(@upload.source, @upload.referer_url) %>
-        <%= render "uploads/related_posts", source: @source %>
-        <%= render_source_data(@source) %>
+      <%= render "uploads/related_posts", source: @upload.source_strategy %>
+
+      <% if @upload.source_strategy.present? %>
+        <%= render_source_data(@upload.source_strategy) %>
       <% end %>
 
       <%= edit_form_for(@post, html: { id: "form" }) do |f| %>
+        <%= f.input :upload_id, as: :hidden, input_html: { value: @upload.id } %>
+        <%= f.input :media_asset_id, as: :hidden, input_html: { value: @media_asset.id } %>
         <%= f.input :upload_media_asset_id, as: :hidden, input_html: { value: @upload_media_asset.id } %>
 
         <%= f.input :source, as: :string %>