2023-04-23 13:58:17 +02:00
|
|
|
use crate::{
|
2023-04-24 20:58:45 +02:00
|
|
|
api::limits::{Limit, LimitType, Limits, LimitsMutRef},
|
2023-04-23 13:58:17 +02:00
|
|
|
errors::InstanceServerError,
|
|
|
|
};
|
2023-04-10 17:20:02 +02:00
|
|
|
|
2023-04-12 18:33:16 +02:00
|
|
|
use reqwest::{Client, RequestBuilder, Response};
|
2023-04-24 19:49:26 +02:00
|
|
|
use std::collections::VecDeque;
|
2023-04-07 21:01:48 +02:00
|
|
|
|
2023-04-08 14:51:36 +02:00
|
|
|
// Note: There seem to be some overlapping request limiters. We need to make sure that sending a
|
|
|
|
// request checks for all the request limiters that apply, and blocks if any of the limiters are 0
|
|
|
|
|
2023-04-12 17:10:12 +02:00
|
|
|
#[allow(dead_code)]
|
2023-04-14 23:38:36 +02:00
|
|
|
#[derive(Debug)]
|
2023-04-14 18:29:20 +02:00
|
|
|
pub struct TypedRequest {
|
|
|
|
request: RequestBuilder,
|
|
|
|
limit_type: LimitType,
|
|
|
|
}
|
2023-04-14 23:38:36 +02:00
|
|
|
|
|
|
|
#[derive(Debug)]
|
2023-04-07 21:51:50 +02:00
|
|
|
pub struct LimitedRequester {
|
2023-04-07 21:01:48 +02:00
|
|
|
http: Client,
|
2023-04-14 18:29:20 +02:00
|
|
|
requests: VecDeque<TypedRequest>,
|
2023-04-07 21:01:48 +02:00
|
|
|
}
|
|
|
|
|
2023-04-07 21:51:50 +02:00
|
|
|
impl LimitedRequester {
|
|
|
|
/// Create a new `LimitedRequester`. `LimitedRequester`s use a `VecDeque` to store requests and
|
|
|
|
/// send them to the server using a `Client`. It keeps track of the remaining requests that can
|
|
|
|
/// be send within the `Limit` of an external API Ratelimiter, and looks at the returned request
|
|
|
|
/// headers to see if it can find Ratelimit info to update itself.
|
2023-04-12 17:10:12 +02:00
|
|
|
#[allow(dead_code)]
|
2023-04-25 17:41:14 +02:00
|
|
|
pub async fn new() -> Self {
|
2023-04-07 21:51:50 +02:00
|
|
|
LimitedRequester {
|
2023-04-07 21:01:48 +02:00
|
|
|
http: Client::new(),
|
2023-04-07 21:51:50 +02:00
|
|
|
requests: VecDeque::new(),
|
2023-04-07 21:01:48 +02:00
|
|
|
}
|
|
|
|
}
|
2023-04-12 18:33:16 +02:00
|
|
|
|
2023-04-14 23:01:48 +02:00
|
|
|
/**
|
|
|
|
# send_request
|
|
|
|
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`](reqwest::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`](crate::api::limits::LimitType) enum.
|
|
|
|
|
|
|
|
## Returns
|
|
|
|
- `Response`: The [`Response`](`reqwest::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`](`core::option::Option`)
|
|
|
|
- `None`: [`None`](`core::option::Option`) will be returned if the rate limit has been hit, and
|
|
|
|
the request could therefore not have been sent.
|
|
|
|
|
|
|
|
## Errors
|
|
|
|
|
|
|
|
This method will panic, if:
|
|
|
|
- The supplied [`RequestBuilder`](reqwest::RequestBuilder) contains invalid or incomplete
|
|
|
|
information
|
|
|
|
- There has been an error with processing (unwrapping) the [`Response`](`reqwest::Response`)
|
|
|
|
- The call to [`update_limits`](`crate::limits::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(
|
|
|
|
&mut self,
|
|
|
|
request: RequestBuilder,
|
|
|
|
limit_type: LimitType,
|
2023-04-24 20:58:45 +02:00
|
|
|
instance_rate_limits: &mut Limits,
|
|
|
|
user_rate_limits: &mut Limits,
|
2023-04-23 13:58:17 +02:00
|
|
|
) -> Result<Response, InstanceServerError> {
|
2023-04-24 20:58:45 +02:00
|
|
|
if self.can_send_request(limit_type, instance_rate_limits, user_rate_limits) {
|
2023-04-14 22:22:23 +02:00
|
|
|
let built_request = request
|
|
|
|
.build()
|
|
|
|
.unwrap_or_else(|e| panic!("Error while building the Request for sending: {}", e));
|
|
|
|
let result = self.http.execute(built_request).await;
|
|
|
|
let response = match result {
|
|
|
|
Ok(is_response) => is_response,
|
|
|
|
Err(e) => panic!("An error occured while processing the response: {}", e),
|
|
|
|
};
|
2023-04-24 20:58:45 +02:00
|
|
|
self.update_limits(
|
|
|
|
&response,
|
|
|
|
limit_type,
|
|
|
|
instance_rate_limits,
|
|
|
|
user_rate_limits,
|
|
|
|
);
|
2023-04-25 17:41:14 +02:00
|
|
|
Ok(response)
|
2023-04-14 22:22:23 +02:00
|
|
|
} else {
|
|
|
|
self.requests.push_back(TypedRequest {
|
2023-04-25 17:41:14 +02:00
|
|
|
request,
|
|
|
|
limit_type,
|
2023-04-14 22:22:23 +02:00
|
|
|
});
|
2023-04-25 17:41:14 +02:00
|
|
|
Err(InstanceServerError::RateLimited)
|
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(
|
|
|
|
&mut self,
|
|
|
|
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,
|
|
|
|
]
|
|
|
|
.to_vec();
|
|
|
|
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
|
|
|
|
.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
|
|
|
|
.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(
|
|
|
|
&mut self,
|
|
|
|
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-04-14 23:38:36 +02:00
|
|
|
use super::*;
|
2023-04-15 18:16:26 +02:00
|
|
|
use crate::{api::limits::Config, URLBundle};
|
2023-04-14 23:38:36 +02:00
|
|
|
#[tokio::test]
|
2023-04-15 17:07:46 +02:00
|
|
|
|
2023-04-15 13:27:43 +02:00
|
|
|
async fn create_limited_requester() {
|
2023-05-03 16:38:42 +02:00
|
|
|
let _urls = URLBundle::new(
|
2023-04-14 23:38:36 +02:00
|
|
|
String::from("http://localhost:3001/api/"),
|
|
|
|
String::from("wss://localhost:3001/"),
|
|
|
|
String::from("http://localhost:3001/cdn"),
|
|
|
|
);
|
2023-04-25 17:41:14 +02:00
|
|
|
let _requester = LimitedRequester::new().await;
|
2023-04-14 23:38:36 +02:00
|
|
|
}
|
2023-04-15 17:07:46 +02:00
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
async fn run_into_limit() {
|
|
|
|
let urls = URLBundle::new(
|
|
|
|
String::from("http://localhost:3001/api/"),
|
|
|
|
String::from("wss://localhost:3001/"),
|
|
|
|
String::from("http://localhost:3001/cdn"),
|
|
|
|
);
|
2023-04-25 17:41:14 +02:00
|
|
|
let mut requester = LimitedRequester::new().await;
|
2023-04-23 13:58:17 +02:00
|
|
|
let mut request: Option<Result<Response, InstanceServerError>> = 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";
|
|
|
|
let request_builder = requester.http.get(request_path);
|
2023-04-23 13:58:17 +02:00
|
|
|
request = Some(
|
|
|
|
requester
|
2023-04-24 20:58:45 +02:00
|
|
|
.send_request(
|
|
|
|
request_builder,
|
|
|
|
LimitType::Channel,
|
|
|
|
&mut instance_rate_limits,
|
|
|
|
&mut user_rate_limits,
|
|
|
|
)
|
2023-04-23 13:58:17 +02:00
|
|
|
.await,
|
|
|
|
);
|
2023-04-15 17:07:46 +02:00
|
|
|
}
|
2023-04-23 13:58:17 +02:00
|
|
|
if request.is_some() {
|
|
|
|
match request.unwrap() {
|
|
|
|
Ok(_) => assert!(false),
|
|
|
|
Err(_) => assert!(true),
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
assert!(false)
|
2023-04-15 17:07:46 +02:00
|
|
|
}
|
|
|
|
}
|
2023-04-15 18:16:26 +02:00
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
async fn test_send_request() {
|
|
|
|
let urls = URLBundle::new(
|
|
|
|
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-04-25 17:41:14 +02:00
|
|
|
let mut requester = LimitedRequester::new().await;
|
2023-04-15 18:16:26 +02:00
|
|
|
let request_path = urls.api.clone() + "/policies/instance/limits";
|
|
|
|
let request_builder = requester.http.get(request_path);
|
|
|
|
let request = requester
|
2023-04-24 20:58:45 +02:00
|
|
|
.send_request(
|
|
|
|
request_builder,
|
|
|
|
LimitType::Channel,
|
|
|
|
&mut instance_rate_limits,
|
|
|
|
&mut user_rate_limits,
|
|
|
|
)
|
2023-04-15 18:16:26 +02:00
|
|
|
.await;
|
|
|
|
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
|
|
|
}
|