2023-06-11 13:52:31 +02:00
|
|
|
use reqwest::{Client, RequestBuilder, Response};
|
|
|
|
|
2023-04-23 13:58:17 +02:00
|
|
|
use crate::{
|
2023-06-11 13:52:53 +02:00
|
|
|
api::limits::{Limit, LimitType, Limits, LimitsMutRef},
|
2023-06-08 22:16:23 +02:00
|
|
|
errors::ChorusLibError,
|
2023-04-23 13:58:17 +02:00
|
|
|
};
|
2023-04-10 17:20:02 +02:00
|
|
|
|
2023-04-14 23:38:36 +02:00
|
|
|
#[derive(Debug)]
|
2023-06-12 18:51:54 +02:00
|
|
|
pub struct LimitedRequester;
|
2023-04-07 21:01:48 +02:00
|
|
|
|
2023-04-07 21:51:50 +02:00
|
|
|
impl LimitedRequester {
|
2023-06-20 22:20:06 +02:00
|
|
|
/// Checks if a request can be sent without hitting API rate limits and sends it, if true.
|
|
|
|
/// Will automatically update the rate limits of the LimitedRequester the request has been
|
|
|
|
/// sent with.
|
|
|
|
///
|
|
|
|
/// # Arguments
|
|
|
|
///
|
|
|
|
/// * `request`: A `RequestBuilder` that contains a request ready to be sent. Unfinished or
|
|
|
|
/// invalid requests will result in the method panicing.
|
|
|
|
/// * `limit_type`: Because this library does not yet implement a way to check for which rate
|
|
|
|
/// limit will be used when the request gets send, you will have to specify this manually using
|
|
|
|
/// a `LimitType` enum.
|
|
|
|
///
|
|
|
|
/// # Returns
|
|
|
|
///
|
|
|
|
/// * `Response`: The `Response` gotten from sending the request to the server. This will be
|
|
|
|
/// returned if the Request was built and send successfully. Is wrapped in an `Option`.
|
|
|
|
/// * `None`: `None` will be returned if the rate limit has been hit, and the request could
|
|
|
|
/// therefore not have been sent.
|
|
|
|
///
|
|
|
|
/// # Errors
|
|
|
|
///
|
|
|
|
/// This method will error if:
|
|
|
|
///
|
|
|
|
/// * The request does not return a success status code (200-299)
|
|
|
|
/// * The supplied `RequestBuilder` contains invalid or incomplete information
|
|
|
|
/// * There has been an error with processing (unwrapping) the `Response`
|
|
|
|
/// * The call to `update_limits` yielded errors. Read the methods' Errors section for more
|
|
|
|
/// information.
|
2023-04-14 22:40:13 +02:00
|
|
|
pub async fn send_request(
|
|
|
|
request: RequestBuilder,
|
|
|
|
limit_type: LimitType,
|
2023-04-24 20:58:45 +02:00
|
|
|
instance_rate_limits: &mut Limits,
|
|
|
|
user_rate_limits: &mut Limits,
|
2023-06-08 22:16:23 +02:00
|
|
|
) -> Result<Response, ChorusLibError> {
|
2023-06-12 18:51:54 +02:00
|
|
|
if LimitedRequester::can_send_request(limit_type, instance_rate_limits, user_rate_limits) {
|
2023-05-13 22:10:46 +02:00
|
|
|
let built_request = match request.build() {
|
|
|
|
Ok(request) => request,
|
|
|
|
Err(e) => {
|
2023-06-08 22:16:23 +02:00
|
|
|
return Err(ChorusLibError::RequestErrorError {
|
2023-05-13 22:10:46 +02:00
|
|
|
url: "".to_string(),
|
|
|
|
error: e.to_string(),
|
2023-06-11 13:52:31 +02:00
|
|
|
});
|
2023-05-13 22:10:46 +02:00
|
|
|
}
|
|
|
|
};
|
2023-06-12 18:51:54 +02:00
|
|
|
let result = Client::new().execute(built_request).await;
|
2023-04-14 22:22:23 +02:00
|
|
|
let response = match result {
|
|
|
|
Ok(is_response) => is_response,
|
2023-05-13 22:06:44 +02:00
|
|
|
Err(e) => {
|
2023-06-08 22:16:23 +02:00
|
|
|
return Err(ChorusLibError::ReceivedErrorCodeError {
|
2023-05-13 22:06:44 +02:00
|
|
|
error_code: e.to_string(),
|
2023-06-11 13:52:31 +02:00
|
|
|
});
|
2023-05-13 22:06:44 +02:00
|
|
|
}
|
2023-04-14 22:22:23 +02:00
|
|
|
};
|
2023-06-12 18:51:54 +02:00
|
|
|
LimitedRequester::update_limits(
|
2023-04-24 20:58:45 +02:00
|
|
|
&response,
|
|
|
|
limit_type,
|
|
|
|
instance_rate_limits,
|
|
|
|
user_rate_limits,
|
|
|
|
);
|
2023-05-13 22:06:44 +02:00
|
|
|
if !response.status().is_success() {
|
2023-05-14 13:07:46 +02:00
|
|
|
match response.status().as_u16() {
|
2023-06-19 10:27:32 +02:00
|
|
|
401 => Err(ChorusLibError::TokenExpired),
|
|
|
|
403 => Err(ChorusLibError::TokenExpired),
|
|
|
|
_ => Err(ChorusLibError::ReceivedErrorCodeError {
|
|
|
|
error_code: response.status().as_str().to_string(),
|
|
|
|
}),
|
2023-05-14 13:07:46 +02:00
|
|
|
}
|
2023-05-13 22:06:44 +02:00
|
|
|
} else {
|
|
|
|
Ok(response)
|
|
|
|
}
|
2023-04-14 22:22:23 +02:00
|
|
|
} else {
|
2023-06-08 22:16:23 +02:00
|
|
|
Err(ChorusLibError::RateLimited {
|
2023-05-14 13:12:02 +02:00
|
|
|
bucket: limit_type.to_string(),
|
|
|
|
})
|
2023-04-14 22:22:23 +02:00
|
|
|
}
|
2023-04-12 18:33:16 +02:00
|
|
|
}
|
|
|
|
|
2023-04-14 17:09:50 +02:00
|
|
|
fn update_limit_entry(entry: &mut Limit, reset: u64, remaining: u64, limit: u64) {
|
|
|
|
if reset != entry.reset {
|
|
|
|
entry.reset = reset;
|
|
|
|
entry.remaining = limit;
|
|
|
|
entry.limit = limit;
|
|
|
|
} else {
|
|
|
|
entry.remaining = remaining;
|
|
|
|
entry.limit = limit;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-04-24 20:58:45 +02:00
|
|
|
fn can_send_request(
|
|
|
|
limit_type: LimitType,
|
|
|
|
instance_rate_limits: &Limits,
|
|
|
|
user_rate_limits: &Limits,
|
|
|
|
) -> bool {
|
2023-04-20 19:47:08 +02:00
|
|
|
// Check if all of the limits in this vec have at least one remaining request
|
2023-04-24 20:58:45 +02:00
|
|
|
|
|
|
|
let rate_limits = Limits::combine(instance_rate_limits, user_rate_limits);
|
|
|
|
|
2023-04-20 19:47:08 +02:00
|
|
|
let constant_limits: Vec<&LimitType> = [
|
|
|
|
&LimitType::Error,
|
|
|
|
&LimitType::Global,
|
|
|
|
&LimitType::Ip,
|
|
|
|
&limit_type,
|
|
|
|
]
|
2023-06-11 13:52:53 +02:00
|
|
|
.to_vec();
|
2023-04-20 19:47:08 +02:00
|
|
|
for limit in constant_limits.iter() {
|
2023-04-25 17:41:14 +02:00
|
|
|
match rate_limits.to_hash_map().get(limit) {
|
2023-04-20 19:47:08 +02:00
|
|
|
Some(limit) => {
|
|
|
|
if limit.remaining == 0 {
|
|
|
|
return false;
|
|
|
|
}
|
2023-04-20 20:11:12 +02:00
|
|
|
// AbsoluteRegister and AuthRegister can cancel each other out.
|
|
|
|
if limit.bucket == LimitType::AbsoluteRegister
|
2023-04-24 19:49:26 +02:00
|
|
|
&& rate_limits
|
2023-06-11 13:52:53 +02:00
|
|
|
.to_hash_map()
|
|
|
|
.get(&LimitType::AuthRegister)
|
|
|
|
.unwrap()
|
|
|
|
.remaining
|
|
|
|
== 0
|
2023-04-20 20:11:12 +02:00
|
|
|
{
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
if limit.bucket == LimitType::AuthRegister
|
2023-04-24 19:49:26 +02:00
|
|
|
&& rate_limits
|
2023-06-11 13:52:53 +02:00
|
|
|
.to_hash_map()
|
|
|
|
.get(&LimitType::AbsoluteRegister)
|
|
|
|
.unwrap()
|
|
|
|
.remaining
|
|
|
|
== 0
|
2023-04-20 20:11:12 +02:00
|
|
|
{
|
|
|
|
return false;
|
|
|
|
}
|
2023-04-14 18:29:20 +02:00
|
|
|
}
|
2023-04-20 19:47:08 +02:00
|
|
|
None => return false,
|
2023-04-14 18:29:20 +02:00
|
|
|
}
|
|
|
|
}
|
2023-04-25 17:41:14 +02:00
|
|
|
true
|
2023-04-14 17:09:50 +02:00
|
|
|
}
|
|
|
|
|
2023-04-24 19:49:26 +02:00
|
|
|
fn update_limits(
|
|
|
|
response: &Response,
|
|
|
|
limit_type: LimitType,
|
2023-04-24 20:58:45 +02:00
|
|
|
instance_rate_limits: &mut Limits,
|
|
|
|
user_rate_limits: &mut Limits,
|
2023-04-24 19:49:26 +02:00
|
|
|
) {
|
2023-04-24 20:58:45 +02:00
|
|
|
let mut rate_limits = LimitsMutRef::combine_mut_ref(instance_rate_limits, user_rate_limits);
|
|
|
|
|
2023-04-12 18:33:16 +02:00
|
|
|
let remaining = match response.headers().get("X-RateLimit-Remaining") {
|
2023-04-12 22:59:08 +02:00
|
|
|
Some(remaining) => remaining.to_str().unwrap().parse::<u64>().unwrap(),
|
2023-04-24 19:49:26 +02:00
|
|
|
None => rate_limits.get_limit_mut_ref(&limit_type).remaining - 1,
|
2023-04-12 18:33:16 +02:00
|
|
|
};
|
|
|
|
let limit = match response.headers().get("X-RateLimit-Limit") {
|
2023-04-12 22:59:08 +02:00
|
|
|
Some(limit) => limit.to_str().unwrap().parse::<u64>().unwrap(),
|
2023-04-24 19:49:26 +02:00
|
|
|
None => rate_limits.get_limit_mut_ref(&limit_type).limit,
|
2023-04-12 18:33:16 +02:00
|
|
|
};
|
|
|
|
let reset = match response.headers().get("X-RateLimit-Reset") {
|
2023-04-12 22:59:08 +02:00
|
|
|
Some(reset) => reset.to_str().unwrap().parse::<u64>().unwrap(),
|
2023-04-24 19:49:26 +02:00
|
|
|
None => rate_limits.get_limit_mut_ref(&limit_type).reset,
|
2023-04-12 18:33:16 +02:00
|
|
|
};
|
2023-04-12 22:59:08 +02:00
|
|
|
|
|
|
|
let status = response.status();
|
|
|
|
let status_str = status.as_str();
|
|
|
|
|
2023-04-25 17:41:14 +02:00
|
|
|
if status_str.starts_with('4') {
|
2023-04-24 19:49:26 +02:00
|
|
|
rate_limits
|
|
|
|
.get_limit_mut_ref(&LimitType::Error)
|
2023-04-15 17:07:46 +02:00
|
|
|
.add_remaining(-1);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
|
2023-04-24 19:49:26 +02:00
|
|
|
rate_limits
|
|
|
|
.get_limit_mut_ref(&LimitType::Global)
|
2023-04-15 17:07:46 +02:00
|
|
|
.add_remaining(-1);
|
|
|
|
|
2023-04-24 19:49:26 +02:00
|
|
|
rate_limits
|
|
|
|
.get_limit_mut_ref(&LimitType::Ip)
|
2023-04-15 17:07:46 +02:00
|
|
|
.add_remaining(-1);
|
|
|
|
|
2023-04-12 22:59:08 +02:00
|
|
|
match limit_type {
|
2023-04-15 17:12:33 +02:00
|
|
|
LimitType::Error => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::Error);
|
2023-04-15 17:12:33 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
|
|
|
}
|
|
|
|
LimitType::Global => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::Global);
|
2023-04-15 17:12:33 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
|
|
|
}
|
|
|
|
LimitType::Ip => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::Ip);
|
2023-04-15 17:12:33 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
|
|
|
}
|
2023-04-12 22:59:08 +02:00
|
|
|
LimitType::AuthLogin => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::AuthLogin);
|
2023-04-15 14:55:50 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::AbsoluteRegister => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::AbsoluteRegister);
|
2023-04-15 14:55:50 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
// AbsoluteRegister and AuthRegister both need to be updated, if a Register event
|
|
|
|
// happens.
|
2023-04-24 19:49:26 +02:00
|
|
|
rate_limits
|
|
|
|
.get_limit_mut_ref(&LimitType::AuthRegister)
|
2023-04-12 22:59:08 +02:00
|
|
|
.remaining -= 1;
|
|
|
|
}
|
|
|
|
LimitType::AuthRegister => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::AuthRegister);
|
2023-04-15 14:55:50 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
// AbsoluteRegister and AuthRegister both need to be updated, if a Register event
|
|
|
|
// happens.
|
2023-04-24 19:49:26 +02:00
|
|
|
rate_limits
|
|
|
|
.get_limit_mut_ref(&LimitType::AbsoluteRegister)
|
2023-04-12 22:59:08 +02:00
|
|
|
.remaining -= 1;
|
|
|
|
}
|
|
|
|
LimitType::AbsoluteMessage => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::AbsoluteMessage);
|
2023-04-15 14:55:50 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::Channel => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::Channel);
|
2023-04-15 14:55:50 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::Guild => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::Guild);
|
2023-04-15 14:55:50 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::Webhook => {
|
2023-04-24 19:49:26 +02:00
|
|
|
let entry = rate_limits.get_limit_mut_ref(&LimitType::Webhook);
|
2023-04-15 14:55:50 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
}
|
2023-04-12 18:33:16 +02:00
|
|
|
}
|
2023-04-07 21:01:48 +02:00
|
|
|
}
|
2023-04-14 23:38:36 +02:00
|
|
|
|
|
|
|
#[cfg(test)]
|
2023-04-15 13:27:43 +02:00
|
|
|
mod rate_limit {
|
2023-04-15 18:16:26 +02:00
|
|
|
use serde_json::from_str;
|
|
|
|
|
2023-06-20 02:59:18 +02:00
|
|
|
use crate::{api::limits::Config, UrlBundle};
|
2023-04-15 17:07:46 +02:00
|
|
|
|
2023-06-11 13:52:31 +02:00
|
|
|
use super::*;
|
|
|
|
|
2023-04-15 17:07:46 +02:00
|
|
|
#[tokio::test]
|
|
|
|
async fn run_into_limit() {
|
2023-06-20 02:59:18 +02:00
|
|
|
let urls = UrlBundle::new(
|
2023-04-15 17:07:46 +02:00
|
|
|
String::from("http://localhost:3001/api/"),
|
|
|
|
String::from("wss://localhost:3001/"),
|
|
|
|
String::from("http://localhost:3001/cdn"),
|
|
|
|
);
|
2023-06-08 22:16:23 +02:00
|
|
|
let mut request: Option<Result<Response, ChorusLibError>> = None;
|
2023-04-24 20:58:45 +02:00
|
|
|
let mut instance_rate_limits = Limits::check_limits(urls.api.clone()).await;
|
|
|
|
let mut user_rate_limits = Limits::check_limits(urls.api.clone()).await;
|
2023-04-15 17:07:46 +02:00
|
|
|
|
2023-04-23 13:58:17 +02:00
|
|
|
for _ in 0..=50 {
|
2023-04-15 17:07:46 +02:00
|
|
|
let request_path = urls.api.clone() + "/some/random/nonexisting/path";
|
2023-06-12 18:51:54 +02:00
|
|
|
let request_builder = Client::new().get(request_path);
|
2023-04-23 13:58:17 +02:00
|
|
|
request = Some(
|
2023-06-12 18:51:54 +02:00
|
|
|
LimitedRequester::send_request(
|
|
|
|
request_builder,
|
|
|
|
LimitType::Channel,
|
|
|
|
&mut instance_rate_limits,
|
|
|
|
&mut user_rate_limits,
|
|
|
|
)
|
|
|
|
.await,
|
2023-04-23 13:58:17 +02:00
|
|
|
);
|
2023-04-15 17:07:46 +02:00
|
|
|
}
|
2023-06-19 10:27:32 +02:00
|
|
|
assert!(matches!(request, Some(Err(_))));
|
2023-04-15 17:07:46 +02:00
|
|
|
}
|
2023-04-15 18:16:26 +02:00
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
async fn test_send_request() {
|
2023-06-20 02:59:18 +02:00
|
|
|
let urls = UrlBundle::new(
|
2023-04-15 18:16:26 +02:00
|
|
|
String::from("http://localhost:3001/api/"),
|
|
|
|
String::from("wss://localhost:3001/"),
|
|
|
|
String::from("http://localhost:3001/cdn"),
|
|
|
|
);
|
2023-04-24 20:58:45 +02:00
|
|
|
let mut instance_rate_limits = Limits::check_limits(urls.api.clone()).await;
|
|
|
|
let mut user_rate_limits = Limits::check_limits(urls.api.clone()).await;
|
2023-06-12 18:58:53 +02:00
|
|
|
let _requester = LimitedRequester;
|
2023-04-15 18:16:26 +02:00
|
|
|
let request_path = urls.api.clone() + "/policies/instance/limits";
|
2023-06-12 18:51:54 +02:00
|
|
|
let request_builder = Client::new().get(request_path);
|
|
|
|
let request = LimitedRequester::send_request(
|
|
|
|
request_builder,
|
|
|
|
LimitType::Channel,
|
|
|
|
&mut instance_rate_limits,
|
|
|
|
&mut user_rate_limits,
|
|
|
|
)
|
|
|
|
.await;
|
2023-04-15 18:16:26 +02:00
|
|
|
let result = match request {
|
2023-04-23 13:58:17 +02:00
|
|
|
Ok(result) => result,
|
|
|
|
Err(_) => panic!("Request failed"),
|
2023-04-15 18:16:26 +02:00
|
|
|
};
|
2023-05-12 12:35:06 +02:00
|
|
|
let _config: Config = from_str(result.text().await.unwrap().as_str()).unwrap();
|
2023-04-15 18:16:26 +02:00
|
|
|
}
|
2023-04-14 23:38:36 +02:00
|
|
|
}
|