feat: ratelimit user edit route and discriminator changes

Cargo.lock

@@ -2843,7 +2843,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-bonfire"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "async-std",
  "async-tungstenite",
@@ -2860,7 +2860,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-database"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "async-recursion",
  "async-std",
@@ -2891,7 +2891,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-delta"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "async-channel",
  "async-std",
@@ -2931,7 +2931,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-models"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "revolt-permissions",
  "revolt_optional_struct",
@@ -2942,7 +2942,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-permissions"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "async-std",
  "async-trait",
@@ -2956,7 +2956,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-presence"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "async-std",
  "log",
@@ -2967,7 +2967,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-quark"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "async-lock",
  "async-recursion",
@@ -2998,6 +2998,7 @@ dependencies = [
  "redis-kiss",
  "regex",
  "reqwest",
+ "revolt-database",
  "revolt-models",
  "revolt-presence",
  "revolt-result",
@@ -3019,7 +3020,7 @@ dependencies = [
 
 [[package]]
 name = "revolt-result"
-version = "0.6.2"
+version = "0.6.3"
 dependencies = [
  "revolt_okapi",
  "revolt_rocket_okapi",

crates/bonfire/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-bonfire"
-version = "0.6.2"
+version = "0.6.3"
 license = "AGPL-3.0-or-later"
 edition = "2021"
 

crates/core/database/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-database"
-version = "0.6.2"
+version = "0.6.3"
 edition = "2021"
 license = "AGPL-3.0-or-later"
 authors = [ "Paul Makles <me@insrt.uk>" ]
@@ -22,10 +22,10 @@ default = [ "mongodb", "async-std-runtime" ]
 
 [dependencies]
 # Core
-revolt-result = { version = "0.6.2", path = "../result" }
+revolt-result = { version = "0.6.3", path = "../result" }
-revolt-models = { version = "0.6.2", path = "../models" }
+revolt-models = { version = "0.6.3", path = "../models" }
-revolt-presence = { version = "0.6.2", path = "../presence" }
+revolt-presence = { version = "0.6.3", path = "../presence" }
-revolt-permissions = { version = "0.6.2", path = "../permissions", features = [ "serde", "bson" ] }
+revolt-permissions = { version = "0.6.3", path = "../permissions", features = [ "serde", "bson" ] }
 
 # Utility
 log = "0.4"

crates/core/database/src/models/admin_migrations/ops/mongodb/init.rs

@@ -76,6 +76,10 @@ pub async fn create_database(db: &MongoDb) {
         .await
         .expect("Failed to create bots collection.");
 
+    db.create_collection("ratelimit_events", None)
+        .await
+        .expect("Failed to create ratelimit_events collection.");
+
     db.create_collection(
         "pubsub",
         CreateCollectionOptions::builder()
@@ -209,5 +213,24 @@ pub async fn create_database(db: &MongoDb) {
         .await
         .expect("Failed to save migration info.");
 
+    db.run_command(
+        doc! {
+            "createIndexes": "ratelimit_events",
+            "indexes": [
+                {
+                    "key": {
+                        "_id": 1_i32,
+                        "target_id": 1_i32,
+                        "event_type": 1_i32,
+                    },
+                    "name": "compound_key"
+                }
+            ]
+        },
+        None,
+    )
+    .await
+    .expect("Failed to create ratelimit_events index.");
+
     info!("Created database.");
 }

crates/core/database/src/models/admin_migrations/ops/mongodb/scripts.rs

@@ -18,7 +18,7 @@ struct MigrationInfo {
     revision: i32,
 }
 
-pub const LATEST_REVISION: i32 = 25;
+pub const LATEST_REVISION: i32 = 26;
 
 pub async fn migrate_database(db: &MongoDb) {
     let migrations = db.col::<Document>("migrations");
@@ -947,8 +947,37 @@ pub async fn run_migrations(db: &MongoDb, revision: i32) -> i32 {
             .expect("Failed to create username index.");
     }
 
+    if revision <= 25 {
+        info!("Running migration [revision 25 / 15-06-2023]: Add collection `ratelimit_events` with index.");
+
+        db.db()
+            .create_collection("ratelimit_events", None)
+            .await
+            .ok();
+
+        db.db()
+            .run_command(
+                doc! {
+                    "createIndexes": "ratelimit_events",
+                    "indexes": [
+                        {
+                            "key": {
+                                "_id": 1_i32,
+                                "target_id": 1_i32,
+                                "event_type": 1_i32,
+                            },
+                            "name": "compound_key"
+                        }
+                    ]
+                },
+                None,
+            )
+            .await
+            .expect("Failed to create ratelimit_events index.");
+    }
+
     // Need to migrate fields on attachments, change `user_id`, `object_id`, etc to `parent`.
 
     // Reminder to update LATEST_REVISION when adding new migrations.
-    LATEST_REVISION
+    LATEST_REVISION.max(revision)
 }

crates/core/database/src/models/mod.rs

@@ -3,6 +3,7 @@ mod bots;
 mod channel_webhooks;
 mod channels;
 mod files;
+mod ratelimit_events;
 mod safety_strikes;
 mod server_members;
 mod servers;
@@ -14,6 +15,7 @@ pub use bots::*;
 pub use channel_webhooks::*;
 pub use channels::*;
 pub use files::*;
+pub use ratelimit_events::*;
 pub use safety_strikes::*;
 pub use server_members::*;
 pub use servers::*;
@@ -30,6 +32,7 @@ pub trait AbstractDatabase:
     + channels::AbstractChannels
     + channel_webhooks::AbstractWebhooks
     + files::AbstractAttachments
+    + ratelimit_events::AbstractRatelimitEvents
     + safety_strikes::AbstractAccountStrikes
     + server_members::AbstractServerMembers
     + servers::AbstractServers

crates/core/database/src/models/ratelimit_events/mod.rs

@@ -0,0 +1,5 @@
+mod model;
+mod ops;
+
+pub use model::*;
+pub use ops::*;

crates/core/database/src/models/ratelimit_events/model.rs

@@ -0,0 +1,25 @@
+use std::fmt;
+
+auto_derived!(
+    /// Ratelimit Event
+    pub struct RatelimitEvent {
+        /// Id
+        #[serde(rename = "_id")]
+        pub id: String,
+        /// Relevant Object Id
+        pub target_id: String,
+        /// Type of event
+        pub event_type: RatelimitEventType,
+    }
+
+    /// Event type
+    pub enum RatelimitEventType {
+        DiscriminatorChange,
+    }
+);
+
+impl fmt::Display for RatelimitEventType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Debug::fmt(self, f)
+    }
+}

crates/core/database/src/models/ratelimit_events/ops.rs

@@ -0,0 +1,20 @@
+use std::time::Duration;
+
+use crate::{revolt_result::Result, RatelimitEvent, RatelimitEventType};
+mod mongodb;
+mod reference;
+
+#[async_trait]
+pub trait AbstractRatelimitEvents: Sync + Send {
+    /// Insert a new ratelimit event
+    async fn insert_ratelimit_event(&self, event: &RatelimitEvent) -> Result<()>;
+
+    /// Count number of events in given duration and check if we've hit the limit
+    async fn has_ratelimited(
+        &self,
+        target_id: &str,
+        event_type: RatelimitEventType,
+        period: Duration,
+        count: usize,
+    ) -> Result<bool>;
+}

crates/core/database/src/models/ratelimit_events/ops/mongodb.rs

@@ -0,0 +1,40 @@
+use std::time::{Duration, SystemTime};
+
+use super::AbstractRatelimitEvents;
+use crate::{MongoDb, RatelimitEvent, RatelimitEventType};
+use revolt_result::Result;
+use ulid::Ulid;
+
+static COL: &str = "ratelimit_events";
+
+#[async_trait]
+impl AbstractRatelimitEvents for MongoDb {
+    /// Insert a new ratelimit event
+    async fn insert_ratelimit_event(&self, event: &RatelimitEvent) -> Result<()> {
+        query!(self, insert_one, COL, &event).map(|_| ())
+    }
+
+    /// Count number of events in given duration and check if we've hit the limit
+    async fn has_ratelimited(
+        &self,
+        target_id: &str,
+        event_type: RatelimitEventType,
+        period: Duration,
+        count: usize,
+    ) -> Result<bool> {
+        self.col::<RatelimitEvent>(COL)
+            .count_documents(
+                doc! {
+                    "_id": {
+                        "$gte": Ulid::from_datetime(SystemTime::now() - period).to_string()
+                    },
+                    "target_id": target_id,
+                    "event_type": event_type.to_string()
+                },
+                None,
+            )
+            .await
+            .map(|c| c as usize >= count)
+            .map_err(|_| create_database_error!("count_documents", COL))
+    }
+}

crates/core/database/src/models/ratelimit_events/ops/reference.rs

@@ -0,0 +1,28 @@
+use std::time::Duration;
+
+use super::AbstractRatelimitEvents;
+use crate::RatelimitEvent;
+use crate::RatelimitEventType;
+use crate::ReferenceDb;
+use revolt_result::Result;
+
+#[async_trait]
+impl AbstractRatelimitEvents for ReferenceDb {
+    /// Insert a new ratelimit event
+    async fn insert_ratelimit_event(&self, _event: &RatelimitEvent) -> Result<()> {
+        // TODO: implement
+        unimplemented!()
+    }
+
+    /// Count number of events in given duration and check if we've hit the limit
+    async fn has_ratelimited(
+        &self,
+        _target_id: &str,
+        _event_type: RatelimitEventType,
+        _period: Duration,
+        _count: usize,
+    ) -> Result<bool> {
+        // TODO: implement
+        unimplemented!()
+    }
+}

crates/core/models/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-models"
-version = "0.6.2"
+version = "0.6.3"
 edition = "2021"
 license = "AGPL-3.0-or-later"
 authors = [ "Paul Makles <me@insrt.uk>" ]
@@ -18,7 +18,7 @@ default = [ "serde", "partials" ]
 
 [dependencies]
 # Core
-revolt-permissions = { version = "0.6.2", path = "../permissions" }
+revolt-permissions = { version = "0.6.3", path = "../permissions" }
 
 # Serialisation
 revolt_optional_struct = { version = "0.2.0", optional = true }

crates/core/permissions/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-permissions"
-version = "0.6.2"
+version = "0.6.3"
 edition = "2021"
 license = "AGPL-3.0-or-later"
 authors = [ "Paul Makles <me@insrt.uk>" ]

crates/core/presence/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-presence"
-version = "0.6.2"
+version = "0.6.3"
 edition = "2021"
 license = "AGPL-3.0-or-later"
 authors = [ "Paul Makles <me@insrt.uk>" ]

crates/core/result/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-result"
-version = "0.6.2"
+version = "0.6.3"
 edition = "2021"
 license = "AGPL-3.0-or-later"
 authors = [ "Paul Makles <me@insrt.uk>" ]

crates/delta/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-delta"
-version = "0.6.2"
+version = "0.6.3"
 license = "AGPL-3.0-or-later"
 authors = ["Paul Makles <paulmakles@gmail.com>"]
 edition = "2018"

crates/quark/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "revolt-quark"
-version = "0.6.2"
+version = "0.6.3"
 edition = "2021"
 license = "AGPL-3.0-or-later"
 
@@ -94,4 +94,5 @@ sentry = "0.25.0"
 # Core
 revolt-result = { path = "../core/result", features = [ "serde", "schemas" ] }
 revolt-presence = { path = "../core/presence", features = [ "redis-is-patched" ] }
+revolt-database = { path = "../core/database" }
 revolt-models = { path = "../core/models" }

crates/quark/src/database.rs

@@ -71,3 +71,14 @@ impl From<Database> for authifier::Database {
         }
     }
 }
+
+impl From<Database> for revolt_database::Database {
+    fn from(val: Database) -> Self {
+        match val {
+            Database::Dummy(_) => revolt_database::Database::Reference(Default::default()),
+            Database::MongoDb(MongoDb(client)) => revolt_database::Database::MongoDb(
+                revolt_database::MongoDb(client, "revolt".to_string()),
+            ),
+        }
+    }
+}

crates/quark/src/impl/generic/users/user.rs

@@ -10,9 +10,11 @@ use futures::try_join;
 use impl_ops::impl_op_ex_commutative;
 use once_cell::sync::Lazy;
 use rand::seq::SliceRandom;
+use revolt_database::RatelimitEventType;
 use revolt_presence::filter_online;
 use std::collections::HashSet;
 use std::ops;
+use std::time::Duration;
 
 impl_op_ex_commutative!(+ |a: &i32, b: &Badges| -> i32 { *a | *b as i32 });
 
@@ -207,7 +209,7 @@ impl User {
     pub async fn find_discriminator(
         db: &Database,
         username: &str,
-        preferred: Option<String>,
+        preferred: Option<(String, String)>,
     ) -> Result<String> {
         let search_space: &HashSet<String> = &DISCRIMINATOR_SEARCH_SPACE_QUARK;
         let used_discriminators: HashSet<String> = db
@@ -223,9 +225,31 @@ impl User {
             return Err(Error::UsernameTaken);
         }
 
-        if let Some(preferred) = preferred {
+        if let Some((preferred, target_id)) = preferred {
             if available_discriminators.contains(&&preferred) {
                 return Ok(preferred);
+            } else {
+                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)?;
             }
         }
 
@@ -257,7 +281,7 @@ impl User {
                         User::find_discriminator(
                             db,
                             &username,
-                            Some(self.discriminator.to_string()),
+                            Some((self.discriminator.to_string(), self.id.clone())),
                         )
                         .await?,
                     ),

crates/quark/src/util/result.rs

@@ -31,6 +31,7 @@ pub enum Error {
     // ? User related errors
     UsernameTaken,
     InvalidUsername,
+    DiscriminatorChangeRatelimited,
     UnknownUser,
     AlreadyFriends,
     AlreadySentRequest,
@@ -165,6 +166,7 @@ impl<'r> Responder<'r, 'static> for Error {
 
             Error::UnknownUser => Status::NotFound,
             Error::InvalidUsername => Status::BadRequest,
+            Error::DiscriminatorChangeRatelimited => Status::TooManyRequests,
             Error::UsernameTaken => Status::Conflict,
             Error::AlreadyFriends => Status::Conflict,
             Error::AlreadySentRequest => Status::Conflict,

crates/quark/src/web/ratelimiter.rs

@@ -101,16 +101,19 @@ pub struct Ratelimiter {
 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", _) => {
+        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),
+            ("bots", _, _) => ("bots", None),
-            ("channels", Some(id)) => {
+            ("channels", Some(id), _) => {
                 if request.method() == Method::Post {
                     if let Some("messages") = request.routed_segment(2) {
                         return ("messaging", Some(id));
@@ -119,17 +122,17 @@ fn resolve_bucket<'r>(request: &'r rocket::Request<'_>) -> (&'r str, Option<&'r
 
                 ("channels", Some(id))
             }
-            ("servers", Some(id)) => ("servers", Some(id)),
+            ("servers", Some(id), _) => ("servers", Some(id)),
-            ("auth", _) => {
+            ("auth", _, _) => {
                 if request.method() == Method::Delete {
                     ("auth_delete", None)
                 } else {
                     ("auth", None)
                 }
             }
-            ("swagger", _) => ("swagger", None),
+            ("swagger", _, _) => ("swagger", None),
-            ("safety", Some("report")) => ("safety_report", Some("report")),
+            ("safety", Some("report"), _) => ("safety_report", Some("report")),
-            ("safety", _) => ("safety", None),
+            ("safety", _, _) => ("safety", None),
             _ => ("any", None),
         }
     } else {
@@ -140,6 +143,7 @@ fn resolve_bucket<'r>(request: &'r rocket::Request<'_>) -> (&'r str, Option<&'r
 /// Resolve per-bucket limits
 fn resolve_bucket_limit(bucket: &str) -> u8 {
     match bucket {
+        "user_edit" => 2,
         "users" => 20,
         "bots" => 10,
         "messaging" => 10,