Skip to content

Commit

Permalink
only use iqdb if we have a usable mime type
Browse files Browse the repository at this point in the history
  • Loading branch information
DonovanDMC committed Jul 8, 2024
1 parent c0c03e3 commit 80e2036
Showing 1 changed file with 24 additions and 18 deletions.
42 changes: 24 additions & 18 deletions src/util/Sauce.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@ import type { Post } from "e621";
import type { JSONResponse } from "yiffy";
import { Strings } from "@uwu-codes/utils";
import { fetch } from "undici";
import { fileTypeFromBuffer } from "file-type";
import { STATUS_CODES } from "node:http";
import { Blob } from "node:buffer";
import { fileTypeFromBuffer } from "file-type";

export class PreCheckError extends Error {
override name = "PreCheckError";
Expand Down Expand Up @@ -107,23 +107,29 @@ export default async function Sauce(input: string, simularity = 80, skipCheck =

iqdb: if (!method) {
const img = await RequestProxy.get(input);
if (!img.ok) break iqdb;

const content = Buffer.from(await img.response.arrayBuffer())
const type = (await fileTypeFromBuffer(content))!;
const result = await fetch(`https://e621.net/iqdb_queries.json?search[score_cutoff]=${simularity}`, {
method: "POST",
body: new Blob([content], { type: type.mime })
});

if (result.status !== 200) break iqdb;

const res = (await result.json()) as Array<{ post_id: number; score: number; }>;

if(res.length > 0) {
method = "iqdb";
post = await E621.posts.get(res[0].post_id);
saucePercent = res[0].score;
if (!img.ok) {
break iqdb;
}

const content = Buffer.from(await img.response.arrayBuffer());
const type = await fileTypeFromBuffer(content);
if (type && ["image/png", "image/jpeg"].includes(type.mime)) {
const result = await fetch(`https://e621.net/iqdb_queries.json?search[score_cutoff]=${simularity}`, {
method: "POST",
body: new Blob([content], { type: type.mime })
});

if (result.status !== 200) {
break iqdb;
}

const res = (await result.json()) as Array<{ post_id: number; score: number; }>;

if (res.length !== 0) {
method = "iqdb";
post = await E621.posts.get(res[0].post_id);
saucePercent = res[0].score;
}
}
}

Expand Down

0 comments on commit 80e2036

Please sign in to comment.