HMC/handmade-revolt-backend
6
0
Fork 0
forked from jmug/stoatchat
Code Pull Requests Actions Packages Releases 2 Activity

refactor: remove quark/web

Browse Source
This commit is contained in:
Paul Makles
2023-08-26 16:19:00 +01:00
parent 3a55d00c6a
commit 51c26e324d
19 changed files with 148 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
crates/quark/src/impl/generic/channels/channel.rs
View File

@@ -1,5 +1,6 @@
use std::collections::HashSet;
use revolt_database::util::idempotency::IdempotencyKey;
use ulid::Ulid;
use crate::{
@@ -11,8 +12,7 @@ use crate::{
},
tasks::{ack::AckEvent, process_embeds},
types::push::MessageAuthor,
variables::delta::{MAX_ATTACHMENT_COUNT, MAX_REPLY_COUNT, MAX_EMBED_COUNT},
web::idempotency::IdempotencyKey,
variables::delta::{MAX_ATTACHMENT_COUNT, MAX_EMBED_COUNT, MAX_REPLY_COUNT},
Database, Error, OverrideField, Ref, Result,
};
@@ -413,7 +413,10 @@ impl Channel {
) -> Result<Message> {
Message::validate_sum(&data.content, data.embeds.as_deref().unwrap_or_default())?;
idempotency.consume_nonce(data.nonce).await?;
idempotency
.consume_nonce(data.nonce)
.await
.map_err(|_| Error::InvalidOperation)?;
// Check the message is not empty
if (data.content.as_ref().map_or(true, |v| v.is_empty()))
@@ -497,16 +500,24 @@ impl Channel {
// Add attachments to message.
let mut attachments = vec![];
if data.attachments.as_ref().is_some_and(|v| v.len() > *MAX_ATTACHMENT_COUNT) {
if data
.attachments
.as_ref()
.is_some_and(|v| v.len() > *MAX_ATTACHMENT_COUNT)
{
return Err(Error::TooManyAttachments {
max: *MAX_ATTACHMENT_COUNT,
});
}
if data.embeds.as_ref().is_some_and(|v| v.len() > *MAX_EMBED_COUNT) {
if data
.embeds
.as_ref()
.is_some_and(|v| v.len() > *MAX_EMBED_COUNT)
{
return Err(Error::TooManyEmbeds {
max: *MAX_EMBED_COUNT
})
max: *MAX_EMBED_COUNT,
});
}
for attachment_id in data.attachments.as_deref().unwrap_or_default() {

19
crates/quark/src/web/cors.rs
View File

@@ -1,19 +0,0 @@
use std::str::FromStr;
pub use rocket_cors::catch_all_options_routes;
use rocket_cors::{AllowedOrigins, Cors};
pub fn new() -> Cors {
rocket_cors::CorsOptions {
allowed_origins: AllowedOrigins::All,
allowed_methods: [
"Get", "Put", "Post", "Delete", "Options", "Head", "Trace", "Connect", "Patch",
]
.iter()
.map(|s| FromStr::from_str(s).unwrap())
.collect(),
..Default::default()
}
.to_cors()
.expect("Failed to create CORS.")
}

102
crates/quark/src/web/idempotency.rs
View File

@@ -1,102 +0,0 @@
use crate::{Error, Result};
use async_std::sync::Mutex;
use revolt_rocket_okapi::gen::OpenApiGenerator;
use revolt_rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput};
use revolt_rocket_okapi::revolt_okapi::openapi3::{Parameter, ParameterValue};
use rocket::http::Status;
use rocket::request::{FromRequest, Outcome};
use schemars::schema::{InstanceType, SchemaObject, SingleOrVec};
use serde::{Deserialize, Serialize};
use validator::Validate;
use once_cell::sync::Lazy;
#[derive(Validate, Serialize, Deserialize)]
pub struct IdempotencyKey {
#[validate(length(min = 1, max = 64))]
key: String,
}
static TOKEN_CACHE: Lazy<Mutex<lru::LruCache<String, ()>>> = Lazy::new(|| Mutex::new(lru::LruCache::new(100)));
impl IdempotencyKey {
// Backwards compatibility.
// Issue #109
pub async fn consume_nonce(&mut self, v: Option<String>) -> Result<()> {
if let Some(v) = v {
let mut cache = TOKEN_CACHE.lock().await;
if cache.get(&v).is_some() {
return Err(Error::DuplicateNonce);
}
cache.put(v.clone(), ());
self.key = v;
}
Ok(())
}
pub fn into_key(self) -> String {
self.key
}
}
impl<'r> OpenApiFromRequest<'r> for IdempotencyKey {
fn from_request_input(
_gen: &mut OpenApiGenerator,
_name: String,
_required: bool,
) -> revolt_rocket_okapi::Result<RequestHeaderInput> {
Ok(RequestHeaderInput::Parameter(Parameter {
name: "Idempotency-Key".to_string(),
description: Some("Unique key to prevent duplicate requests".to_string()),
allow_empty_value: false,
required: false,
deprecated: false,
extensions: schemars::Map::new(),
location: "header".to_string(),
value: ParameterValue::Schema {
allow_reserved: false,
example: None,
examples: None,
explode: None,
style: None,
schema: SchemaObject {
instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::String))),
..Default::default()
},
},
}))
}
}
#[async_trait]
impl<'r> FromRequest<'r> for IdempotencyKey {
type Error = Error;
async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
if let Some(key) = request
.headers()
.get("Idempotency-Key")
.next()
.map(|k| k.to_string())
{
let idempotency = IdempotencyKey { key };
if let Err(error) = idempotency.validate() {
return Outcome::Failure((Status::BadRequest, Error::FailedValidation { error }));
}
let mut cache = TOKEN_CACHE.lock().await;
if cache.get(&idempotency.key).is_some() {
return Outcome::Failure((Status::Conflict, Error::DuplicateNonce));
}
cache.put(idempotency.key.clone(), ());
return Outcome::Success(idempotency);
}
Outcome::Success(IdempotencyKey {
key: ulid::Ulid::new().to_string(),
})
}
}

5
crates/quark/src/web/mod.rs
View File

@@ -1,10 +1,5 @@
use crate::Database;
use rocket::State;
pub mod cors;
pub mod idempotency;
pub mod ratelimiter;
pub mod swagger;
pub use rocket_empty::EmptyResponse;
pub type Db = State<Database>;

331
crates/quark/src/web/ratelimiter.rs
View File

@@ -1,331 +0,0 @@
//! Pulled from lightspeed-tv/backend.
//!
//! This will be replaced again in the near future since
//! I don't want duplication between two different projects.
use std::collections::hash_map::DefaultHasher;
use std::hash::Hasher;
use std::ops::Add;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use crate::authifier::models::Session;
use rocket::fairing::{Fairing, Info, Kind};
use rocket::http::uri::Origin;
use rocket::http::{Method, Status};
use rocket::request::{FromRequest, Outcome};
use rocket::serde::json::Json;
use rocket::{Data, Request, Response};
use revolt_rocket_okapi::gen::OpenApiGenerator;
use revolt_rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput};
use serde::Serialize;
use dashmap::DashMap;
use once_cell::sync::Lazy;
/// Ratelimit Bucket
#[derive(Clone, Copy, Debug)]
struct Entry {
used: u8,
reset: u128,
}
static MAP: Lazy<DashMap<u64, Entry>> = Lazy::new(DashMap::new);
/// Get the current time from Unix Epoch as a Duration
fn now() -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time went backwards...")
}
impl Entry {
/// Find bucket by its key
pub fn from(key: u64) -> Entry {
MAP.get(&key).map(|x| *x).unwrap_or_else(|| Entry {
used: 0,
reset: now().add(Duration::from_secs(10)).as_millis(),
})
}
/// Deduct one unit from the bucket and save
pub fn deduct(&mut self) {
let current_time = now().as_millis();
if current_time > self.reset {
self.used = 1;
self.reset = now().add(Duration::from_secs(10)).as_millis();
} else {
self.used += 1;
}
}
/// Save information
pub fn save(self, key: u64) {
MAP.insert(key, self);
}
/// Get remaining units in the bucket
pub fn get_remaining(&self, limit: u8) -> u8 {
if now().as_millis() > self.reset {
limit
} else {
limit - self.used
}
}
/// Get how long bucket has until reset
pub fn left_until_reset(&self) -> u128 {
let current_time = now().as_millis();
if current_time > self.reset {
0
} else {
self.reset - current_time
}
}
}
/// Ratelimit Guard
#[derive(Serialize, Clone, Copy, Debug)]
#[allow(dead_code)]
pub struct Ratelimiter {
key: u64,
limit: u8,
remaining: u8,
reset: u128,
}
/// Find bucket from given request
///
/// Optionally, include a resource id to hash against.
fn resolve_bucket<'r>(request: &'r rocket::Request<'_>) -> (&'r str, Option<&'r str>) {
if let Some(segment) = request.routed_segment(0) {
let resource = request.routed_segment(1);
let method = request.method();
match (segment, resource, method) {
("users", target, Method::Patch) => ("user_edit", target),
("users", _, _) => {
if let Some("default_avatar") = request.routed_segment(2) {
return ("default_avatar", None);
}
("users", None)
}
("bots", _, _) => ("bots", None),
("channels", Some(id), _) => {
if request.method() == Method::Post {
if let Some("messages") = request.routed_segment(2) {
return ("messaging", Some(id));
}
}
("channels", Some(id))
}
("servers", Some(id), _) => ("servers", Some(id)),
("auth", _, _) => {
if request.method() == Method::Delete {
("auth_delete", None)
} else {
("auth", None)
}
}
("swagger", _, _) => ("swagger", None),
("safety", Some("report"), _) => ("safety_report", Some("report")),
("safety", _, _) => ("safety", None),
_ => ("any", None),
}
} else {
("any", None)
}
}
/// Resolve per-bucket limits
fn resolve_bucket_limit(bucket: &str) -> u8 {
match bucket {
"user_edit" => 2,
"users" => 20,
"bots" => 10,
"messaging" => 10,
"channels" => 15,
"servers" => 5,
"auth" => 15,
"auth_delete" => 255,
"default_avatar" => 255,
"swagger" => 100,
"safety" => 15,
"safety_report" => 3,
_ => 20,
}
}
/// Find the remote IP of the client
fn to_ip(request: &'_ rocket::Request<'_>) -> String {
request
.remote()
.map(|x| x.ip().to_string())
.unwrap_or_default()
}
/// Find the actual IP of the client
fn to_real_ip(request: &'_ rocket::Request<'_>) -> String {
if let Ok(true) = std::env::var("TRUST_CLOUDFLARE").map(|x| x == "1") {
request
.headers()
.get_one("CF-Connecting-IP")
.map(|x| x.to_string())
.unwrap_or_else(|| to_ip(request))
} else {
to_ip(request)
}
}
impl Ratelimiter {
/// Generate guard from identifier and target bucket
pub fn from(
identifier: &str,
(bucket, resource): (&str, Option<&str>),
) -> Result<Ratelimiter, u128> {
let mut key = DefaultHasher::new();
key.write(identifier.as_bytes());
key.write(bucket.as_bytes());
if let Some(id) = resource {
key.write(id.as_bytes());
}
let key = key.finish();
let limit = resolve_bucket_limit(bucket);
let mut entry = Entry::from(key);
let remaining = entry.get_remaining(limit);
if remaining > 0 {
entry.deduct();
let reset = entry.left_until_reset();
entry.save(key);
Ok(Ratelimiter {
key,
limit,
remaining: remaining - 1,
reset,
})
} else {
Err(entry.left_until_reset())
}
}
}
#[async_trait]
impl<'r> FromRequest<'r> for Ratelimiter {
type Error = u128;
async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
let ratelimiter = request
.local_cache_async(async {
use rocket::outcome::Outcome;
let identifier = if let Outcome::Success(session) = request.guard::<Session>().await
{
session.id
} else {
to_real_ip(request)
};
Ratelimiter::from(&identifier, resolve_bucket(request))
})
.await;
match ratelimiter {
Ok(ratelimiter) => Outcome::Success(*ratelimiter),
Err(retry_after) => Outcome::Failure((Status::TooManyRequests, *retry_after)),
}
}
}
impl<'r> OpenApiFromRequest<'r> for Ratelimiter {
fn from_request_input(
_gen: &mut OpenApiGenerator,
_name: String,
_required: bool,
) -> revolt_rocket_okapi::Result<RequestHeaderInput> {
Ok(RequestHeaderInput::None)
}
}
/// Attach ratelimiter to the Rocket application
pub struct RatelimitFairing;
#[rocket::async_trait]
impl Fairing for RatelimitFairing {
fn info(&self) -> Info {
Info {
name: "Ratelimiter",
kind: Kind::Request | Kind::Response,
}
}
async fn on_request(&self, request: &mut Request<'_>, _: &mut Data<'_>) {
use rocket::outcome::Outcome;
if let Outcome::Failure(_) = request.guard::<Ratelimiter>().await {
info!(
"User rate-limited on route {}! (IP = {:?})",
request.uri(),
to_real_ip(request)
);
request.set_method(Method::Get);
request.set_uri(Origin::parse("/ratelimit").unwrap())
}
}
async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
use rocket::outcome::Outcome;
match request.guard::<Ratelimiter>().await {
Outcome::Success(ratelimiter) => {
let Ratelimiter {
key,
limit,
remaining,
reset,
} = ratelimiter;
response.set_raw_header("X-RateLimit-Limit", limit.to_string());
response.set_raw_header("X-RateLimit-Bucket", key.to_string());
response.set_raw_header("X-RateLimit-Remaining", remaining.to_string());
response.set_raw_header("X-RateLimit-Reset-After", reset.to_string());
}
Outcome::Failure(_) => response.set_status(Status::TooManyRequests),
Outcome::Forward(_) => unreachable!(),
}
}
}
#[derive(Serialize)]
#[serde(untagged)]
pub enum RatelimitInformation {
Success(Ratelimiter),
Failure { retry_after: u128 },
}
#[async_trait]
impl<'r> FromRequest<'r> for RatelimitInformation {
type Error = u128;
async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
Outcome::Success(match request.guard::<Ratelimiter>().await {
Outcome::Success(ratelimiter) => RatelimitInformation::Success(ratelimiter),
Outcome::Failure((_, retry_after)) => RatelimitInformation::Failure { retry_after },
_ => unreachable!(),
})
}
}
#[rocket::get("/ratelimit")]
fn ratelimit_info(info: RatelimitInformation) -> Json<RatelimitInformation> {
Json(info)
}
pub fn routes() -> Vec<rocket::Route> {
rocket::routes![ratelimit_info]
}

9
crates/quark/src/web/swagger.rs
View File

@@ -1,9 +0,0 @@
use rocket::Route;
pub fn routes() -> Vec<Route> {
revolt_rocket_okapi::swagger_ui::make_swagger_ui(&revolt_rocket_okapi::swagger_ui::SwaggerUIConfig {
url: "../openapi.json".to_owned(),
..Default::default()
})
.into()
}