Implement Limits::new() and more logic

This commit is contained in:
bitfl0wer 2023-04-24 19:38:42 +02:00
parent 3932df005b
commit 3e9aebf7ee
No known key found for this signature in database
GPG Key ID: 84BBB60DF895ABF2
1 changed files with 97 additions and 128 deletions

View File

@ -166,6 +166,7 @@ pub mod limits {
} }
} }
#[derive(Debug, Clone)]
pub struct Limits { pub struct Limits {
pub limit_absolute_messages: Limit, pub limit_absolute_messages: Limit,
pub limit_absolute_register: Limit, pub limit_absolute_register: Limit,
@ -180,19 +181,49 @@ pub mod limits {
} }
impl Limits { impl Limits {
pub fn iter(&self) -> std::vec::IntoIter<Limit> { pub fn get_limit_ref(&self, limit_type: &LimitType) -> &Limit {
let mut limits: Vec<Limit> = Vec::new(); match limit_type {
limits.push(self.limit_absolute_messages.clone()); &LimitType::AbsoluteMessage => &self.limit_absolute_messages,
limits.push(self.limit_absolute_register.clone()); &LimitType::AbsoluteRegister => &self.limit_absolute_register,
limits.push(self.limit_auth_login.clone()); &LimitType::AuthLogin => &self.limit_auth_login,
limits.push(self.limit_auth_register.clone()); &LimitType::AuthRegister => &self.limit_auth_register,
limits.push(self.limit_ip.clone()); &LimitType::Channel => &self.limit_channel,
limits.push(self.limit_global.clone()); &LimitType::Error => &self.limit_error,
limits.push(self.limit_error.clone()); &LimitType::Global => &self.limit_global,
limits.push(self.limit_guild.clone()); &LimitType::Guild => &self.limit_guild,
limits.push(self.limit_webhook.clone()); &LimitType::Ip => &self.limit_ip,
limits.push(self.limit_channel.clone()); &LimitType::Webhook => &self.limit_webhook,
limits.into_iter() }
}
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,
}
}
pub fn to_hash_map(&self) -> HashMap<LimitType, Limit> {
let mut map: HashMap<LimitType, Limit> = HashMap::new();
map.insert(LimitType::AbsoluteMessage, self.limit_absolute_messages);
map.insert(LimitType::AbsoluteRegister, self.limit_absolute_register);
map.insert(LimitType::AuthLogin, self.limit_auth_login);
map.insert(LimitType::AuthRegister, self.limit_auth_register);
map.insert(LimitType::Ip, self.limit_ip);
map.insert(LimitType::Global, self.limit_global);
map.insert(LimitType::Error, self.limit_error);
map.insert(LimitType::Guild, self.limit_guild);
map.insert(LimitType::Webhook, self.limit_webhook);
map.insert(LimitType::Channel, self.limit_channel);
map
} }
/// check_limits uses the API to get the current request limits of the instance. /// check_limits uses the API to get the current request limits of the instance.
@ -201,7 +232,7 @@ pub mod limits {
/// # Errors /// # Errors
/// This function will panic if the request fails or if the response body cannot be parsed. /// This function will panic if the request fails or if the response body cannot be parsed.
/// TODO: Change this to return a Result and handle the errors properly. /// TODO: Change this to return a Result and handle the errors properly.
pub async fn check_limits(api_url: String) -> HashMap<LimitType, Limit> { pub async fn check_limits(api_url: String) -> Limits {
let client = Client::new(); let client = Client::new();
let url_parsed = crate::URLBundle::parse_url(api_url) + "/policies/instance/limits"; let url_parsed = crate::URLBundle::parse_url(api_url) + "/policies/instance/limits";
let result = client let result = client
@ -219,213 +250,151 @@ pub mod limits {
}); });
let config: Config = from_str(&result).unwrap(); let config: Config = from_str(&result).unwrap();
// If config.rate.enabled is false, then add return a Limits struct with all limits set to u64::MAX // If config.rate.enabled is false, then add return a Limits struct with all limits set to u64::MAX
let mut limits: HashMap<LimitType, Limit> = HashMap::new(); let mut limits: Limits;
if config.rate.enabled == false { if config.rate.enabled == false {
limits.insert( limits = Limits {
LimitType::AbsoluteMessage, limit_absolute_messages: Limit {
Limit {
bucket: LimitType::AbsoluteMessage, bucket: LimitType::AbsoluteMessage,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_absolute_register: Limit {
limits.insert(
LimitType::AbsoluteRegister,
Limit {
bucket: LimitType::AbsoluteRegister, bucket: LimitType::AbsoluteRegister,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_auth_login: Limit {
limits.insert(
LimitType::AuthLogin,
Limit {
bucket: LimitType::AuthLogin, bucket: LimitType::AuthLogin,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_auth_register: Limit {
limits.insert(
LimitType::AuthRegister,
Limit {
bucket: LimitType::AuthRegister, bucket: LimitType::AuthRegister,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_ip: Limit {
limits.insert(
LimitType::Ip,
Limit {
bucket: LimitType::Ip, bucket: LimitType::Ip,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_global: Limit {
limits.insert(
LimitType::Global,
Limit {
bucket: LimitType::Global, bucket: LimitType::Global,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_error: Limit {
limits.insert(
LimitType::Error,
Limit {
bucket: LimitType::Error, bucket: LimitType::Error,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_guild: Limit {
limits.insert(
LimitType::Guild,
Limit {
bucket: LimitType::Guild, bucket: LimitType::Guild,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_webhook: Limit {
limits.insert(
LimitType::Webhook,
Limit {
bucket: LimitType::Webhook, bucket: LimitType::Webhook,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); limit_channel: Limit {
limits.insert(
LimitType::Channel,
Limit {
bucket: LimitType::Channel, bucket: LimitType::Channel,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, },
); };
} else { } else {
limits.insert( limits = Limits {
LimitType::AbsoluteMessage, limit_absolute_messages: Limit {
Limit {
bucket: LimitType::AbsoluteMessage, bucket: LimitType::AbsoluteMessage,
limit: config.absoluteRate.sendMessage.limit, limit: config.absoluteRate.sendMessage.limit,
remaining: config.absoluteRate.sendMessage.limit, remaining: config.absoluteRate.sendMessage.limit,
reset: config.absoluteRate.sendMessage.window, reset: config.absoluteRate.sendMessage.window,
}, },
); limit_absolute_register: Limit {
limits.insert(
LimitType::AbsoluteRegister,
Limit {
bucket: LimitType::AbsoluteRegister, bucket: LimitType::AbsoluteRegister,
limit: config.absoluteRate.register.limit, limit: config.absoluteRate.register.limit,
remaining: config.absoluteRate.register.limit, remaining: config.absoluteRate.register.limit,
reset: config.absoluteRate.register.window, reset: config.absoluteRate.register.window,
}, },
); limit_auth_login: Limit {
limits.insert(
LimitType::AuthLogin,
Limit {
bucket: LimitType::AuthLogin, bucket: LimitType::AuthLogin,
limit: config.rate.routes.auth.login.count, limit: config.rate.routes.auth.login.count,
remaining: config.rate.routes.auth.login.count, remaining: config.rate.routes.auth.login.count,
reset: config.rate.routes.auth.login.window, reset: config.rate.routes.auth.login.window,
}, },
); limit_auth_register: Limit {
limits.insert(
LimitType::AuthRegister,
Limit {
bucket: LimitType::AuthRegister, bucket: LimitType::AuthRegister,
limit: config.rate.routes.auth.register.count, limit: config.rate.routes.auth.register.count,
remaining: config.rate.routes.auth.register.count, remaining: config.rate.routes.auth.register.count,
reset: config.rate.routes.auth.register.window, reset: config.rate.routes.auth.register.window,
}, },
); limit_ip: Limit {
limits.insert(
LimitType::Guild,
Limit {
bucket: LimitType::Guild,
limit: config.rate.routes.guild.count,
remaining: config.rate.routes.guild.count,
reset: config.rate.routes.guild.window,
},
);
limits.insert(
LimitType::Webhook,
Limit {
bucket: LimitType::Webhook,
limit: config.rate.routes.webhook.count,
remaining: config.rate.routes.webhook.count,
reset: config.rate.routes.webhook.window,
},
);
limits.insert(
LimitType::Channel,
Limit {
bucket: LimitType::Channel,
limit: config.rate.routes.channel.count,
remaining: config.rate.routes.channel.count,
reset: config.rate.routes.channel.window,
},
);
limits.insert(
LimitType::Ip,
Limit {
bucket: LimitType::Ip, bucket: LimitType::Ip,
limit: config.rate.ip.count, limit: config.rate.ip.count,
remaining: config.rate.ip.count, remaining: config.rate.ip.count,
reset: config.rate.ip.window, reset: config.rate.ip.window,
}, },
); limit_global: Limit {
limits.insert(
LimitType::Global,
Limit {
bucket: LimitType::Global, bucket: LimitType::Global,
limit: config.rate.global.count, limit: config.rate.global.count,
remaining: config.rate.global.count, remaining: config.rate.global.count,
reset: config.rate.global.window, reset: config.rate.global.window,
}, },
); limit_error: Limit {
limits.insert(
LimitType::Error,
Limit {
bucket: LimitType::Error, bucket: LimitType::Error,
limit: config.rate.error.count, limit: config.rate.error.count,
remaining: config.rate.error.count, remaining: config.rate.error.count,
reset: config.rate.error.window, reset: config.rate.error.window,
}, },
); limit_guild: Limit {
bucket: LimitType::Guild,
limit: config.rate.routes.guild.count,
remaining: config.rate.routes.guild.count,
reset: config.rate.routes.guild.window,
},
limit_webhook: Limit {
bucket: LimitType::Webhook,
limit: config.rate.routes.webhook.count,
remaining: config.rate.routes.webhook.count,
reset: config.rate.routes.webhook.window,
},
limit_channel: Limit {
bucket: LimitType::Channel,
limit: config.rate.routes.channel.count,
remaining: config.rate.routes.channel.count,
reset: config.rate.routes.channel.window,
},
};
} }
if !config.absoluteRate.register.enabled { if !config.absoluteRate.register.enabled {
limits.insert( limits.limit_absolute_register = Limit {
LimitType::AbsoluteRegister,
Limit {
bucket: LimitType::AbsoluteRegister, bucket: LimitType::AbsoluteRegister,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, };
);
} }
if !config.absoluteRate.sendMessage.enabled { if !config.absoluteRate.sendMessage.enabled {
limits.insert( limits.limit_absolute_messages = Limit {
LimitType::AbsoluteMessage,
Limit {
bucket: LimitType::AbsoluteMessage, bucket: LimitType::AbsoluteMessage,
limit: u64::MAX, limit: u64::MAX,
remaining: u64::MAX, remaining: u64::MAX,
reset: u64::MAX, reset: u64::MAX,
}, };
);
} }
return limits; return limits;