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

chore: refactor generic web server code into quark

Browse Source
This commit is contained in:
Paul Makles
2022-06-04 19:02:10 +01:00
parent 7390b3c087
commit d660127c14
38 changed files with 170 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
crates/delta/src/main.rs
View File

@@ -6,101 +6,21 @@ extern crate rocket_okapi;
extern crate serde_json;
#[macro_use]
extern crate lazy_static;
extern crate ctrlc;
pub mod routes;
pub mod util;
pub mod version;
use log::info;
use revolt_quark::rauth::config::{
Captcha, Config, EmailVerificationConfig, SMTPSettings, Template, Templates,
};
use revolt_quark::rauth::RAuth;
use revolt_quark::variables::delta::{
APP_URL, HCAPTCHA_KEY, INVITE_ONLY, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME,
USE_EMAIL, USE_HCAPTCHA,
};
use revolt_quark::DatabaseInfo;
use rocket_cors::AllowedOrigins;
use std::str::FromStr;
#[launch]
async fn rocket() -> _ {
let _guard = revolt_quark::setup_logging();
info!(
"Starting Revolt server [version {}].",
crate::version::VERSION
);
// Configure logging and environment
revolt_quark::configure!();
// Ensure environment variables are present
revolt_quark::variables::delta::preflight_checks();
#[cfg(debug_assertions)]
ctrlc::set_handler(move || {
// Force ungraceful exit to avoid hang.
std::process::exit(0);
})
.expect("Error setting Ctrl-C handler");
let 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.");
let mut config = Config {
email_verification: if *USE_EMAIL {
EmailVerificationConfig::Enabled {
smtp: SMTPSettings {
from: (*SMTP_FROM).to_string(),
host: (*SMTP_HOST).to_string(),
username: (*SMTP_USERNAME).to_string(),
password: (*SMTP_PASSWORD).to_string(),
reply_to: Some("support@revolt.chat".into()),
port: None,
use_tls: None,
},
expiry: Default::default(),
templates: Templates {
verify: Template {
title: "Verify your Revolt account.".into(),
text: include_str!(crate::asset!("templates/verify.txt")).into(),
url: format!("{}/login/verify/", *APP_URL),
html: None,
},
reset: Template {
title: "Reset your Revolt password.".into(),
text: include_str!(crate::asset!("templates/reset.txt")).into(),
url: format!("{}/login/reset/", *APP_URL),
html: None,
},
welcome: None,
},
}
} else {
EmailVerificationConfig::Disabled
},
..Default::default()
};
if *INVITE_ONLY {
config.invite_only = true;
}
if *USE_HCAPTCHA {
config.captcha = Captcha::HCaptcha {
secret: HCAPTCHA_KEY.clone(),
};
}
// Setup database
let db = DatabaseInfo::Auto.connect().await.unwrap();
db.migrate_database().await.unwrap();
@@ -108,7 +28,7 @@ async fn rocket() -> _ {
// Setup rAuth
let rauth = RAuth {
database: db.clone().into(),
config,
config: revolt_quark::util::rauth::config(),
};
// Launch background task workers.
@@ -117,27 +37,11 @@ async fn rocket() -> _ {
// Configure Rocket
let rocket = rocket::build();
routes::mount(rocket)
.mount("/", rocket_cors::catch_all_options_routes())
.mount("/", util::ratelimiter::routes())
.mount(
"/swagger/",
rocket_okapi::swagger_ui::make_swagger_ui(&rocket_okapi::swagger_ui::SwaggerUIConfig {
url: "../openapi.json".to_owned(),
..Default::default()
}),
)
.mount("/", revolt_quark::web::cors::catch_all_options_routes())
.mount("/", revolt_quark::web::ratelimiter::routes())
.mount("/swagger/", revolt_quark::web::swagger::routes())
.manage(rauth)
.manage(db)
.manage(cors.clone())
.attach(util::ratelimiter::RatelimitFairing)
.attach(cors)
.attach(revolt_quark::web::ratelimiter::RatelimitFairing)
.attach(revolt_quark::web::cors::fairing())
}
/// Resolve asset
macro_rules! asset {
($path:literal) => {
concat!(env!("CARGO_MANIFEST_DIR"), "/assets/", $path)
};
}
pub(crate) use asset;

6
crates/delta/src/routes/channels/message_send.rs
View File

@@ -5,7 +5,9 @@ use revolt_quark::{
message::{Masquerade, Reply, SendableEmbed},
Message, User,
},
perms, Db, Error, Permission, Ref, Result,
perms,
web::idempotency::IdempotencyKey,
Db, Error, Permission, Ref, Result,
};
use regex::Regex;
@@ -14,8 +16,6 @@ use serde::{Deserialize, Serialize};
use ulid::Ulid;
use validator::Validate;
use crate::util::idempotency::IdempotencyKey;
#[derive(Validate, Serialize, Deserialize, JsonSchema)]
pub struct DataMessageSend {
/// Unique token to prevent duplicate message sending

2
crates/delta/src/routes/root.rs
View File

@@ -76,7 +76,7 @@ pub struct RevoltConfig {
#[get("/")]
pub async fn root() -> Result<Json<RevoltConfig>> {
Ok(Json(RevoltConfig {
revolt: crate::version::VERSION.to_string(),
revolt: env!("CARGO_PKG_VERSION").to_string(),
features: RevoltFeatures {
captcha: CaptchaFeature {
enabled: *USE_HCAPTCHA,

14
crates/delta/src/routes/users/get_default_avatar.rs
View File

@@ -59,18 +59,6 @@ impl rocket_okapi::response::OpenApiResponderInner for CachedFile {
pub async fn req(target: String) -> CachedFile {
CachedFile((
ContentType::PNG,
match target.chars().last().unwrap() {
// 0123456789ABCDEFGHJKMNPQRSTVWXYZ
'0' | '1' | '2' | '3' | 'S' | 'Z' => {
include_bytes!(crate::asset!("user/2.png")).to_vec()
}
'4' | '5' | '6' | '7' | 'T' => include_bytes!(crate::asset!("user/3.png")).to_vec(),
'8' | '9' | 'A' | 'B' => include_bytes!(crate::asset!("user/4.png")).to_vec(),
'C' | 'D' | 'E' | 'F' | 'V' => include_bytes!(crate::asset!("user/5.png")).to_vec(),
'G' | 'H' | 'J' | 'K' | 'W' => include_bytes!(crate::asset!("user/6.png")).to_vec(),
'M' | 'N' | 'P' | 'Q' | 'X' => include_bytes!(crate::asset!("user/7.png")).to_vec(),
/*'0' | '1' | '2' | '3' | 'R' | 'Y'*/
_ => include_bytes!(crate::asset!("user/1.png")).to_vec(),
},
revolt_quark::util::pfp::avatar(target.chars().last().unwrap()),
))
}

102
crates/delta/src/util/idempotency.rs
View File

@@ -1,102 +0,0 @@
use async_std::sync::Mutex;
use revolt_quark::{Error, Result};
use rocket::http::Status;
use rocket::request::{FromRequest, Outcome};
use rocket_okapi::gen::OpenApiGenerator;
use rocket_okapi::okapi::openapi3::{Parameter, ParameterValue};
use rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput};
use schemars::schema::{InstanceType, SchemaObject, SingleOrVec};
use serde::{Deserialize, Serialize};
use validator::Validate;
#[derive(Validate, Serialize, Deserialize)]
pub struct IdempotencyKey {
#[validate(length(min = 1, max = 64))]
key: String,
}
lazy_static! {
static ref TOKEN_CACHE: Mutex<lru::LruCache<String, ()>> = 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,
) -> 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(),
})
}
}

2
crates/delta/src/util/mod.rs
View File

@@ -1,3 +1 @@
pub mod idempotency;
pub mod ratelimiter;
pub mod regex;

314
crates/delta/src/util/ratelimiter.rs
View File

@@ -1,314 +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 revolt_quark::rauth::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 rocket_okapi::gen::OpenApiGenerator;
use rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput};
use serde::Serialize;
use dashmap::DashMap;
use log::info;
/// Ratelimit Bucket
#[derive(Clone, Copy)]
struct Entry {
used: u8,
reset: u128,
}
lazy_static! {
static ref MAP: DashMap<u64, Entry> = 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, key: u64) {
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;
}
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);
match (segment, resource) {
("users", _) => ("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),
_ => ("any", None),
}
} else {
("any", None)
}
}
/// Resolve per-bucket limits
fn resolve_bucket_limit(bucket: &str) -> u8 {
match bucket {
"users" => 20,
"bots" => 10,
"messaging" => 10,
"channels" => 15,
"servers" => 5,
"auth" => 15,
"auth_delete" => 255,
"swagger" => 100,
_ => 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 entry = Entry::from(key);
let remaining = entry.get_remaining(limit);
let reset = entry.left_until_reset();
if remaining > 0 {
entry.deduct(key);
Ok(Ratelimiter {
key,
limit,
remaining: remaining - 1,
reset,
})
} else {
Err(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,
) -> 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!(),
})
}
}
#[get("/ratelimit")]
fn ratelimit_info(info: RatelimitInformation) -> Json<RatelimitInformation> {
Json(info)
}
pub fn routes() -> Vec<rocket::Route> {
routes![ratelimit_info]
}

1
crates/delta/src/version.rs
View File

@@ -1 +0,0 @@
pub const VERSION: &str = "0.5.3-5-patch.2";