diff --git a/crates/core/database/src/drivers/reference.rs b/crates/core/database/src/drivers/reference.rs index 6f6dd9c2..bfef6490 100644 --- a/crates/core/database/src/drivers/reference.rs +++ b/crates/core/database/src/drivers/reference.rs @@ -4,7 +4,7 @@ use futures::lock::Mutex; use crate::{ Bot, Channel, ChannelCompositeKey, ChannelUnread, Emoji, File, Invite, Member, - MemberCompositeKey, Message, Server, ServerBan, User, UserSettings, Webhook, + MemberCompositeKey, Message, RatelimitEvent, Server, ServerBan, User, UserSettings, Webhook, }; database_derived!( @@ -19,6 +19,7 @@ database_derived!( pub emojis: Arc>>, pub files: Arc>>, pub messages: Arc>>, + pub ratelimit_events: Arc>>, pub user_settings: Arc>>, pub users: Arc>>, pub server_bans: Arc>>, diff --git a/crates/core/database/src/models/ratelimit_events/ops/reference.rs b/crates/core/database/src/models/ratelimit_events/ops/reference.rs index 711dec51..751055ef 100644 --- a/crates/core/database/src/models/ratelimit_events/ops/reference.rs +++ b/crates/core/database/src/models/ratelimit_events/ops/reference.rs @@ -1,28 +1,46 @@ +use std::cmp::Ordering; use std::time::Duration; +use std::time::SystemTime; use super::AbstractRatelimitEvents; use crate::RatelimitEvent; use crate::RatelimitEventType; use crate::ReferenceDb; use revolt_result::Result; +use ulid::Ulid; #[async_trait] impl AbstractRatelimitEvents for ReferenceDb { /// Insert a new ratelimit event - async fn insert_ratelimit_event(&self, _event: &RatelimitEvent) -> Result<()> { - // TODO: implement - unimplemented!() + async fn insert_ratelimit_event(&self, event: &RatelimitEvent) -> Result<()> { + let mut ratelimit_events = self.ratelimit_events.lock().await; + if ratelimit_events.contains_key(&event.id) { + Err(create_database_error!("insert", "message")) + } else { + ratelimit_events.insert(event.id.to_string(), event.clone()); + Ok(()) + } } /// 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, + target_id: &str, + event_type: RatelimitEventType, + period: Duration, + count: usize, ) -> Result { - // TODO: implement - unimplemented!() + let ratelimit_events = self.ratelimit_events.lock().await; + let gte_cmp_id = Ulid::from_datetime(SystemTime::now() - period).to_string(); + + Ok(ratelimit_events + .iter() + .filter(|(id, event)| { + id.cmp(&>e_cmp_id) == Ordering::Greater + && event.target_id == target_id + && event.event_type == event_type + }) + .count() + >= count) } }