From 658a0233348a43d1abebcbfd37f0fd0663baa165 Mon Sep 17 00:00:00 2001
From: Syfaro <syfaro@huefox.com>
Date: Tue, 21 Jan 2020 03:08:15 -0600
Subject: [PATCH] Some improvements.

---
 src/filters.rs  |  1 +
 src/handlers.rs | 64 +++++++++++++------------------------------------
 src/main.rs     |  5 +++-
 src/models.rs   | 53 ++++++++++++++++++++++++++++++++++++++++
 src/types.rs    | 14 +++++++++++
 5 files changed, 89 insertions(+), 48 deletions(-)

diff --git a/src/filters.rs b/src/filters.rs
index b81221e..25f6ca7 100644
--- a/src/filters.rs
+++ b/src/filters.rs
@@ -20,6 +20,7 @@ pub fn search_image(db: Pool) -> impl Filter<Extract = impl Reply, Error = Rejec
     warp::path("image")
         .and(warp::post())
         .and(warp::multipart::form().max_length(1024 * 1024 * 10))
+        .and(warp::query::<ImageSearchOpts>())
         .and(with_pool(db))
         .and(with_api_key())
         .and_then(handlers::search_image)
diff --git a/src/handlers.rs b/src/handlers.rs
index 217196f..c543ac3 100644
--- a/src/handlers.rs
+++ b/src/handlers.rs
@@ -1,7 +1,8 @@
+use crate::models::image_query;
 use crate::types::*;
 use crate::utils::{extract_e621_rows, extract_fa_rows};
 use crate::{rate_limit, Pool};
-use log::{info, debug};
+use log::{debug, info};
 use warp::{reject, Rejection, Reply};
 
 fn map_bb8_err(err: bb8::RunError<tokio_postgres::Error>) -> Rejection {
@@ -37,6 +38,7 @@ impl warp::reject::Reject for Error {}
 
 pub async fn search_image(
     form: warp::multipart::FormData,
+    opts: ImageSearchOpts,
     db: Pool,
     api_key: String,
 ) -> Result<impl Reply, Rejection> {
@@ -77,52 +79,20 @@ pub async fn search_image(
 
     debug!("Matching hash {}", num);
 
-    let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = vec![&num];
-
-    let fa = db.query(
-        "SELECT
-                submission.id,
-                submission.url,
-                submission.filename,
-                submission.file_id,
-                submission.hash,
-                submission.hash_int,
-                artist.name
-            FROM
-                submission
-            JOIN artist
-                ON artist.id = submission.artist_id
-            WHERE
-                hash_int <@ ($1, 10)",
-        &params,
-    );
-
-    let e621 = db.query(
-        "SELECT
-                e621.id,
-                e621.hash,
-                e621.data->>'file_url' url,
-                e621.data->>'md5' md5,
-                sources.list sources,
-                artists.list artists,
-                (e621.data->>'md5') || '.' || (e621.data->>'file_ext') filename
-            FROM
-                e621,
-                LATERAL (
-                    SELECT array_agg(s) list
-                    FROM jsonb_array_elements_text(data->'sources') s
-                ) sources,
-                LATERAL (
-                    SELECT array_agg(s) list
-                    FROM jsonb_array_elements_text(data->'artist') s
-                ) artists
-            WHERE
-                hash <@ ($1, 10)",
-        &params,
-    );
-
-    let results = futures::future::join(fa, e621).await;
-    let (fa_results, e621_results) = (results.0.unwrap(), results.1.unwrap());
+    let (fa_results, e621_results) = {
+        if opts.search_type == Some(ImageSearchType::Force) {
+            image_query(&db, num, 10).await.unwrap()
+        } else {
+            let (fa_results, e621_results) = image_query(&db, num, 0).await.unwrap();
+            if fa_results.len() + e621_results.len() == 0
+                && opts.search_type != Some(ImageSearchType::Exact)
+            {
+                image_query(&db, num, 10).await.unwrap()
+            } else {
+                (fa_results, e621_results)
+            }
+        }
+    };
 
     let mut items = Vec::with_capacity(fa_results.len() + e621_results.len());
 
diff --git a/src/main.rs b/src/main.rs
index 5a3ba44..ab9c7f8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -27,9 +27,12 @@ async fn main() {
     let log = warp::log("fuzzysearch");
     let cors = warp::cors()
         .allow_any_origin()
+        .allow_headers(vec!["x-api-key"])
         .allow_methods(vec!["GET", "POST"]);
 
-    let api = filters::search(db_pool);
+    let options = warp::options().map(|| "✓");
+
+    let api = options.or(filters::search(db_pool));
     let routes = api
         .or(warp::path::end()
             .map(|| warp::redirect(warp::http::Uri::from_static("https://fuzzysearch.net"))))
diff --git a/src/models.rs b/src/models.rs
index 1b4aa1f..1e260f3 100644
--- a/src/models.rs
+++ b/src/models.rs
@@ -34,3 +34,56 @@ pub async fn lookup_api_key(key: &str, db: DB<'_>) -> Option<ApiKey> {
         _ => None,
     }
 }
+
+pub async fn image_query(
+    db: DB<'_>,
+    num: i64,
+    distance: i64,
+) -> Result<(Vec<tokio_postgres::Row>, Vec<tokio_postgres::Row>), tokio_postgres::Error> {
+    let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = vec![&num, &distance];
+
+    let fa = db.query(
+        "SELECT
+                submission.id,
+                submission.url,
+                submission.filename,
+                submission.file_id,
+                submission.hash,
+                submission.hash_int,
+                artist.name
+            FROM
+                submission
+            JOIN artist
+                ON artist.id = submission.artist_id
+            WHERE
+                hash_int <@ ($1, $2)",
+        &params,
+    );
+
+    let e621 = db.query(
+        "SELECT
+                e621.id,
+                e621.hash,
+                e621.data->>'file_url' url,
+                e621.data->>'md5' md5,
+                sources.list sources,
+                artists.list artists,
+                (e621.data->>'md5') || '.' || (e621.data->>'file_ext') filename
+            FROM
+                e621,
+                LATERAL (
+                    SELECT array_agg(s) list
+                    FROM jsonb_array_elements_text(data->'sources') s
+                ) sources,
+                LATERAL (
+                    SELECT array_agg(s) list
+                    FROM jsonb_array_elements_text(data->'artist') s
+                ) artists
+            WHERE
+                hash <@ ($1, $2)",
+        &params,
+    );
+
+    let results = futures::future::join(fa, e621).await;
+    Ok((results.0?, results.1?))
+}
diff --git a/src/types.rs b/src/types.rs
index bae316f..1aace91 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -68,6 +68,20 @@ pub struct FileSearchOpts {
     pub url: Option<String>,
 }
 
+#[derive(Debug, Deserialize)]
+pub struct ImageSearchOpts {
+    #[serde(rename = "type")]
+    pub search_type: Option<ImageSearchType>,
+}
+
+#[derive(Debug, Deserialize, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum ImageSearchType {
+    Close,
+    Exact,
+    Force,
+}
+
 #[derive(Debug, Serialize)]
 pub struct ImageSimilarity {
     pub hash: i64,