refactor limit.rs, moved some logic
This commit is contained in:
parent
2d42423a1f
commit
317c8e801e
|
@ -1,5 +1,7 @@
|
||||||
pub mod limits {
|
pub mod limits {
|
||||||
|
use reqwest::Client;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::from_str;
|
||||||
#[derive(Debug, Deserialize, Serialize)]
|
#[derive(Debug, Deserialize, Serialize)]
|
||||||
#[allow(non_snake_case)]
|
#[allow(non_snake_case)]
|
||||||
pub struct User {
|
pub struct User {
|
||||||
|
@ -97,4 +99,205 @@ pub mod limits {
|
||||||
pub rate: Rate,
|
pub rate: Rate,
|
||||||
pub absoluteRate: AbsoluteRate,
|
pub absoluteRate: AbsoluteRate,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub struct Limit {
|
||||||
|
pub bucket: String,
|
||||||
|
pub limit: u64,
|
||||||
|
pub remaining: u64,
|
||||||
|
pub reset: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for Limit {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
write!(
|
||||||
|
f,
|
||||||
|
"Bucket: {}, Limit: {}, Remaining: {}, Reset: {}",
|
||||||
|
self.bucket, self.limit, self.remaining, self.reset
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
pub struct Limits {
|
||||||
|
pub limit_absolute_messages: Limit,
|
||||||
|
pub limit_absolute_register: Limit,
|
||||||
|
pub limit_auth_login: Limit,
|
||||||
|
pub limit_auth_register: Limit,
|
||||||
|
pub limit_ip: Limit,
|
||||||
|
pub limit_global: Limit,
|
||||||
|
pub limit_error: Limit,
|
||||||
|
pub limit_guild: Limit,
|
||||||
|
pub limit_webhook: Limit,
|
||||||
|
pub limit_channel: Limit,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Limits {
|
||||||
|
/// check_limits uses the API to get the current request limits of the instance.
|
||||||
|
/// It returns a `Limits` struct containing all the limits.
|
||||||
|
pub async fn check_limits(api_url: String) -> Limits {
|
||||||
|
let client = Client::new();
|
||||||
|
let url_parsed = crate::URLBundle::parse_url(api_url) + "/policies/instance/limits";
|
||||||
|
let result = client
|
||||||
|
.get(url_parsed)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|e| panic!("An error occured while performing the request: {}", e))
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|e| {
|
||||||
|
panic!(
|
||||||
|
"An error occured while parsing the request body string: {}",
|
||||||
|
e
|
||||||
|
)
|
||||||
|
});
|
||||||
|
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
|
||||||
|
let mut limits: Limits;
|
||||||
|
if config.rate.enabled == false {
|
||||||
|
limits = Limits {
|
||||||
|
limit_absolute_messages: Limit {
|
||||||
|
bucket: "absolute_messages".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_absolute_register: Limit {
|
||||||
|
bucket: "absolute_register".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_auth_login: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_auth_register: Limit {
|
||||||
|
bucket: "auth_register".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_ip: Limit {
|
||||||
|
bucket: "ip".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_global: Limit {
|
||||||
|
bucket: "global".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_error: Limit {
|
||||||
|
bucket: "error".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_guild: Limit {
|
||||||
|
bucket: "guild".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_webhook: Limit {
|
||||||
|
bucket: "webhook".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
limit_channel: Limit {
|
||||||
|
bucket: "channel".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
} else {
|
||||||
|
limits = Limits {
|
||||||
|
limit_absolute_messages: Limit {
|
||||||
|
bucket: "absolute_messages".to_string(),
|
||||||
|
limit: config.absoluteRate.sendMessage.limit,
|
||||||
|
remaining: config.absoluteRate.sendMessage.limit,
|
||||||
|
reset: config.absoluteRate.sendMessage.window,
|
||||||
|
},
|
||||||
|
limit_absolute_register: Limit {
|
||||||
|
bucket: "absolute_register".to_string(),
|
||||||
|
limit: config.absoluteRate.register.limit,
|
||||||
|
remaining: config.absoluteRate.register.limit,
|
||||||
|
reset: config.absoluteRate.register.window,
|
||||||
|
},
|
||||||
|
limit_auth_login: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.routes.auth.login.count,
|
||||||
|
remaining: config.rate.routes.auth.login.count,
|
||||||
|
reset: config.rate.routes.auth.login.window,
|
||||||
|
},
|
||||||
|
limit_auth_register: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.routes.auth.register.count,
|
||||||
|
remaining: config.rate.routes.auth.register.count,
|
||||||
|
reset: config.rate.routes.auth.register.window,
|
||||||
|
},
|
||||||
|
limit_guild: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.routes.guild.count,
|
||||||
|
remaining: config.rate.routes.guild.count,
|
||||||
|
reset: config.rate.routes.guild.window,
|
||||||
|
},
|
||||||
|
limit_webhook: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.routes.webhook.count,
|
||||||
|
remaining: config.rate.routes.webhook.count,
|
||||||
|
reset: config.rate.routes.webhook.window,
|
||||||
|
},
|
||||||
|
limit_channel: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.routes.channel.count,
|
||||||
|
remaining: config.rate.routes.channel.count,
|
||||||
|
reset: config.rate.routes.channel.window,
|
||||||
|
},
|
||||||
|
limit_ip: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.ip.count,
|
||||||
|
remaining: config.rate.ip.count,
|
||||||
|
reset: config.rate.ip.window,
|
||||||
|
},
|
||||||
|
limit_global: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.global.count,
|
||||||
|
remaining: config.rate.global.count,
|
||||||
|
reset: config.rate.global.window,
|
||||||
|
},
|
||||||
|
limit_error: Limit {
|
||||||
|
bucket: "auth_login".to_string(),
|
||||||
|
limit: config.rate.error.count,
|
||||||
|
remaining: config.rate.error.count,
|
||||||
|
reset: config.rate.error.window,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !config.absoluteRate.register.enabled {
|
||||||
|
limits.limit_absolute_register = Limit {
|
||||||
|
bucket: "absolute_messages".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if !config.absoluteRate.sendMessage.enabled {
|
||||||
|
limits.limit_absolute_messages = Limit {
|
||||||
|
bucket: "absolute_register".to_string(),
|
||||||
|
limit: u64::MAX,
|
||||||
|
remaining: u64::MAX,
|
||||||
|
reset: u64::MAX,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return limits;
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
137
src/limit.rs
137
src/limit.rs
|
@ -1,4 +1,4 @@
|
||||||
use crate::api::limits::Config;
|
use crate::api::limits::{Config, Limit, Limits};
|
||||||
|
|
||||||
use reqwest::{Client, Request, RequestBuilder};
|
use reqwest::{Client, Request, RequestBuilder};
|
||||||
use serde_json::from_str;
|
use serde_json::from_str;
|
||||||
|
@ -7,28 +7,11 @@ use std::collections::VecDeque;
|
||||||
// Note: There seem to be some overlapping request limiters. We need to make sure that sending a
|
// 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
|
// request checks for all the request limiters that apply, and blocks if any of the limiters are 0
|
||||||
|
|
||||||
pub struct Limit {
|
|
||||||
limit: u64,
|
|
||||||
remaining: u64,
|
|
||||||
reset: u64,
|
|
||||||
bucket: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::fmt::Display for Limit {
|
|
||||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
||||||
write!(
|
|
||||||
f,
|
|
||||||
"Bucket: {}, Limit: {}, Remaining: {}, Reset: {}",
|
|
||||||
self.bucket, self.limit, self.remaining, self.reset
|
|
||||||
)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct LimitedRequester {
|
pub struct LimitedRequester {
|
||||||
http: Client,
|
http: Client,
|
||||||
limits: Vec<Limit>, // TODO: Replace with all Limits
|
|
||||||
requests: VecDeque<Request>,
|
requests: VecDeque<Request>,
|
||||||
last_reset_epoch: i64,
|
last_reset_epoch: i64,
|
||||||
|
limits_rate: Limits,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl LimitedRequester {
|
impl LimitedRequester {
|
||||||
|
@ -38,124 +21,10 @@ impl LimitedRequester {
|
||||||
/// headers to see if it can find Ratelimit info to update itself.
|
/// headers to see if it can find Ratelimit info to update itself.
|
||||||
pub async fn new(api_url: String) -> Self {
|
pub async fn new(api_url: String) -> Self {
|
||||||
LimitedRequester {
|
LimitedRequester {
|
||||||
limits: LimitedRequester::check_limits(api_url).await,
|
|
||||||
http: Client::new(),
|
http: Client::new(),
|
||||||
requests: VecDeque::new(),
|
requests: VecDeque::new(),
|
||||||
last_reset_epoch: chrono::Utc::now().timestamp(),
|
last_reset_epoch: chrono::Utc::now().timestamp(),
|
||||||
|
limits_rate: Limits::check_limits(api_url).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// check_limits uses the API to get the current request limits of the instance.
|
|
||||||
/// It returns a `Vec` of `Limit`s, which can be used to check if a request can be sent.
|
|
||||||
pub async fn check_limits(api_url: String) -> Vec<Limit> {
|
|
||||||
let client = Client::new();
|
|
||||||
let url_parsed = crate::URLBundle::parse_url(api_url) + "/policies/instance/limits";
|
|
||||||
let result = client
|
|
||||||
.get(url_parsed)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.unwrap_or_else(|e| panic!("An error occured while performing the request: {}", e))
|
|
||||||
.text()
|
|
||||||
.await
|
|
||||||
.unwrap_or_else(|e| {
|
|
||||||
panic!(
|
|
||||||
"An error occured while parsing the request body string: {}",
|
|
||||||
e
|
|
||||||
)
|
|
||||||
});
|
|
||||||
let config: Config = from_str(&result).unwrap();
|
|
||||||
|
|
||||||
let mut limit_vector = Vec::new();
|
|
||||||
if !config.rate.enabled {
|
|
||||||
// The different rate limit buckets, except for the absoluteRate ones. These will be
|
|
||||||
// handled seperately.
|
|
||||||
let types = [
|
|
||||||
"ip",
|
|
||||||
"auth.login",
|
|
||||||
"auth.register",
|
|
||||||
"global",
|
|
||||||
"error",
|
|
||||||
"guild",
|
|
||||||
"webhook",
|
|
||||||
"channel",
|
|
||||||
];
|
|
||||||
for type_ in types.iter() {
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: u64::MAX,
|
|
||||||
remaining: u64::MAX,
|
|
||||||
reset: 1,
|
|
||||||
bucket: type_.to_string(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.ip.count,
|
|
||||||
remaining: config.rate.ip.count,
|
|
||||||
reset: config.rate.ip.window,
|
|
||||||
bucket: String::from("ip"),
|
|
||||||
});
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.global.count,
|
|
||||||
remaining: config.rate.global.count,
|
|
||||||
reset: config.rate.global.window,
|
|
||||||
bucket: String::from("global"),
|
|
||||||
});
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.error.count,
|
|
||||||
remaining: config.rate.error.count,
|
|
||||||
reset: config.rate.error.window,
|
|
||||||
bucket: String::from("error"),
|
|
||||||
});
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.routes.guild.count,
|
|
||||||
remaining: config.rate.routes.guild.count,
|
|
||||||
reset: config.rate.routes.guild.window,
|
|
||||||
bucket: String::from("guild"),
|
|
||||||
});
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.routes.webhook.count,
|
|
||||||
remaining: config.rate.routes.webhook.count,
|
|
||||||
reset: config.rate.routes.webhook.window,
|
|
||||||
bucket: String::from("webhook"),
|
|
||||||
});
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.routes.channel.count,
|
|
||||||
remaining: config.rate.routes.channel.count,
|
|
||||||
reset: config.rate.routes.channel.window,
|
|
||||||
bucket: String::from("channel"),
|
|
||||||
});
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.routes.auth.login.count,
|
|
||||||
remaining: config.rate.routes.auth.login.count,
|
|
||||||
reset: config.rate.routes.auth.login.window,
|
|
||||||
bucket: String::from("auth.login"),
|
|
||||||
});
|
|
||||||
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.rate.routes.auth.register.count,
|
|
||||||
remaining: config.rate.routes.auth.register.count,
|
|
||||||
reset: config.rate.routes.auth.register.window,
|
|
||||||
bucket: String::from("auth.register"),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
if config.absoluteRate.register.enabled {
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.absoluteRate.register.limit,
|
|
||||||
remaining: config.absoluteRate.register.limit,
|
|
||||||
reset: config.absoluteRate.register.window,
|
|
||||||
bucket: String::from("absoluteRate.register"),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
if config.absoluteRate.sendMessage.enabled {
|
|
||||||
limit_vector.push(Limit {
|
|
||||||
limit: config.absoluteRate.sendMessage.limit,
|
|
||||||
remaining: config.absoluteRate.sendMessage.limit,
|
|
||||||
reset: config.absoluteRate.sendMessage.window,
|
|
||||||
bucket: String::from("absoluteRate.messages"),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
limit_vector
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue