edshot99/eBooru
1
0
Fork 0
forked from e621ng/e621ng
Code Issues 11 Releases 1 Activity

[IQDB] Switch to new implementation

Browse Source
This commit is contained in:
Earlopain 2023-04-24 14:46:26 +02:00
parent edae5a6b77
commit e87d617c6f
No known key found for this signature in database
GPG Key ID: 6CFB948E15246897
14 changed files with 99 additions and 242 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.env.sample
View File

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

8
.gitignore vendored
View File

@ -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

52
app/controllers/iqdb_queries_controller.rb
View File

@ -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

2
app/jobs/iqdb_remove_job_new.rb
View File

@ -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

2
app/jobs/iqdb_update_job_new.rb
View File

@ -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

96
app/logical/iqdb_proxy.rb
View File

@ -1,24 +1,69 @@
class IqdbProxy
# frozen_string_literal: true
module IqdbProxy
class Error < StandardError; end
def self.query(image_url)
response = make_request(url: image_url)
decorate_posts(response)
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 self.query_file(image)
response = make_request(file: image)
decorate_posts(response)
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 self.query_path(image_path)
File.open(image_path) do |f|
query_file(f)
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 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)
raise NotImplementedError if Danbooru.config.iqdbs_server.blank?
def generate_thumbnail(file_path)
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)
url.path = "/similar"
json = HTTParty.post(url.to_s, { body: params }.merge(Danbooru.config.httparty_options))
return {} if json.code != 200
json.parsed_response
rescue Errno::ECONNREFUSED, Errno::EADDRNOTAVAIL
raise Error, "This service is temporarily unavailable. Please try again later."
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

95
app/logical/iqdb_proxy_new.rb
View File

@ -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

2
app/models/post.rb
View File

@ -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)

2
app/views/posts/partials/show/_options.html.erb
View File

@ -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>

3
config/danbooru_default_config.rb
View File

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

21
docker-compose.yml
View File

@ -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_new:5588
DANBOORU_IQDB_SERVER: http://iqdb:5588
# These are just development secrets, do not use them in production
SECRET_TOKEN: 1c58518a891eff4520cadc59afa9e378a9325f1247544ff258096e497f095f45
SESSION_SECRET_KEY: 44b4f44e9f253c406cbe727d403d500c1cecff943e4d2aea8f5447f28846fffe
@ -107,24 +106,6 @@ services:
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:
image: ghcr.io/e621ng/iqdb:d4fed9d9a51184e72d2f14d4ec461d7830bd177a
command: iqdb http 0.0.0.0 5588 /iqdb/e621_v2.db
volumes:

34
docker/iqdb/Dockerfile
View File

@ -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"]

13
docker/iqdb/iqdb.patch
View File

@ -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; }
};

8
test/functional/iqdb_queries_controller_test.rb
View File

@ -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
IqdbProxy.expects(:query_path).with(@posts[0].preview_file_path).returns(@mocked_response)
should "redirect to iqdb" do
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