From c9011ac692601ffefb097dd808602c6b917dc4f2 Mon Sep 17 00:00:00 2001 From: Paul Makles Date: Sat, 5 Aug 2023 11:24:02 +0100 Subject: [PATCH] chore: migrate bots create route --- Cargo.lock | 5 + clippy.toml | 16 +- crates/core/database/Cargo.toml | 5 +- crates/core/database/src/drivers/mongodb.rs | 2 +- crates/core/database/src/models/users/mod.rs | 4 + .../core/database/src/models/users/model.rs | 141 +++++++++++++++--- crates/core/database/src/models/users/ops.rs | 3 + .../database/src/models/users/ops/mongodb.rs | 33 ++++ .../src/models/users/ops/reference.rs | 12 ++ .../core/database/src/models/users/schema.rs | 31 ++++ crates/core/models/Cargo.toml | 6 +- crates/core/models/src/v0/bots.rs | 8 +- crates/core/models/src/v0/users.rs | 9 ++ crates/core/result/src/lib.rs | 1 + crates/core/result/src/rocket.rs | 1 + crates/delta/src/routes/bots/create.rs | 30 ++-- 16 files changed, 272 insertions(+), 35 deletions(-) create mode 100644 crates/core/database/src/models/users/schema.rs diff --git a/Cargo.lock b/Cargo.lock index 9d23d056..af401d26 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2882,6 +2882,7 @@ dependencies = [ "async-trait", "authifier", "bson", + "decancer", "futures", "indexmap", "iso8601-timestamp 0.2.11", @@ -2896,7 +2897,9 @@ dependencies = [ "revolt-permissions", "revolt-presence", "revolt-result", + "revolt_okapi", "revolt_optional_struct", + "revolt_rocket_okapi", "rocket", "schemars", "serde", @@ -2952,6 +2955,8 @@ version = "0.6.5" dependencies = [ "indexmap", "iso8601-timestamp 0.2.11", + "once_cell", + "regex", "revolt-permissions", "revolt_optional_struct", "schemars", diff --git a/clippy.toml b/clippy.toml index a1c69535..00ef1dce 100644 --- a/clippy.toml +++ b/clippy.toml @@ -3,13 +3,23 @@ disallowed-methods = [ "revolt_database::models::bots::model::Bot::remove_field", # Prefer to use Object::create() - "revolt_database::models::safety_strikes::ops::AbstractAccountStrikes::insert_account_strike", + "revolt_database::models::bots::ops::AbstractBots::*", + "revolt_database::models::bots::ops::AbstractBots::insert_bot", + "revolt_database::models::bots::ops::AbstractChannelInvites::insert_invite", + "revolt_database::models::bots::ops::AbstractChannelUnreads::acknowledge_message", + "revolt_database::models::bots::ops::AbstractChannels::insert_channel", + "revolt_database::models::bots::ops::AbstractEmojis::insert_emoji", + "revolt_database::models::bots::ops::AbstractMessages::insert_message", + "revolt_database::models::bots::ops::AbstractRatelimitEvents::insert_ratelimit_event", + "revolt_database::models::bots::ops::AbstractRatelimitEvents::insert_ratelimit_event", + "revolt_database::models::bots::ops::AbstractServerBans::insert_ban", + "revolt_database::models::bots::ops::AbstractServerMembers::insert_member", + "revolt_database::models::bots::ops::AbstractServers::insert_server", + "revolt_database::models::bots::ops::AbstractUsers::insert_user", # Prefer to use Object::update(&self) "revolt_database::models::bots::ops::AbstractBots::update_bot", - "revolt_database::models::safety_strikes::ops::AbstractAccountStrikes::update_account_strike", # Prefer to use Object::delete(&self) "revolt_database::models::bots::ops::AbstractBots::delete_bot", - "revolt_database::models::safety_strikes::ops::AbstractAccountStrikes::delete_account_strike", ] diff --git a/crates/core/database/Cargo.toml b/crates/core/database/Cargo.toml index f4880144..9524f406 100644 --- a/crates/core/database/Cargo.toml +++ b/crates/core/database/Cargo.toml @@ -14,7 +14,7 @@ mongodb = ["dep:mongodb", "bson"] # ... Other async-std-runtime = ["async-std"] -rocket-impl = ["rocket", "schemars"] +rocket-impl = ["rocket", "schemars", "revolt_okapi", "revolt_rocket_okapi"] redis-is-patched = ["revolt-presence/redis-is-patched"] # Default Features @@ -37,6 +37,7 @@ ulid = "1.0.0" nanoid = "0.4.0" once_cell = "1.17" indexmap = "1.9.1" +decancer = "1.6.2" # Serialisation serde_json = "1" @@ -68,6 +69,8 @@ schemars = { version = "0.8.8", optional = true } rocket = { version = "0.5.0-rc.2", default-features = false, features = [ "json", ], optional = true } +revolt_okapi = { version = "0.9.1", optional = true } +revolt_rocket_okapi = { version = "0.9.1", optional = true } # Authifier authifier = { version = "1.0" } diff --git a/crates/core/database/src/drivers/mongodb.rs b/crates/core/database/src/drivers/mongodb.rs index 74e5c6ef..a4bbe56b 100644 --- a/crates/core/database/src/drivers/mongodb.rs +++ b/crates/core/database/src/drivers/mongodb.rs @@ -238,6 +238,6 @@ pub fn prefix_keys(t: &T, prefix: &str) -> HashMap = serde_json::from_str(&v).unwrap(); v.into_iter() .filter(|(_k, v)| !v.is_null()) - .map(|(k, v)| (prefix.to_owned() + &k, v)) + .map(|(k, v)| (format!("{}{}", prefix.to_owned(), k), v)) .collect() } diff --git a/crates/core/database/src/models/users/mod.rs b/crates/core/database/src/models/users/mod.rs index 04963173..84935c6f 100644 --- a/crates/core/database/src/models/users/mod.rs +++ b/crates/core/database/src/models/users/mod.rs @@ -2,8 +2,12 @@ mod model; mod ops; #[cfg(feature = "rocket-impl")] mod rocket; +#[cfg(feature = "rocket-impl")] +mod schema; #[cfg(feature = "rocket-impl")] pub use self::rocket::*; +#[cfg(feature = "rocket-impl")] +pub use self::schema::*; pub use model::*; pub use ops::*; diff --git a/crates/core/database/src/models/users/model.rs b/crates/core/database/src/models/users/model.rs index 1e49e022..4287fd44 100644 --- a/crates/core/database/src/models/users/model.rs +++ b/crates/core/database/src/models/users/model.rs @@ -1,9 +1,11 @@ -use std::collections::HashSet; +use std::{collections::HashSet, time::Duration}; -use crate::{Database, File}; +use crate::{Database, File, RatelimitEvent}; use once_cell::sync::Lazy; -use revolt_result::{Error, ErrorType, Result}; +use rand::seq::SliceRandom; +use revolt_result::{create_error, Error, ErrorType, Result}; +use ulid::Ulid; auto_derived_partial!( /// # User @@ -117,7 +119,126 @@ auto_derived!( } ); +pub static DISCRIMINATOR_SEARCH_SPACE: Lazy> = Lazy::new(|| { + let mut set = (2..9999) + .map(|v| format!("{:0>4}", v)) + .collect::>(); + + for discrim in [ + 123, 1234, 1111, 2222, 3333, 4444, 5555, 6666, 7777, 8888, 9999, + ] { + set.remove(&format!("{:0>4}", discrim)); + } + + set.into_iter().collect() +}); + impl User { + /// Sanitise and validate a username can be used + pub fn validate_username(username: String) -> Result { + // Copy the username for validation + let username_lowercase = username.to_lowercase(); + + // Block homoglyphs + if decancer::cure(&username_lowercase).into_str() != username_lowercase { + return Err(create_error!(InvalidUsername)); + } + + // Ensure the username itself isn't blocked + const BLOCKED_USERNAMES: &[&str] = &["admin", "revolt"]; + + for username in BLOCKED_USERNAMES { + if username_lowercase == *username { + return Err(create_error!(InvalidUsername)); + } + } + + // Ensure none of the following substrings show up in the username + const BLOCKED_SUBSTRINGS: &[&str] = &["```"]; + + for substr in BLOCKED_SUBSTRINGS { + if username_lowercase.contains(substr) { + return Err(create_error!(InvalidUsername)); + } + } + + Ok(username) + } + + // Find a free discriminator for a given username + pub async fn find_discriminator( + db: &Database, + username: &str, + preferred: Option<(String, String)>, + ) -> Result { + let search_space: &HashSet = &DISCRIMINATOR_SEARCH_SPACE; + let used_discriminators: HashSet = db + .fetch_discriminators_in_use(username) + .await? + .into_iter() + .collect(); + + let available_discriminators: Vec<&String> = + search_space.difference(&used_discriminators).collect(); + + if available_discriminators.is_empty() { + return Err(create_error!(UsernameTaken)); + } + + if let Some((preferred, target_id)) = preferred { + if available_discriminators.contains(&&preferred) { + return Ok(preferred); + } else { + if db + .has_ratelimited( + &target_id, + crate::RatelimitEventType::DiscriminatorChange, + Duration::from_secs(60 * 60 * 24), + 1, + ) + .await? + { + return Err(create_error!(DiscriminatorChangeRatelimited)); + } + + db.insert_ratelimit_event(&RatelimitEvent { + id: Ulid::new().to_string(), + target_id, + event_type: crate::RatelimitEventType::DiscriminatorChange, + }) + .await?; + + /* let rvdb: revolt_database::Database = db.clone().into(); + if rvdb + .has_ratelimited( + &target_id, + RatelimitEventType::DiscriminatorChange, + Duration::from_secs(60 * 60 * 24), + 1, + ) + .await + .map_err(Error::from_core)? + { + return Err(Error::DiscriminatorChangeRatelimited); + } + + rvdb.insert_ratelimit_event(&revolt_database::RatelimitEvent { + id: ulid::Ulid::new().to_string(), + target_id, + event_type: RatelimitEventType::DiscriminatorChange, + }) + .await + .map_err(Error::from_core)?; */ + } + } + + let mut rng = rand::thread_rng(); + Ok(available_discriminators + .choose(&mut rng) + .expect("we can assert this has an element") + .to_string()) + } + /// Check whether a username is already in use by another user #[allow(dead_code)] async fn is_username_taken(db: &Database, username: &str) -> Result { @@ -203,17 +324,3 @@ impl User { .await } } - -pub static DISCRIMINATOR_SEARCH_SPACE: Lazy> = Lazy::new(|| { - let mut set = (2..9999) - .map(|v| format!("{:0>4}", v)) - .collect::>(); - - for discrim in [ - 123, 1234, 1111, 2222, 3333, 4444, 5555, 6666, 7777, 8888, 9999, - ] { - set.remove(&format!("{:0>4}", discrim)); - } - - set.into_iter().collect() -}); diff --git a/crates/core/database/src/models/users/ops.rs b/crates/core/database/src/models/users/ops.rs index 28421ec0..a07ceb87 100644 --- a/crates/core/database/src/models/users/ops.rs +++ b/crates/core/database/src/models/users/ops.rs @@ -22,6 +22,9 @@ pub trait AbstractUsers: Sync + Send { /// Fetch multiple users by their ids async fn fetch_users<'a>(&self, ids: &'a [String]) -> Result>; + /// Fetch all discriminators in use for a username + async fn fetch_discriminators_in_use(&self, username: &str) -> Result>; + /// Fetch ids of users that both users are friends with async fn fetch_mutual_user_ids(&self, user_a: &str, user_b: &str) -> Result>; diff --git a/crates/core/database/src/models/users/ops/mongodb.rs b/crates/core/database/src/models/users/ops/mongodb.rs index 6bea5077..9212c806 100644 --- a/crates/core/database/src/models/users/ops/mongodb.rs +++ b/crates/core/database/src/models/users/ops/mongodb.rs @@ -87,6 +87,39 @@ impl AbstractUsers for MongoDb { .await) } + /// Fetch all discriminators in use for a username + async fn fetch_discriminators_in_use(&self, username: &str) -> Result> { + #[derive(Deserialize)] + struct UserDocument { + discriminator: String, + } + + Ok(self + .col::(COL) + .find( + doc! { + "username": username + }, + FindOptions::builder() + .collation( + Collation::builder() + .locale("en") + .strength(CollationStrength::Secondary) + .build(), + ) + .projection(doc! { "_id": 0, "discriminator": 1 }) + .build(), + ) + .await + .map_err(|_| create_database_error!("find", COL))? + .filter_map(|s| async { s.ok() }) + .collect::>() + .await + .into_iter() + .map(|user| user.discriminator) + .collect::>()) + } + /// Fetch ids of users that both users are friends with async fn fetch_mutual_user_ids(&self, user_a: &str, user_b: &str) -> Result> { Ok(self diff --git a/crates/core/database/src/models/users/ops/reference.rs b/crates/core/database/src/models/users/ops/reference.rs index d75af7a6..c557bea5 100644 --- a/crates/core/database/src/models/users/ops/reference.rs +++ b/crates/core/database/src/models/users/ops/reference.rs @@ -56,6 +56,18 @@ impl AbstractUsers for ReferenceDb { .collect() } + /// Fetch all discriminators in use for a username + async fn fetch_discriminators_in_use(&self, username: &str) -> Result> { + let users = self.users.lock().await; + let lowercase = username.to_lowercase(); + Ok(users + .values() + .filter(|user| user.username.to_lowercase() == lowercase) + .map(|user| &user.discriminator) + .cloned() + .collect()) + } + /// Fetch ids of users that both users are friends with async fn fetch_mutual_user_ids(&self, _user_a: &str, _user_b: &str) -> Result> { todo!() diff --git a/crates/core/database/src/models/users/schema.rs b/crates/core/database/src/models/users/schema.rs new file mode 100644 index 00000000..d887b895 --- /dev/null +++ b/crates/core/database/src/models/users/schema.rs @@ -0,0 +1,31 @@ +use revolt_okapi::openapi3::{SecurityScheme, SecuritySchemeData}; +use revolt_rocket_okapi::{ + gen::OpenApiGenerator, + request::{OpenApiFromRequest, RequestHeaderInput}, +}; + +use crate::User; + +impl<'r> OpenApiFromRequest<'r> for User { + fn from_request_input( + _gen: &mut OpenApiGenerator, + _name: String, + _required: bool, + ) -> revolt_rocket_okapi::Result { + let mut requirements = schemars::Map::new(); + requirements.insert("Session Token".to_owned(), vec![]); + + Ok(RequestHeaderInput::Security( + "Session Token".to_owned(), + SecurityScheme { + data: SecuritySchemeData::ApiKey { + name: "x-session-token".to_owned(), + location: "header".to_owned(), + }, + description: Some("Used to authenticate as a user.".to_owned()), + extensions: schemars::Map::new(), + }, + requirements, + )) + } +} diff --git a/crates/core/models/Cargo.toml b/crates/core/models/Cargo.toml index 992f524f..b844283c 100644 --- a/crates/core/models/Cargo.toml +++ b/crates/core/models/Cargo.toml @@ -20,8 +20,12 @@ default = ["serde", "partials"] # Core revolt-permissions = { version = "0.6.5", path = "../permissions" } -# Serialisation +# Utility +regex = "1" indexmap = "1.9.3" +once_cell = "1.17.1" + +# Serialisation revolt_optional_struct = { version = "0.2.0", optional = true } serde = { version = "1", features = ["derive"], optional = true } iso8601-timestamp = { version = "0.2.11", features = ["schema", "bson"] } diff --git a/crates/core/models/src/v0/bots.rs b/crates/core/models/src/v0/bots.rs index 3b6a6051..72edc21f 100644 --- a/crates/core/models/src/v0/bots.rs +++ b/crates/core/models/src/v0/bots.rs @@ -1,7 +1,10 @@ use super::User; +use validator::Validate; + auto_derived!( /// Bot + #[derive(Default)] pub struct Bot { /// Bot Id #[cfg_attr(feature = "serde", serde(rename = "_id"))] @@ -87,12 +90,13 @@ auto_derived!( } /// Bot Details + #[cfg_attr(feature = "validator", derive(Validate))] pub struct DataCreateBot { /// Bot username #[cfg_attr( feature = "validator", - validate(length(min = 2, max = 32), regex = "RE_USERNAME") + validate(length(min = 2, max = 32), regex = "super::RE_USERNAME") )] - name: String, + pub name: String, } ); diff --git a/crates/core/models/src/v0/users.rs b/crates/core/models/src/v0/users.rs index 82fd9a91..51f1907e 100644 --- a/crates/core/models/src/v0/users.rs +++ b/crates/core/models/src/v0/users.rs @@ -1,5 +1,14 @@ +use once_cell::sync::Lazy; +use regex::Regex; + use super::File; +/// Regex for valid usernames +/// +/// Block zero width space +/// Block lookalike characters +pub static RE_USERNAME: Lazy = Lazy::new(|| Regex::new(r"^(\p{L}|[\d_.-])+$").unwrap()); + auto_derived!( /// User pub struct User { diff --git a/crates/core/result/src/lib.rs b/crates/core/result/src/lib.rs index 7115405d..9db79cf1 100644 --- a/crates/core/result/src/lib.rs +++ b/crates/core/result/src/lib.rs @@ -43,6 +43,7 @@ pub enum ErrorType { // ? User related errors UsernameTaken, InvalidUsername, + DiscriminatorChangeRatelimited, UnknownUser, AlreadyFriends, AlreadySentRequest, diff --git a/crates/core/result/src/rocket.rs b/crates/core/result/src/rocket.rs index a1723905..43b20373 100644 --- a/crates/core/result/src/rocket.rs +++ b/crates/core/result/src/rocket.rs @@ -19,6 +19,7 @@ impl<'r> Responder<'r, 'static> for Error { ErrorType::UnknownUser => Status::NotFound, ErrorType::InvalidUsername => Status::BadRequest, ErrorType::UsernameTaken => Status::Conflict, + ErrorType::DiscriminatorChangeRatelimited => Status::TooManyRequests, ErrorType::AlreadyFriends => Status::Conflict, ErrorType::AlreadySentRequest => Status::Conflict, ErrorType::Blocked => Status::Conflict, diff --git a/crates/delta/src/routes/bots/create.rs b/crates/delta/src/routes/bots/create.rs index 4a59ffa4..3462b9cd 100644 --- a/crates/delta/src/routes/bots/create.rs +++ b/crates/delta/src/routes/bots/create.rs @@ -1,9 +1,10 @@ -use crate::util::regex::RE_USERNAME; - use nanoid::nanoid; +use revolt_database::{Bot, BotInformation, Database, User}; +use revolt_models::v0; +use revolt_result::{create_error, Result}; use rocket::serde::json::Json; -use serde::Deserialize; +use rocket::State; use ulid::Ulid; use validator::Validate; @@ -12,17 +13,26 @@ use validator::Validate; /// Create a new Revolt bot. #[openapi(tag = "Bots")] #[post("/create", data = "")] -pub async fn create_bot(db: &Db, user: User, info: Json) -> Result> { +pub async fn create_bot( + db: &State, + user: User, + info: Json, +) -> Result> { if user.bot.is_some() { - return Err(Error::IsBot); + return Err(create_error!(IsBot)); } let info = info.into_inner(); - info.validate() - .map_err(|error| Error::FailedValidation { error })?; + info.validate().map_err(|error| { + create_error!(FailedValidation { + error: error.to_string() + }) + })?; - if db.get_number_of_bots_by_user(&user.id).await? >= *MAX_BOT_COUNT { - return Err(Error::ReachedMaximumBots); + // TODO: config + let max_bot_count = 5; + if db.get_number_of_bots_by_user(&user.id).await? >= max_bot_count { + return Err(create_error!(ReachedMaximumBots)); } let id = Ulid::new().to_string(); @@ -46,5 +56,5 @@ pub async fn create_bot(db: &Db, user: User, info: Json) -> Resul db.insert_user(&bot_user).await?; db.insert_bot(&bot).await?; - Ok(Json(bot)) + Ok(Json(bot.into())) }