chore: refactor generic web server code into quark
This commit is contained in:
Paul Makles
@@ -18,7 +18,6 @@ linkify = "0.6.0"
|
||||
once_cell = "1.4.1"
|
||||
env_logger = "0.7.1"
|
||||
lazy_static = "1.4.0"
|
||||
ctrlc = { version = "3.0", features = ["termination"] }
|
||||
|
||||
# Lang. Utilities
|
||||
regex = "1"
|
||||
@@ -54,7 +53,6 @@ mobc-redis = { version = "0.7.0", default-features = false, features = ["async-s
|
||||
rocket = { version = "0.5.0-rc.2", default-features = false, features = ["json"] }
|
||||
rocket_empty = { git = "https://github.com/insertish/rocket_empty", branch = "master" }
|
||||
rocket_rauth = { git = "https://github.com/insertish/rauth", rev = "a82cb5d23da4787e75d6f3c795189ce31e1590c4" }
|
||||
rocket_cors = { git = "https://github.com/lawliet89/rocket_cors", rev = "5843861a88958c16bfaa0b40f0d8910772bcd2f6" }
|
||||
|
||||
# spec generation
|
||||
schemars = "0.8.8"
|
||||
|
||||
@@ -1,63 +0,0 @@
|
||||
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
|
||||
<html xmlns="http://www.w3.org/1999/xhtml" lang="en-GB">
|
||||
<head>
|
||||
<meta http-equiv="Content-Type" content="text/html; charset=UTF-8" />
|
||||
<title>Reset your password.</title>
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0"/>
|
||||
|
||||
<style type="text/css">
|
||||
a[x-apple-data-detectors] {color: inherit !important;}
|
||||
</style>
|
||||
|
||||
</head>
|
||||
<body style="margin: 0; padding: 0;">
|
||||
<table border="0" cellpadding="0" cellspacing="0" width="100%">
|
||||
<tr>
|
||||
<td style="padding: 20px 0 30px 0;">
|
||||
|
||||
<table align="center" border="0" cellpadding="0" cellspacing="0" width="600" style="border-collapse: collapse; border: 1px solid #cccccc;">
|
||||
<tr>
|
||||
<td align="center" bgcolor="#ff4654">
|
||||
<img src="https://revolt.chat/header.png" alt="Revolt logo" width="600" height="168" style="display: block;" />
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td bgcolor="#ffffff" style="padding: 40px 30px 40px 30px;">
|
||||
<table border="0" cellpadding="0" cellspacing="0" width="100%" style="border-collapse: collapse;">
|
||||
<tr>
|
||||
<td style="color: #153643; font-family: Arial, sans-serif;">
|
||||
<h1 style="font-size: 24px; margin: 0;">Reset your password!</h1>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="color: #153643; font-family: Arial, sans-serif; font-size: 16px; line-height: 24px; padding: 20px 0 0 0;">
|
||||
<p>
|
||||
You requested a password reset, if you didn't perform this action you can safely ignore this email.
|
||||
</p>
|
||||
<p style="margin: 0;">
|
||||
Reset your password by navigating to <a href="{{url}}">{{url}}</a>.
|
||||
</p>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td bgcolor="#ff4654" style="padding: 30px 30px;">
|
||||
<table border="0" cellpadding="0" cellspacing="0" width="100%" style="border-collapse: collapse;">
|
||||
<tr>
|
||||
<td style="color: #ffffff; font-family: Arial, sans-serif; font-size: 14px;">
|
||||
<p style="margin: 0;">Sent by Revolt. · Website: <a style="color: white;" href="https://revolt.chat">https://revolt.chat</a></p>
|
||||
<p>Revolt is a user-first chat platform built with modern web technologies.</p>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1,5 +0,0 @@
|
||||
You requested a password reset, if you did not perform this action you can safely ignore this email.
|
||||
|
||||
Reset your password here: {{url}}
|
||||
|
||||
Sent by Revolt.
|
||||
@@ -1,14 +0,0 @@
|
||||
<h2>Reset your password.</h2>
|
||||
<p>
|
||||
You requested a password reset, if you did not perform this action you can safely ignore this email.
|
||||
</p>
|
||||
<p>
|
||||
Reset your password here: <a href="{{url}}">{{url}}</a>
|
||||
</p>
|
||||
<br/>
|
||||
<p>
|
||||
Sent by Revolt. · Website: <a href="https://revolt.chat">https://revolt.chat</a>
|
||||
</p>
|
||||
<p>
|
||||
Revolt is a user-first chat platform built with modern web technologies.
|
||||
</p>
|
||||
@@ -1,63 +0,0 @@
|
||||
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
|
||||
<html xmlns="http://www.w3.org/1999/xhtml" lang="en-GB">
|
||||
<head>
|
||||
<meta http-equiv="Content-Type" content="text/html; charset=UTF-8" />
|
||||
<title>Verify your account.</title>
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0"/>
|
||||
|
||||
<style type="text/css">
|
||||
a[x-apple-data-detectors] {color: inherit !important;}
|
||||
</style>
|
||||
|
||||
</head>
|
||||
<body style="margin: 0; padding: 0;">
|
||||
<table border="0" cellpadding="0" cellspacing="0" width="100%">
|
||||
<tr>
|
||||
<td style="padding: 20px 0 30px 0;">
|
||||
|
||||
<table align="center" border="0" cellpadding="0" cellspacing="0" width="600" style="border-collapse: collapse; border: 1px solid #cccccc;">
|
||||
<tr>
|
||||
<td align="center" bgcolor="#ff4654">
|
||||
<img src="https://revolt.chat/header.png" alt="Revolt logo" width="600" height="168" style="display: block;" />
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td bgcolor="#ffffff" style="padding: 40px 30px 40px 30px;">
|
||||
<table border="0" cellpadding="0" cellspacing="0" width="100%" style="border-collapse: collapse;">
|
||||
<tr>
|
||||
<td style="color: #153643; font-family: Arial, sans-serif;">
|
||||
<h1 style="font-size: 24px; margin: 0;">You're almost there!</h1>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="color: #153643; font-family: Arial, sans-serif; font-size: 16px; line-height: 24px; padding: 20px 0 0 0;">
|
||||
<p>
|
||||
Verify your account to be able to log into the platform. If you didn't perform this action you can safely ignore this email.
|
||||
</p>
|
||||
<p style="margin: 0;">
|
||||
Please verify your account by navigating to <a href="{{url}}">{{url}}</a>.
|
||||
</p>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td bgcolor="#ff4654" style="padding: 30px 30px;">
|
||||
<table border="0" cellpadding="0" cellspacing="0" width="100%" style="border-collapse: collapse;">
|
||||
<tr>
|
||||
<td style="color: #ffffff; font-family: Arial, sans-serif; font-size: 14px;">
|
||||
<p style="margin: 0;">Sent by Revolt. · Website: <a style="color: white;" href="https://revolt.chat">https://revolt.chat</a></p>
|
||||
<p>Revolt is a user-first chat platform built with modern web technologies.</p>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1,6 +0,0 @@
|
||||
You're almost there!
|
||||
If you did not perform this action you can safely ignore this email.
|
||||
|
||||
Please verify your account here: {{url}}
|
||||
|
||||
Sent by Revolt.
|
||||
@@ -1,15 +0,0 @@
|
||||
<h2>You're almost there!</h2>
|
||||
<p>
|
||||
Verify your account to be able to log into the platform.<br/>
|
||||
If you did not perform this action you can safely ignore this email.
|
||||
</p>
|
||||
<p>
|
||||
Please verify your account here: <a href="{{url}}">{{url}}</a>
|
||||
</p>
|
||||
<br/>
|
||||
<p>
|
||||
Sent by Revolt. · Website: <a href="https://revolt.chat">https://revolt.chat</a>
|
||||
</p>
|
||||
<p>
|
||||
Revolt is a user-first chat platform built with modern web technologies.
|
||||
</p>
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 6.7 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 6.3 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 6.4 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 6.3 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 6.6 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 6.6 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 5.5 KiB |
@@ -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;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1 @@
|
||||
pub mod idempotency;
|
||||
pub mod ratelimiter;
|
||||
pub mod regex;
|
||||
|
||||
@@ -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 +0,0 @@
|
||||
pub const VERSION: &str = "0.5.3-5-patch.2";
|
||||
Reference in New Issue
Block a user