refactor limit.rs, moved some logic
This commit is contained in:
bitfl0wer
parent
2d42423a1f
commit
317c8e801e
|
|
@ -1,5 +1,7 @@
|
|||
pub mod limits {
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::from_str;
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[allow(non_snake_case)]
|
||||
pub struct User {
|
||||
|
|
@ -97,4 +99,205 @@ pub mod limits {
|
|||
pub rate: Rate,
|
||||
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 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
|
||||
// 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 {
|
||||
http: Client,
|
||||
limits: Vec<Limit>, // TODO: Replace with all Limits
|
||||
requests: VecDeque<Request>,
|
||||
last_reset_epoch: i64,
|
||||
limits_rate: Limits,
|
||||
}
|
||||
|
||||
impl LimitedRequester {
|
||||
|
|
@ -38,124 +21,10 @@ impl LimitedRequester {
|
|||
/// headers to see if it can find Ratelimit info to update itself.
|
||||
pub async fn new(api_url: String) -> Self {
|
||||
LimitedRequester {
|
||||
limits: LimitedRequester::check_limits(api_url).await,
|
||||
http: Client::new(),
|
||||
requests: VecDeque::new(),
|
||||
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