[IQDB] Add integration for the new iqdb server

Use param v2=1 to use it, for now

Procfile

@@ -1,4 +1,4 @@
 server: bin/rails server -p 9000 -b 0.0.0.0
 # server: bundle exec unicorn -c config/unicorn/development.rb
-jobs: SIDEKIQ_QUEUES="low_prio:1;tags:2;default:3;high_prio:5" bundle exec sidekiq
+jobs: SIDEKIQ_QUEUES="low_prio:1;iqdb_new:1;tags:2;default:3;high_prio:5" bundle exec sidekiq
 cron: run-parts /etc/periodic/daily && crond -f

app/controllers/iqdb_queries_controller.rb

@@ -3,39 +3,63 @@ class IqdbQueriesController < ApplicationController
   before_action :detect_xhr, :throttle
 
   def show
-    if params[:file]
+    if params[:v2].present?
-      @matches = IqdbProxy.query_file(params[:file].tempfile)
+      new_version
-    elsif params[:url].present?
+    else
-      parsed_url = Addressable::URI.heuristic_parse(params[:url]) rescue nil
+      old_version
-      raise User::PrivilegeError.new("Invalid URL") unless parsed_url
-      whitelist_result = UploadWhitelist.is_whitelisted?(parsed_url)
-      raise User::PrivilegeError.new("Not allowed to request content from this URL") unless whitelist_result[0]
-      @matches = IqdbProxy.query(params[:url])
-    elsif params[:post_id]
-      @matches = IqdbProxy.query_path(Post.find(params[:post_id]).preview_file_path)
     end
 
     respond_with(@matches) do |fmt|
       fmt.html do |html|
-        html.xhr { render layout: false}
+        html.xhr { render layout: false }
       end
 
       fmt.json do
-        render json: @matches, root: 'posts'
+        render json: @matches, root: "posts"
       end
     end
-  rescue IqdbProxy::Error => e
+  rescue IqdbProxy::Error, IqdbProxyNew::Error => e
     render_expected_error(500, e.message)
   end
 
-private
+  private
+
+  def old_version
+    if params[:file]
+      @matches = IqdbProxy.query_file(params[:file].tempfile)
+    elsif params[:url].present?
+      parsed_url = Addressable::URI.heuristic_parse(params[:url]) rescue nil
+      raise User::PrivilegeError "Invalid URL" unless parsed_url
+      whitelist_result = UploadWhitelist.is_whitelisted?(parsed_url)
+      raise User::PrivilegeError "Not allowed to request content from this URL" unless whitelist_result[0]
+      @matches = IqdbProxy.query(params[:url])
+    elsif params[:post_id]
+      @matches = IqdbProxy.query_path(Post.find(params[:post_id]).preview_file_path)
+    end
+  end
+
+  def new_version
+    if params[:file]
+      @matches = IqdbProxyNew.query_file(params[:file].tempfile)
+    elsif params[:url].present?
+      parsed_url = Addressable::URI.heuristic_parse(params[:url]) rescue nil
+      raise User::PrivilegeError, "Invalid URL" unless parsed_url
+      whitelist_result = UploadWhitelist.is_whitelisted?(parsed_url)
+      raise User::PrivilegeError, "Not allowed to request content from this URL" unless whitelist_result[0]
+      @matches = IqdbProxyNew.query_url(params[:url])
+    elsif params[:post_id]
+      @matches = IqdbProxyNew.query_post(Post.find(params[:post_id]))
+    elsif params[:hash]
+      @matches = IqdbProxyNew.query_hash(params[:hash])
+    end
+  end
 
   def throttle
     if params[:file] || params[:url] || params[:post_id]
-      unless RateLimiter.check_limit("img:#{CurrentUser.ip_addr}", 1, 2.seconds)
+      if RateLimiter.check_limit("img:#{CurrentUser.ip_addr}", 1, 2.seconds) && !Danbooru.config.disable_throttles?
-        RateLimiter.hit("img:#{CurrentUser.ip_addr}", 2.seconds)
+        raise APIThrottled
       else
-        raise APIThrottled.new
+        RateLimiter.hit("img:#{CurrentUser.ip_addr}", 2.seconds)
       end
     end
   end

app/jobs/iqdb_remove_job_new.rb

@@ -0,0 +1,7 @@
+class IqdbRemoveJobNew < ApplicationJob
+  queue_as :iqdb_new
+
+  def perform(post_id)
+    IqdbProxyNew.remove_post(post_id)
+  end
+end

app/jobs/iqdb_update_job_new.rb

@@ -0,0 +1,10 @@
+class IqdbUpdateJobNew < ApplicationJob
+  queue_as :iqdb_new
+
+  def perform(post_id)
+    post = Post.find_by id: post_id
+    return unless post
+
+    IqdbProxyNew.update_post(post)
+  end
+end

app/logical/iqdb_proxy_new.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module IqdbProxyNew
+  class Error < StandardError; end
+
+  IQDB_NUM_PIXELS = 128
+
+  module_function
+
+  def make_request(path, request_type, params = {})
+    url = URI.parse(Danbooru.config.iqdb_server)
+    url.path = path
+    HTTParty.send(request_type, url, { body: params.to_json, headers: { "Content-Type" => "application/json" } })
+  rescue Errno::ECONNREFUSED, Errno::EADDRNOTAVAIL
+    raise Error, "This service is temporarily unavailable. Please try again later."
+  end
+
+  def update_post(post)
+    return unless post.has_preview?
+
+    thumb = generate_thumbnail(post.preview_file_path)
+    raise Error, "failed to generate thumb for #{post.id}" unless thumb
+
+    response = make_request("/images/#{post.id}", :post, get_channels_data(thumb))
+    raise Error, "iqdb request failed" if response.code != 200
+  end
+
+  def remove_post(post_id)
+    response = make_request("/images/#{post_id}", :delete)
+    raise Error, "iqdb request failed" if response.code != 200
+  end
+
+  def query_url(image_url)
+    file, _strategy = Downloads::File.new(image_url).download!
+    query_file(file)
+  end
+
+  def query_post(post)
+    return [] unless post.has_preview?
+
+    File.open(post.preview_file_path) do |f|
+      query_file(f)
+    end
+  end
+
+  def query_file(file)
+    thumb = generate_thumbnail(file.path)
+    return [] unless thumb
+
+    response = make_request("/query", :post, get_channels_data(thumb))
+    return [] if response.code != 200
+
+    process_iqdb_result(response.parsed_response)
+  end
+
+  def query_hash(hash)
+    response = make_request "/query", :post, { hash: hash }
+    return [] if response.code != 200
+
+    process_iqdb_result(response.parsed_response)
+  end
+
+  def process_iqdb_result(json, score_cutoff = 80)
+    raise Error, "Server returned an error. Most likely the url is not found." unless json.is_a?(Array)
+
+    json.filter! { |entry| (entry["score"] || 0) >= score_cutoff }
+    json.map do |x|
+      x["post"] = Post.find(x["post_id"])
+      x
+    rescue ActiveRecord::RecordNotFound
+      nil
+    end.compact
+  end
+
+  def generate_thumbnail(file_path)
+    Vips::Image.thumbnail(file_path, IQDB_NUM_PIXELS, height: IQDB_NUM_PIXELS, size: :force)
+  rescue Vips::Error
+    nil
+  end
+
+  def get_channels_data(thumbnail)
+    r = []
+    g = []
+    b = []
+    is_grayscale = thumbnail.bands == 1
+    thumbnail.to_a.each do |data|
+      data.each do |rgb|
+        r << rgb[0]
+        g << (is_grayscale ? rgb[0] : rgb[1])
+        b << (is_grayscale ? rgb[0] : rgb[2])
+      end
+    end
+    { channels: { r: r, g: g, b: b } }
+  end
+end

app/models/post.rb

@@ -1551,14 +1551,15 @@ class Post < ApplicationRecord
       def remove_iqdb(post_id)
         if iqdb_enabled?
           IqdbRemoveJob.perform_async(post_id)
+          IqdbRemoveJobNew.perform_later(post_id)
         end
       end
     end
 
     def update_iqdb_async
       if Post.iqdb_enabled? && has_preview?
-        # IqdbUpdateJob.perform_async(id, preview_file_url)
         IqdbUpdateJob.perform_async(id, "md5:#{md5}.jpg")
+        IqdbUpdateJobNew.perform_later(id)
       end
     end
 

config/danbooru_default_config.rb

@@ -723,6 +723,9 @@ module Danbooru
     def iqdbs_server
     end
 
+    def iqdb_server
+    end
+
     def elasticsearch_host
       '127.0.0.1'
     end

docker-compose.yml

@@ -6,6 +6,7 @@ x-environment: &common-env
   DANBOORU_ELASTICSEARCH_HOST: elastic
   DANBOORU_MEMCACHED_SERVERS: memcached
   DANBOORU_IQDBS_SERVER: http://iqdb:4567
+  DANBOORU_IQDB_SERVER: http://iqdb_new:5588
   # These are just development secrets, do not use them in production
   SECRET_TOKEN: 1c58518a891eff4520cadc59afa9e378a9325f1247544ff258096e497f095f45
   SESSION_SECRET_KEY: 44b4f44e9f253c406cbe727d403d500c1cecff943e4d2aea8f5447f28846fffe
@@ -123,6 +124,12 @@ services:
       - post_data:/data
       - iqdb_data:/iqdb
 
+  iqdb_new:
+    image: ghcr.io/e621ng/iqdb:d4fed9d9a51184e72d2f14d4ec461d7830bd177a 
+    command: iqdb http 0.0.0.0 5588 /iqdb/e621_v2.db
+    volumes:
+      - iqdb_data:/iqdb
+
   # Useful for development
 
   tests:

test/functional/moderator/post/posts_controller_test.rb

@@ -67,7 +67,7 @@ module Moderator
 
             post_auth move_favorites_moderator_post_post_path(@child.id), @admin, params: { commit: "Submit" }
             assert_redirected_to(@child)
-            perform_enqueued_jobs
+            perform_enqueued_jobs(only: TransferFavoritesJob)
             @parent.reload
             @child.reload
             as(@admin) do