use crate::api::limits::Config; use reqwest::{Client, Request, RequestBuilder}; use serde_json::from_str; 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, requests: VecDeque, last_reset_epoch: i64, } 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. 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(), } } fn get_limit(&self, bucket: &str) -> Option<&Limit> { for limit in self.limits.iter() { if limit.bucket == bucket { return Some(limit.to_owned()); } } None } /// `last_match` returns the last match of a `Vec<&str>` in a `String`. It returns a `Option<(usize, String)>` /// that contains the index of the match and the match itself. /// If no match is found, it returns `None`. /// # Example /// ```rs /// let string = "https://discord.com/api/v8/channels/1234567890/messages"; /// let matches = ["channels", "messages"]; /// let index_match = last_match(string, matches.to_vec()); /// assert_eq!(index_match, Some((43, String::from("messages")))); /// ``` fn last_match(string: &str, search_for: Vec<&str>) -> Option<(usize, String)> { let mut index_match: (usize, &str) = (0, ""); for _match in search_for.iter() { if !string.contains(_match) { continue; } // Get the index of the match let temp_index_match = string.match_indices(_match).next().unwrap(); // As only the last match is relevant, we only update the index_match if the index is // higher than the previous one. if temp_index_match.0 > index_match.0 { index_match = temp_index_match; } } if index_match.0 == 0 { return None; } else { return Some((index_match.0, String::from(index_match.1))); } } /// `can_send_request` checks if a request can be sent. It returns a `bool` that indicates if /// the request can be sent. fn can_send_request(&mut self, request: RequestBuilder) -> bool { // get the url from request let global_limit = self.get_limit("global").unwrap(); let ip_limit = self.get_limit("ip").unwrap(); if ip_limit.remaining == 0 || global_limit.remaining == 0 { return false; } let url_path = request .try_clone() .unwrap() .build() .unwrap() .url() .path() .to_string(); // Define the different rate limit buckets as they would appear in the URL let matches = [ "login", "register", "webhooks", "channels", "messages", "guilds", ]; let index_match_string = LimitedRequester::last_match(&url_path, matches.to_vec()) .unwrap_or_else(|| (0, String::from(""))); if index_match_string.0 == 0 { return true; } let index_match = (index_match_string.0, index_match_string.1.as_str()); match index_match.1 { "login" => { let auth_limit = self.get_limit("login").unwrap(); if auth_limit.remaining != 0 { return true; } return false; } "register" => { let auth_limit = self.get_limit("auth.register").unwrap(); let absolute_limit = self.get_limit("absoluteRate.register").unwrap(); if auth_limit.remaining != 0 && absolute_limit.remaining != 0 { return true; } return false; } "messages" => { let absolute_limit = self.get_limit("absoluteRate.sendMessages").unwrap(); let request_method = request .try_clone() .unwrap() .build() .unwrap() .method() .as_str() .to_owned(); if absolute_limit.remaining != 0 || request_method != "POST" || request_method != "PUT" || request_method != "PATCH" { return true; } return false; } "webhooks" => { let auth_limit = self.get_limit("webhooks").unwrap(); if auth_limit.remaining != 0 { return true; } return false; } "channels" => { let auth_limit = self.get_limit("channels").unwrap(); if auth_limit.remaining != 0 { return true; } return false; } "guilds" => { let auth_limit = self.get_limit("guilds").unwrap(); if auth_limit.remaining != 0 { return true; } return false; } &_ => { panic!(); } } } /// `update_limits` updates the `Limit`s of the `LimitedRequester` based on the response headers /// of the request that was just sent. fn update_limits(&mut self, response: &reqwest::Response) { let headers = response.headers(); let mut reset_epoch = 0; for limit in self.limits.iter_mut() { if let Some(value) = headers.get(limit.bucket.clone()) { limit.remaining = value.to_str().unwrap().parse().unwrap(); } if let Some(value) = headers.get("X-RateLimit-Reset".to_string()) { reset_epoch = value.to_str().unwrap().parse().unwrap(); } } if reset_epoch != 0 { self.last_reset_epoch = reset_epoch; } } /// `send_request` sends a request to the server, if `can_send_request()` is true. /// It will then update the `Limit`s by calling `update_limits()`. /// # Example pub async fn send_request(&mut self, request: RequestBuilder) -> reqwest::Response { if !self.can_send_request(request.try_clone().unwrap()) { panic!("429: Rate limited"); } let response = request.send().await.unwrap(); self.update_limits(&response); // TODO: This does not use the request queue at all. Implement! >:3 response } /// 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 { 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 } }