mirror of
https://github.com/Syfaro/fuzzysearch.git
synced 2024-09-20 22:08:35 +00:00
3ade5aeba9
* Initial progress on searching by URL. * Avoid rejections for error messages. * Handle some more rejections. * Fix build issues. * Remove detailed error messages. * Fix issue with built Docker image. * Add rate limit headers to all responses. * Remove unneeded dependency. * Limit URLs to 10MB.
139 lines
4.0 KiB
Rust
139 lines
4.0 KiB
Rust
use crate::models::DB;
|
|
use crate::types::*;
|
|
|
|
#[macro_export]
|
|
macro_rules! rate_limit {
|
|
($api_key:expr, $db:expr, $limit:tt, $group:expr) => {
|
|
rate_limit!($api_key, $db, $limit, $group, 1)
|
|
};
|
|
|
|
($api_key:expr, $db:expr, $limit:tt, $group:expr, $incr_by:expr) => {{
|
|
let api_key = match crate::models::lookup_api_key($api_key, $db).await {
|
|
Some(api_key) => api_key,
|
|
None => return Ok(Box::new(Error::ApiKey)),
|
|
};
|
|
|
|
let rate_limit = match crate::utils::update_rate_limit(
|
|
$db,
|
|
api_key.id,
|
|
api_key.$limit,
|
|
$group,
|
|
$incr_by,
|
|
)
|
|
.await
|
|
{
|
|
Ok(rate_limit) => rate_limit,
|
|
Err(err) => return Ok(Box::new(Error::Postgres(err))),
|
|
};
|
|
|
|
match rate_limit {
|
|
crate::types::RateLimit::Limited => return Ok(Box::new(Error::RateLimit)),
|
|
crate::types::RateLimit::Available(count) => count,
|
|
}
|
|
}};
|
|
}
|
|
|
|
#[macro_export]
|
|
macro_rules! early_return {
|
|
($val:expr) => {
|
|
match $val {
|
|
Ok(val) => val,
|
|
Err(err) => return Ok(Box::new(Error::from(err))),
|
|
}
|
|
};
|
|
}
|
|
|
|
/// Increment the rate limit for a group.
|
|
///
|
|
/// We need to specify the ID of the API key to increment, the key's limit for
|
|
/// the specified group, the name of the group we're incrementing, and the
|
|
/// amount to increment for this request. This should remain as 1 except for
|
|
/// joined requests.
|
|
#[tracing::instrument(skip(db))]
|
|
pub async fn update_rate_limit(
|
|
db: DB<'_>,
|
|
key_id: i32,
|
|
key_group_limit: i16,
|
|
group_name: &'static str,
|
|
incr_by: i16,
|
|
) -> Result<RateLimit, tokio_postgres::Error> {
|
|
let now = chrono::Utc::now();
|
|
let timestamp = now.timestamp();
|
|
let time_window = timestamp - (timestamp % 60);
|
|
|
|
let rows = db
|
|
.query(
|
|
"INSERT INTO
|
|
rate_limit (api_key_id, time_window, group_name, count)
|
|
VALUES
|
|
($1, $2, $3, $4)
|
|
ON CONFLICT ON CONSTRAINT unique_window
|
|
DO UPDATE set count = rate_limit.count + $4
|
|
RETURNING rate_limit.count",
|
|
&[&key_id, &time_window, &group_name, &incr_by],
|
|
)
|
|
.await?;
|
|
|
|
let count: i16 = rows[0].get(0);
|
|
|
|
if count > key_group_limit {
|
|
Ok(RateLimit::Limited)
|
|
} else {
|
|
Ok(RateLimit::Available((
|
|
key_group_limit - count,
|
|
key_group_limit,
|
|
)))
|
|
}
|
|
}
|
|
|
|
pub fn extract_rows(
|
|
rows: Vec<tokio_postgres::Row>,
|
|
hash: Option<&[u8]>,
|
|
) -> impl IntoIterator<Item = File> + '_ {
|
|
rows.into_iter().map(move |row| {
|
|
let dbhash: i64 = row.get("hash");
|
|
let dbbytes = dbhash.to_be_bytes();
|
|
|
|
let (furaffinity_id, e621_id, twitter_id): (Option<i32>, Option<i32>, Option<i64>) = (
|
|
row.get("furaffinity_id"),
|
|
row.get("e621_id"),
|
|
row.get("twitter_id"),
|
|
);
|
|
|
|
let (site_id, site_info) = if let Some(fa_id) = furaffinity_id {
|
|
(
|
|
fa_id as i64,
|
|
Some(SiteInfo::FurAffinity(FurAffinityFile {
|
|
file_id: row.get("file_id"),
|
|
})),
|
|
)
|
|
} else if let Some(e6_id) = e621_id {
|
|
(
|
|
e6_id as i64,
|
|
Some(SiteInfo::E621(E621File {
|
|
sources: row.get("sources"),
|
|
})),
|
|
)
|
|
} else if let Some(t_id) = twitter_id {
|
|
(t_id, Some(SiteInfo::Twitter))
|
|
} else {
|
|
(-1, None)
|
|
};
|
|
|
|
File {
|
|
id: row.get("id"),
|
|
site_id,
|
|
site_info,
|
|
site_id_str: site_id.to_string(),
|
|
url: row.get("url"),
|
|
hash: Some(dbhash),
|
|
distance: hash
|
|
.map(|hash| hamming::distance_fast(&dbbytes, &hash).ok())
|
|
.flatten(),
|
|
artists: row.get("artists"),
|
|
filename: row.get("filename"),
|
|
searched_hash: None,
|
|
}
|
|
})
|
|
}
|