chore: refactor generic web server code into quark

crates/quark/src/impl/mongo/admin/migrations/init.rs

@@ -2,7 +2,6 @@ use crate::r#impl::mongo::MongoDb;
 
 use super::scripts::LATEST_REVISION;
 
-use log::info;
 use mongodb::bson::doc;
 use mongodb::options::CreateCollectionOptions;
 

crates/quark/src/impl/mongo/admin/migrations/scripts.rs

@@ -2,7 +2,6 @@ use std::time::Duration;
 
 use bson::Bson;
 use futures::StreamExt;
-use log::info;
 use mongodb::{
     bson::{doc, from_bson, from_document, to_document, Document},
     options::FindOptions,

crates/quark/src/lib.rs

@@ -15,7 +15,7 @@ extern crate lazy_static;
 #[macro_use]
 extern crate bitfield;
 #[macro_use]
-pub extern crate bson;
+extern crate bson;
 
 pub use iso8601_timestamp::Timestamp;
 pub use rauth;
@@ -27,6 +27,7 @@ pub mod models;
 pub mod presence;
 pub mod tasks;
 pub mod types;
+pub mod util;
 
 #[cfg(feature = "rocket_impl")]
 pub mod web;
@@ -34,7 +35,6 @@ pub mod web;
 mod database;
 mod permissions;
 mod traits;
-mod util;
 
 pub use database::*;
 pub use traits::*;
@@ -43,10 +43,19 @@ pub use permissions::defn::*;
 pub use permissions::{get_relationship, perms};
 
 pub use util::{
-    log::setup_logging,
     r#ref::Ref,
     result::{Error, Result},
     variables,
 };
 
+#[cfg(feature = "rocket_impl")]
 pub use web::{Db, EmptyResponse};
+
+/// Resolve asset
+macro_rules! asset {
+    ($path:literal) => {
+        concat!(env!("CARGO_MANIFEST_DIR"), "/assets/", $path)
+    };
+}
+
+pub(crate) use asset;

crates/quark/src/tasks/last_message_id.rs

@@ -2,7 +2,6 @@
 use crate::{models::channel::PartialChannel, Database};
 
 use deadqueue::limited::Queue;
-use log::info;
 use mongodb::bson::doc;
 use std::{collections::HashMap, time::Duration};
 

crates/quark/src/tasks/process_embeds.rs

@@ -6,7 +6,6 @@ use crate::{
 };
 
 use deadqueue::limited::Queue;
-use log::error;
 
 /// Task information
 #[derive(Debug)]

crates/quark/src/util/log.rs

@@ -1,5 +1,5 @@
 /// Configure logging and common Rust variables
-pub fn setup_logging() -> sentry::ClientInitGuard {
+pub fn setup_logging(release: &'static str) -> sentry::ClientInitGuard {
     dotenv::dotenv().ok();
 
     if std::env::var("RUST_LOG").is_err() {
@@ -11,12 +11,24 @@ pub fn setup_logging() -> sentry::ClientInitGuard {
     }
 
     pretty_env_logger::init();
+    info!("Starting {release}");
 
     sentry::init((
         "https://62fd0e02c5354905b4e286757f4beb16@sentry.insert.moe/4",
         sentry::ClientOptions {
-            release: sentry::release_name!(),
+            release: Some(release.into()),
             ..Default::default()
         },
     ))
 }
+
+#[macro_export]
+macro_rules! configure {
+    () => {
+        let _sentry = revolt_quark::util::log::setup_logging(concat!(
+            env!("CARGO_PKG_NAME"),
+            "@",
+            env!("CARGO_PKG_VERSION")
+        ));
+    };
+}

crates/quark/src/util/mod.rs

@@ -1,5 +1,7 @@
 pub mod log;
 pub mod manipulation;
+pub mod pfp;
+pub mod rauth;
 pub mod r#ref;
 pub mod result;
 pub mod value;

crates/quark/src/util/pfp.rs

@@ -0,0 +1,13 @@
+pub fn avatar(v: char) -> Vec<u8> {
+    match v {
+        // 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(),
+    }
+}

crates/quark/src/util/rauth.rs

@@ -0,0 +1,57 @@
+use super::variables::delta::{
+    APP_URL, HCAPTCHA_KEY, INVITE_ONLY, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME,
+    USE_EMAIL, USE_HCAPTCHA,
+};
+
+use crate::rauth::config::{
+    Captcha, Config, EmailVerificationConfig, SMTPSettings, Template, Templates,
+};
+
+pub fn config() -> Config {
+    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(),
+        };
+    }
+
+    config
+}

crates/quark/src/util/variables/delta.rs

@@ -1,8 +1,5 @@
 use std::env;
 
-#[cfg(debug_assertions)]
-use log::warn;
-
 lazy_static! {
     // Application Settings
     pub static ref PUBLIC_URL: String =

crates/quark/src/web/cors.rs

@@ -0,0 +1,19 @@
+use std::str::FromStr;
+
+pub use rocket_cors::catch_all_options_routes;
+use rocket_cors::{AllowedOrigins, Cors};
+
+pub fn fairing() -> 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.")
+}

crates/quark/src/web/idempotency.rs

@@ -0,0 +1,103 @@
+use crate::{Error, Result};
+
+use async_std::sync::Mutex;
+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(),
+        })
+    }
+}

crates/quark/src/web/mod.rs

@@ -1,5 +1,10 @@
 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>;

crates/quark/src/web/ratelimiter.rs

@@ -0,0 +1,312 @@
+//! 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::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;
+
+/// 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!(),
+        })
+    }
+}
+
+#[rocket::get("/ratelimit")]
+fn ratelimit_info(info: RatelimitInformation) -> Json<RatelimitInformation> {
+    Json(info)
+}
+
+pub fn routes() -> Vec<rocket::Route> {
+    rocket::routes![ratelimit_info]
+}

crates/quark/src/web/swagger.rs

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