2023-04-12 22:59:08 +02:00
|
|
|
use crate::api::limits::{Limit, LimitType, Limits};
|
2023-04-10 17:20:02 +02:00
|
|
|
|
2023-04-12 18:33:16 +02:00
|
|
|
use reqwest::{Client, RequestBuilder, Response};
|
2023-04-12 22:59:08 +02:00
|
|
|
use std::collections::{HashMap, 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-12 22:59:08 +02:00
|
|
|
limits_rate: HashMap<LimitType, Limit>,
|
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-08 14:51:36 +02:00
|
|
|
pub async fn new(api_url: String) -> 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-12 16:49:18 +02:00
|
|
|
limits_rate: Limits::check_limits(api_url).await,
|
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,
|
|
|
|
) -> Option<Response> {
|
2023-04-14 22:22:23 +02:00
|
|
|
if self.can_send_request(limit_type) {
|
|
|
|
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-14 22:40:13 +02:00
|
|
|
self.update_limits(&response, limit_type);
|
|
|
|
return Some(response);
|
2023-04-14 22:22:23 +02:00
|
|
|
} else {
|
|
|
|
self.requests.push_back(TypedRequest {
|
|
|
|
request: request,
|
|
|
|
limit_type: limit_type,
|
|
|
|
});
|
2023-04-14 22:40:13 +02:00
|
|
|
return None;
|
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-14 22:40:13 +02:00
|
|
|
fn can_send_request(&mut self, limit_type: LimitType) -> bool {
|
2023-04-14 18:29:20 +02:00
|
|
|
let limits = self.limits_rate.get(&limit_type);
|
|
|
|
|
|
|
|
match limits {
|
|
|
|
Some(limit) => {
|
|
|
|
if limit.remaining > 0 {
|
|
|
|
true
|
|
|
|
} else {
|
|
|
|
false
|
|
|
|
}
|
|
|
|
}
|
|
|
|
None => false,
|
|
|
|
}
|
2023-04-14 17:09:50 +02:00
|
|
|
}
|
|
|
|
|
2023-04-14 22:40:13 +02:00
|
|
|
fn update_limits(&mut self, response: &Response, limit_type: LimitType) {
|
2023-04-12 22:59:08 +02:00
|
|
|
// TODO: Make this work
|
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-15 14:55:50 +02:00
|
|
|
None => self.limits_rate.get(&limit_type).unwrap().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-15 14:55:50 +02:00
|
|
|
None => self.limits_rate.get(&limit_type).unwrap().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-15 14:55:50 +02:00
|
|
|
None => self.limits_rate.get(&limit_type).unwrap().reset,
|
2023-04-12 18:33:16 +02:00
|
|
|
};
|
2023-04-12 22:59:08 +02:00
|
|
|
|
|
|
|
let mut limits_copy = self.limits_rate.clone();
|
|
|
|
let status = response.status();
|
|
|
|
let status_str = status.as_str();
|
|
|
|
|
|
|
|
if status_str.chars().next().unwrap() == '4' {
|
|
|
|
limits_copy.get_mut(&LimitType::Error).unwrap().remaining -= 1;
|
|
|
|
}
|
|
|
|
|
|
|
|
limits_copy.get_mut(&LimitType::Global).unwrap().remaining -= 1;
|
|
|
|
limits_copy.get_mut(&LimitType::Ip).unwrap().remaining -= 1;
|
|
|
|
|
|
|
|
match limit_type {
|
|
|
|
// Error, Global and Ip get handled seperately.
|
|
|
|
LimitType::Error => {}
|
|
|
|
LimitType::Global => {}
|
|
|
|
LimitType::Ip => {}
|
|
|
|
LimitType::AuthLogin => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AuthLogin).unwrap();
|
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 => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AbsoluteRegister).unwrap();
|
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.
|
|
|
|
limits_copy
|
|
|
|
.get_mut(&LimitType::AuthRegister)
|
|
|
|
.unwrap()
|
|
|
|
.remaining -= 1;
|
|
|
|
}
|
|
|
|
LimitType::AuthRegister => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AuthRegister).unwrap();
|
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.
|
|
|
|
limits_copy
|
|
|
|
.get_mut(&LimitType::AbsoluteRegister)
|
|
|
|
.unwrap()
|
|
|
|
.remaining -= 1;
|
|
|
|
}
|
|
|
|
LimitType::AbsoluteMessage => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AbsoluteMessage).unwrap();
|
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 => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::Channel).unwrap();
|
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 => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::Guild).unwrap();
|
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 => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::Webhook).unwrap();
|
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-14 23:38:36 +02:00
|
|
|
use super::*;
|
2023-04-15 13:27:43 +02:00
|
|
|
use crate::URLBundle;
|
2023-04-14 23:38:36 +02:00
|
|
|
#[tokio::test]
|
2023-04-15 13:27:43 +02:00
|
|
|
async fn create_limited_requester() {
|
2023-04-14 23:38:36 +02:00
|
|
|
let urls = URLBundle::new(
|
|
|
|
String::from("http://localhost:3001/api/"),
|
|
|
|
String::from("wss://localhost:3001/"),
|
|
|
|
String::from("http://localhost:3001/cdn"),
|
|
|
|
);
|
|
|
|
let requester = LimitedRequester::new(urls.api).await;
|
2023-04-15 13:54:33 +02:00
|
|
|
assert_eq!(
|
|
|
|
requester.limits_rate.get(&LimitType::Ip).unwrap(),
|
|
|
|
&Limit {
|
|
|
|
bucket: LimitType::Ip,
|
|
|
|
limit: 500,
|
|
|
|
remaining: 500,
|
|
|
|
reset: 5
|
|
|
|
}
|
|
|
|
);
|
2023-04-14 23:38:36 +02:00
|
|
|
}
|
|
|
|
}
|