From a5943197d4bee6c435b4cfe8691a927249c2e812 Mon Sep 17 00:00:00 2001 From: bitfl0wer Date: Mon, 24 Apr 2023 20:58:45 +0200 Subject: [PATCH] separate User and Instance limits. --- src/api/auth/login.rs | 11 +++- src/api/auth/register.rs | 13 ++++- src/api/policies/instance/limits.rs | 78 +++++++++++++++++++++++++++++ src/api/schemas.rs | 2 + src/limit.rs | 49 ++++++++++++++---- 5 files changed, 140 insertions(+), 13 deletions(-) diff --git a/src/api/auth/login.rs b/src/api/auth/login.rs index 2cb61bd..a11db3c 100644 --- a/src/api/auth/login.rs +++ b/src/api/auth/login.rs @@ -17,8 +17,17 @@ pub mod login { let client = Client::new(); let endpoint_url = self.urls.get_api().to_string() + "/auth/login"; let request_builder = client.post(endpoint_url).body(json_schema.to_string()); + // We do not have a user yet, and the UserRateLimits will not be affected by a login + // request (since login is an instance wide limit), which is why we are just cloning the + // instances' limits to pass them on as user_rate_limits later. + let mut cloned_limits = self.limits.clone(); let response = requester - .send_request(request_builder, LimitType::AuthRegister, &mut self.limits) + .send_request( + request_builder, + LimitType::AuthRegister, + &mut self.limits, + &mut cloned_limits, + ) .await; if !response.is_ok() { return Err(InstanceServerError::NoResponse); diff --git a/src/api/auth/register.rs b/src/api/auth/register.rs index ef1062d..2e0db06 100644 --- a/src/api/auth/register.rs +++ b/src/api/auth/register.rs @@ -28,8 +28,17 @@ pub mod register { let client = Client::new(); let endpoint_url = self.urls.get_api().to_string() + "/auth/register"; let request_builder = client.post(endpoint_url).body(json_schema.to_string()); + // We do not have a user yet, and the UserRateLimits will not be affected by a login + // request (since register is an instance wide limit), which is why we are just cloning + // the instances' limits to pass them on as user_rate_limits later. + let mut cloned_limits = self.limits.clone(); let response = limited_requester - .send_request(request_builder, LimitType::AuthRegister, &mut self.limits) + .send_request( + request_builder, + LimitType::AuthRegister, + &mut self.limits, + &mut cloned_limits, + ) .await; if !response.is_ok() { return Err(InstanceServerError::NoResponse); @@ -111,7 +120,7 @@ mod test { AuthUsername::new("Hiiii".to_string()).unwrap(), Some(AuthPassword::new("mysupersecurepass123!".to_string()).unwrap()), true, - Some(AuthEmail::new("flori@aaaa.xyz".to_string()).unwrap()), + Some(AuthEmail::new("random978234@aaaa.xyz".to_string()).unwrap()), None, None, Some("2000-01-01".to_string()), diff --git a/src/api/policies/instance/limits.rs b/src/api/policies/instance/limits.rs index 299164b..c3c76a3 100644 --- a/src/api/policies/instance/limits.rs +++ b/src/api/policies/instance/limits.rs @@ -166,6 +166,69 @@ pub mod limits { } } + pub struct LimitsMutRef<'a> { + pub limit_absolute_messages: &'a mut Limit, + pub limit_absolute_register: &'a mut Limit, + pub limit_auth_login: &'a mut Limit, + pub limit_auth_register: &'a mut Limit, + pub limit_ip: &'a mut Limit, + pub limit_global: &'a mut Limit, + pub limit_error: &'a mut Limit, + pub limit_guild: &'a mut Limit, + pub limit_webhook: &'a mut Limit, + pub limit_channel: &'a mut Limit, + } + + impl LimitsMutRef<'_> { + pub fn combine_mut_ref<'a>( + instance_rate_limits: &'a mut Limits, + user_rate_limits: &'a mut Limits, + ) -> LimitsMutRef<'a> { + return LimitsMutRef { + limit_absolute_messages: &mut instance_rate_limits.limit_absolute_messages, + limit_absolute_register: &mut instance_rate_limits.limit_absolute_register, + limit_auth_login: &mut instance_rate_limits.limit_auth_login, + limit_auth_register: &mut instance_rate_limits.limit_auth_register, + limit_channel: &mut user_rate_limits.limit_channel, + limit_error: &mut user_rate_limits.limit_error, + limit_global: &mut instance_rate_limits.limit_global, + limit_guild: &mut user_rate_limits.limit_guild, + limit_ip: &mut instance_rate_limits.limit_ip, + limit_webhook: &mut user_rate_limits.limit_webhook, + }; + } + + pub fn get_limit_ref(&self, limit_type: &LimitType) -> &Limit { + match limit_type { + &LimitType::AbsoluteMessage => &self.limit_absolute_messages, + &LimitType::AbsoluteRegister => &self.limit_absolute_register, + &LimitType::AuthLogin => &self.limit_auth_login, + &LimitType::AuthRegister => &self.limit_auth_register, + &LimitType::Channel => &self.limit_channel, + &LimitType::Error => &self.limit_error, + &LimitType::Global => &self.limit_global, + &LimitType::Guild => &self.limit_guild, + &LimitType::Ip => &self.limit_ip, + &LimitType::Webhook => &self.limit_webhook, + } + } + + pub fn get_limit_mut_ref(&mut self, limit_type: &LimitType) -> &mut Limit { + match limit_type { + &LimitType::AbsoluteMessage => &mut self.limit_absolute_messages, + &LimitType::AbsoluteRegister => &mut self.limit_absolute_register, + &LimitType::AuthLogin => &mut self.limit_auth_login, + &LimitType::AuthRegister => &mut self.limit_auth_register, + &LimitType::Channel => &mut self.limit_channel, + &LimitType::Error => &mut self.limit_error, + &LimitType::Global => &mut self.limit_global, + &LimitType::Guild => &mut self.limit_guild, + &LimitType::Ip => &mut self.limit_ip, + &LimitType::Webhook => &mut self.limit_webhook, + } + } + } + #[derive(Debug, Clone)] pub struct Limits { pub limit_absolute_messages: Limit, @@ -181,6 +244,21 @@ pub mod limits { } impl Limits { + pub fn combine(instance_rate_limits: &Limits, user_rate_limits: &Limits) -> Limits { + return Limits { + limit_absolute_messages: instance_rate_limits.limit_absolute_messages, + limit_absolute_register: instance_rate_limits.limit_absolute_register, + limit_auth_login: instance_rate_limits.limit_auth_login, + limit_auth_register: instance_rate_limits.limit_auth_register, + limit_channel: user_rate_limits.limit_channel, + limit_error: user_rate_limits.limit_error, + limit_global: instance_rate_limits.limit_global, + limit_guild: user_rate_limits.limit_guild, + limit_ip: instance_rate_limits.limit_ip, + limit_webhook: user_rate_limits.limit_webhook, + }; + } + pub fn get_limit_ref(&self, limit_type: &LimitType) -> &Limit { match limit_type { &LimitType::AbsoluteMessage => &self.limit_absolute_messages, diff --git a/src/api/schemas.rs b/src/api/schemas.rs index d28e60e..0feb280 100644 --- a/src/api/schemas.rs +++ b/src/api/schemas.rs @@ -439,6 +439,7 @@ pub mod schemas { logged_in: bool, belongs_to: URLBundle, token: String, + rate_limits: Limits, settings: UserSettings, object: UserObject, ) -> User { @@ -446,6 +447,7 @@ pub mod schemas { logged_in, belongs_to, token, + rate_limits, settings, object, } diff --git a/src/limit.rs b/src/limit.rs index b69ad86..f7f89d5 100644 --- a/src/limit.rs +++ b/src/limit.rs @@ -1,5 +1,5 @@ use crate::{ - api::limits::{Limit, LimitType, Limits}, + api::limits::{Limit, LimitType, Limits, LimitsMutRef}, errors::InstanceServerError, }; @@ -68,9 +68,10 @@ impl LimitedRequester { &mut self, request: RequestBuilder, limit_type: LimitType, - rate_limits: &mut Limits, + instance_rate_limits: &mut Limits, + user_rate_limits: &mut Limits, ) -> Result { - if self.can_send_request(limit_type, rate_limits) { + if self.can_send_request(limit_type, instance_rate_limits, user_rate_limits) { let built_request = request .build() .unwrap_or_else(|e| panic!("Error while building the Request for sending: {}", e)); @@ -79,7 +80,12 @@ impl LimitedRequester { Ok(is_response) => is_response, Err(e) => panic!("An error occured while processing the response: {}", e), }; - self.update_limits(&response, limit_type, rate_limits); + self.update_limits( + &response, + limit_type, + instance_rate_limits, + user_rate_limits, + ); return Ok(response); } else { self.requests.push_back(TypedRequest { @@ -101,8 +107,16 @@ impl LimitedRequester { } } - fn can_send_request(&mut self, limit_type: LimitType, rate_limits: &Limits) -> bool { + fn can_send_request( + &mut self, + limit_type: LimitType, + instance_rate_limits: &Limits, + user_rate_limits: &Limits, + ) -> bool { // Check if all of the limits in this vec have at least one remaining request + + let rate_limits = Limits::combine(instance_rate_limits, user_rate_limits); + let constant_limits: Vec<&LimitType> = [ &LimitType::Error, &LimitType::Global, @@ -148,8 +162,11 @@ impl LimitedRequester { &mut self, response: &Response, limit_type: LimitType, - rate_limits: &mut Limits, + instance_rate_limits: &mut Limits, + user_rate_limits: &mut Limits, ) { + let mut rate_limits = LimitsMutRef::combine_mut_ref(instance_rate_limits, user_rate_limits); + let remaining = match response.headers().get("X-RateLimit-Remaining") { Some(remaining) => remaining.to_str().unwrap().parse::().unwrap(), None => rate_limits.get_limit_mut_ref(&limit_type).remaining - 1, @@ -261,14 +278,20 @@ mod rate_limit { ); let mut requester = LimitedRequester::new(urls.api.clone()).await; let mut request: Option> = None; - let mut limits = Limits::check_limits(urls.api.clone()).await; + let mut instance_rate_limits = Limits::check_limits(urls.api.clone()).await; + let mut user_rate_limits = Limits::check_limits(urls.api.clone()).await; for _ in 0..=50 { let request_path = urls.api.clone() + "/some/random/nonexisting/path"; let request_builder = requester.http.get(request_path); request = Some( requester - .send_request(request_builder, LimitType::Channel, &mut limits) + .send_request( + request_builder, + LimitType::Channel, + &mut instance_rate_limits, + &mut user_rate_limits, + ) .await, ); } @@ -289,12 +312,18 @@ mod rate_limit { String::from("wss://localhost:3001/"), String::from("http://localhost:3001/cdn"), ); - let mut limits = Limits::check_limits(urls.api.clone()).await; + let mut instance_rate_limits = Limits::check_limits(urls.api.clone()).await; + let mut user_rate_limits = Limits::check_limits(urls.api.clone()).await; let mut requester = LimitedRequester::new(urls.api.clone()).await; let request_path = urls.api.clone() + "/policies/instance/limits"; let request_builder = requester.http.get(request_path); let request = requester - .send_request(request_builder, LimitType::Channel, &mut limits) + .send_request( + request_builder, + LimitType::Channel, + &mut instance_rate_limits, + &mut user_rate_limits, + ) .await; let result = match request { Ok(result) => result,