[IQDB] Switch to new implementation

.env.sample

@@ -32,7 +32,6 @@
 # export DANBOORU_APP_NAME=
 # export DANBOORU_VERSION=
 # export DANBOORU_HOSTNAME=
-# export DANBOORU_IQDBS_SERVER=
 
 #
 # Development Only

.gitignore

@@ -6,26 +6,18 @@ config/danbooru_local_config.rb
 config/deploy/*.rb
 config/application.yml
 config/newrelic.yml
-db/*.sqlite3
 doc/*
 log/*.log
 tmp/*
 public/data
 public/assets
 vendor/cache/*.gem
-.sass-cache
 .DS_Store
 coverage
 *~
 *.swp
 tmp/*.jpg
 tmp/*.png
-danbooru.db
-iqdb.db
-danbooru.sublime-project
-danbooru.sublime-workspace
-script/scratch.*
-test/fixtures/vcr_cassettes
 .gitconfig
 test/reports
 

app/controllers/iqdb_queries_controller.rb

@@ -3,42 +3,6 @@ class IqdbQueriesController < ApplicationController
   before_action :detect_xhr, :throttle
 
   def show
-    if params[:v2].present?
-      new_version
-    else
-      old_version
-    end
-
-    respond_with(@matches) do |fmt|
-      fmt.html do |html|
-        html.xhr { render layout: false }
-      end
-
-      fmt.json do
-        render json: @matches, root: "posts"
-      end
-    end
-  rescue IqdbProxy::Error, IqdbProxyNew::Error => e
-    render_expected_error(500, e.message)
-  end
-
-  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 = iqdb_proxy(:query_file, params[:file].tempfile)
     elsif params[:url].present?
@@ -52,10 +16,24 @@ class IqdbQueriesController < ApplicationController
     elsif params[:hash]
       @matches = iqdb_proxy(:query_hash, params[:hash])
     end
+
+    respond_with(@matches) do |fmt|
+      fmt.html do |html|
+        html.xhr { render layout: false }
+      end
+
+      fmt.json do
+        render json: @matches, root: "posts"
+      end
+    end
+  rescue IqdbProxy::Error => e
+    render_expected_error(500, e.message)
   end
 
+  private
+
   def iqdb_proxy(method, value)
-    IqdbProxyNew.send(method, value, params[:score_cutoff])
+    IqdbProxy.send(method, value, params[:score_cutoff])
   end
 
   def throttle

app/jobs/iqdb_remove_job_new.rb

@@ -2,6 +2,6 @@ class IqdbRemoveJobNew < ApplicationJob
   queue_as :iqdb_new
 
   def perform(post_id)
-    IqdbProxyNew.remove_post(post_id)
+    IqdbProxy.remove_post(post_id)
   end
 end

app/jobs/iqdb_update_job_new.rb

@@ -5,6 +5,6 @@ class IqdbUpdateJobNew < ApplicationJob
     post = Post.find_by id: post_id
     return unless post
 
-    IqdbProxyNew.update_post(post)
+    IqdbProxy.update_post(post)
   end
 end

app/logical/iqdb_proxy.rb

@@ -1,24 +1,69 @@
-class IqdbProxy
+# frozen_string_literal: true
+
+module IqdbProxy
   class Error < StandardError; end
 
-  def self.query(image_url)
+  IQDB_NUM_PIXELS = 128
-    response = make_request(url: image_url)
+
-    decorate_posts(response)
+  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 self.query_file(image)
+  def update_post(post)
-    response = make_request(file: image)
+    return unless post.has_preview?
-    decorate_posts(response)
+
+    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 self.query_path(image_path)
+  def remove_post(post_id)
-    File.open(image_path) do |f|
+    response = make_request("/images/#{post_id}", :delete)
-      query_file(f)
+    raise Error, "iqdb request failed" if response.code != 200
+  end
+
+  def query_url(image_url, score_cutoff)
+    file, _strategy = Downloads::File.new(image_url).download!
+    query_file(file, score_cutoff)
+  end
+
+  def query_post(post, score_cutoff)
+    return [] unless post.has_preview?
+
+    File.open(post.preview_file_path) do |f|
+      query_file(f, score_cutoff)
     end
   end
 
-  def self.decorate_posts(json)
+  def query_file(file, score_cutoff)
+    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, score_cutoff)
+  end
+
+  def query_hash(hash, score_cutoff)
+    response = make_request "/query", :post, { hash: hash }
+    return [] if response.code != 200
+
+    process_iqdb_result(response.parsed_response, score_cutoff)
+  end
+
+  def process_iqdb_result(json, score_cutoff)
     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 || 80).to_i }
     json.map do |x|
       x["post"] = Post.find(x["post_id"])
       x
@@ -27,17 +72,24 @@ class IqdbProxy
     end.compact
   end
 
-  def self.make_request(**params)
+  def generate_thumbnail(file_path)
-    raise NotImplementedError if Danbooru.config.iqdbs_server.blank?
+    Vips::Image.thumbnail(file_path, IQDB_NUM_PIXELS, height: IQDB_NUM_PIXELS, size: :force)
+  rescue Vips::Error
+    nil
+  end
 
-    url = URI.parse(Danbooru.config.iqdbs_server)
+  def get_channels_data(thumbnail)
-    url.path = "/similar"
+    r = []
-
+    g = []
-    json = HTTParty.post(url.to_s, { body: params }.merge(Danbooru.config.httparty_options))
+    b = []
-    return {} if json.code != 200
+    is_grayscale = thumbnail.bands == 1
-
+    thumbnail.to_a.each do |data|
-    json.parsed_response
+      data.each do |rgb|
-  rescue Errno::ECONNREFUSED, Errno::EADDRNOTAVAIL
+        r << rgb[0]
-    raise Error, "This service is temporarily unavailable. Please try again later."
+        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/logical/iqdb_proxy_new.rb

@@ -1,95 +0,0 @@
-# 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, score_cutoff)
-    file, _strategy = Downloads::File.new(image_url).download!
-    query_file(file, score_cutoff)
-  end
-
-  def query_post(post, score_cutoff)
-    return [] unless post.has_preview?
-
-    File.open(post.preview_file_path) do |f|
-      query_file(f, score_cutoff)
-    end
-  end
-
-  def query_file(file, score_cutoff)
-    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, score_cutoff)
-  end
-
-  def query_hash(hash, score_cutoff)
-    response = make_request "/query", :post, { hash: hash }
-    return [] if response.code != 200
-
-    process_iqdb_result(response.parsed_response, score_cutoff)
-  end
-
-  def process_iqdb_result(json, score_cutoff)
-    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 || 80).to_i }
-    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

@@ -1545,7 +1545,7 @@ class Post < ApplicationRecord
 
     module ClassMethods
       def iqdb_enabled?
-        Danbooru.config.iqdbs_server.present?
+        Danbooru.config.iqdb_server.present?
       end
 
       def remove_iqdb(post_id)

app/views/posts/partials/show/_options.html.erb

@@ -48,7 +48,7 @@
         <% end %>
 
         <% if CurrentUser.is_admin? %>
-          <% if Danbooru.config.iqdbs_server.present? %>
+          <% if Danbooru.config.iqdb_server.present? %>
             <li><%= link_to "Update IQDB", update_iqdb_post_path(@post) %></li>
           <% end %>
           <li><%= tag.a "Destroy", href: '#', id: 'destroy-post-link', 'data-pid': post.id %></li>

config/danbooru_default_config.rb

@@ -720,9 +720,6 @@ module Danbooru
       true
     end
 
-    def iqdbs_server
-    end
-
     def iqdb_server
     end
 

docker-compose.yml

@@ -5,8 +5,7 @@ x-environment: &common-env
   DANBOORU_REDIS_URL: redis://redis
   DANBOORU_ELASTICSEARCH_HOST: elastic
   DANBOORU_MEMCACHED_SERVERS: memcached
-  DANBOORU_IQDBS_SERVER: http://iqdb:4567
+  DANBOORU_IQDB_SERVER: http://iqdb:5588
-  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
@@ -105,26 +104,8 @@ services:
       interval: 10s
       timeout: 2s
       test: nc -z elastic 9200
-  
-  iqdb:
-    build: ./docker/iqdb
-    environment:
-      - IQDB_HOSTNAME=localhost
-      - IQDB_PORT=62125
-      # Hardcoded in iqdbs Procfile
-      - IQDB_DATABASE_FILE=/iqdb/e621.db
-      - SINATRA_PORT=4567
-      - IMAGES_FOLDER=/data
-      - REDIS_URL=redis://redis
-      # Sinatra only binds to localhost in development, but to 0.0.0.0 for anything else
-      - APP_ENV=docker
-    depends_on:
-      - redis
-    volumes:
-      - post_data:/data
-      - iqdb_data:/iqdb
 
-  iqdb_new:
+  iqdb:
     image: ghcr.io/e621ng/iqdb:d4fed9d9a51184e72d2f14d4ec461d7830bd177a 
     command: iqdb http 0.0.0.0 5588 /iqdb/e621_v2.db
     volumes:

docker/iqdb/Dockerfile

@@ -1,34 +0,0 @@
-FROM alpine:3.17 as builder
-
-COPY iqdb.patch /tmp/iqdb.patch
-ENV IQDB_VERSION=20161008
-ENV IQDB_CHECKSUM=daa4812b35f84bc7e2f07661fd9abf73a06738c76beb92252666a86ebaea3c64
-ARG BUILD_DEPS="build-base libjpeg-turbo-dev gd-dev"
-
-RUN apk --no-cache add $BUILD_DEPS \
-  && cd /tmp \
-  && wget -q  https://iqdb.org/code/iqdb-$IQDB_VERSION.tar.bz2 \
-  && echo "$IQDB_CHECKSUM  iqdb-$IQDB_VERSION.tar.bz2" | sha256sum -c - \
-  && tar xjf iqdb-$IQDB_VERSION.tar.bz2 \
-  && cd iqdb \
-  && patch -N -i /tmp/iqdb.patch \
-  && make EXTRADEFS="-include stdint.h" -j$(nproc) \
-  && apk del $BUILD_DEPS
-
-FROM ruby:3.2.2-alpine3.17
-
-RUN apk --no-cache add git libjpeg-turbo gd tzdata
-
-WORKDIR /iqdbs
-RUN bundle config without production
-
-ADD https://api.github.com/repos/e621ng/iqdbs/git/refs/heads/master /tmp/iqdbs_version.json
-RUN git clone https://github.com/e621ng/iqdbs /iqdbs \
-  && bundler install --gemfile /iqdbs/Gemfile \
-  && gem i foreman
-
-COPY --from=builder /tmp/iqdb/iqdb /usr/bin/
-
-WORKDIR /iqdbs
-
-CMD ["foreman", "start", "-f", "Procfile.dev"]

docker/iqdb/iqdb.patch

@@ -1,13 +0,0 @@
-diff --git a/imglib.h b/imglib.h
-index 2a03f93..e5468c4 100644
---- a/imglib.h
-+++ b/imglib.h
-@@ -98,7 +98,7 @@ union image_id_index {
- 	bool operator==(const image_id_index& other) const { return id == other.id; }
- 	bool operator!=(const image_id_index& other) const { return id != other.id; }
- 
--	operator size_t&() { return index; }
-+	//operator size_t&() { return index; }
- 	operator imageId&() { return id; }
- };
- 

test/functional/iqdb_queries_controller_test.rb

@@ -3,7 +3,7 @@ require 'test_helper'
 class IqdbQueriesControllerTest < ActionDispatch::IntegrationTest
   context "The iqdb controller" do
     setup do
-      Danbooru.config.stubs(:iqdbs_server).returns("https://karasuma.donmai.us")
+      Danbooru.config.stubs(:iqdb_server).returns("https://karasuma.donmai.us")
       @user = create(:user)
       as(@user) do
         @posts = create_list(:post, 2)
@@ -24,7 +24,7 @@ class IqdbQueriesControllerTest < ActionDispatch::IntegrationTest
         end
 
         should "render a response" do
-          IqdbProxy.expects(:query).with(@url).returns(@mocked_response)
+          IqdbProxy.expects(:query_url).with(@url, nil).returns(@mocked_response)
           get_auth iqdb_queries_path(variant: "xhr"), @user, params: @params
           assert_select("#post_#{@posts[0].id}")
         end
@@ -41,8 +41,8 @@ class IqdbQueriesControllerTest < ActionDispatch::IntegrationTest
           }]
         end
 
-        should "redirect to iqdbs" do
+        should "redirect to iqdb" do
-          IqdbProxy.expects(:query_path).with(@posts[0].preview_file_path).returns(@mocked_response)
+          IqdbProxy.expects(:query_post).with(@posts[0], nil).returns(@mocked_response)
           get_auth iqdb_queries_path, @user, params: @params
           assert_select("#post_#{@posts[0].id}")
         end