//! # Requests limit helper //! //! Handle the limitation of requests, depending on threshold criterias use std::sync::{Arc, Mutex}; use crate::constants::LIMIT_COUNTER_LIFETIME; use crate::controllers::routes::{LimitPolicy, Route}; use crate::data::error::{ExecError, ResultBoxError}; use crate::data::http_request_handler::HttpRequestHandler; use crate::utils::date_utils; use crate::utils::date_utils::time; /// Information about a IP address limitation struct IpInfo { ip: String, time_start: u64, uri: String, count: u64, } /// Limits cache type Cache = Vec; static mut LIMITS_CACHE: Option>> = None; /// Initialize limit cache storage pub fn init() { let limits_cache = Vec::new(); let limits_cache = Some(Arc::new(Mutex::new(limits_cache))); unsafe { LIMITS_CACHE = limits_cache; } } /// Get access to the cache. This resource must absolutely be released as quickly as possible fn get_cache() -> ResultBoxError>> { let cache; unsafe { cache = LIMITS_CACHE.as_ref().unwrap().clone(); } Ok(cache) } /// Clean cached information pub fn clean_cache() -> ResultBoxError { let time = date_utils::time(); let mut i = 0; let cache = get_cache()?; let mut cache = cache.lock().unwrap(); while i < cache.len() { if cache[i].time_start + LIMIT_COUNTER_LIFETIME < time { cache.remove(i); } else { i = i + 1; } } Ok(()) } /// Trigger limit helper at the beginning of requests pub fn trigger_before(req: &HttpRequestHandler, route: &Route) -> ResultBoxError { if route.limit_policy.is_none() { return Ok(()); } let max_count = route.limit_policy.get_count(); let ip = req.remote_ip(); let cache = get_cache()?; let found = cache.lock().unwrap() .iter() .find( |k| k.uri.eq(route.uri) && k.count >= max_count && k.ip.eq(&ip)) .is_some(); if found { return Err(ExecError::boxed_new("Limit exceeded!")); } Ok(()) } /// Trigger limit at the end of the request pub fn trigger_after(is_success: bool, req: &HttpRequestHandler, route: &Route) -> ResultBoxError { let need_trigger = match (&route.limit_policy, is_success) { (LimitPolicy::NONE, _) => false, (LimitPolicy::ANY(_), _) => true, (LimitPolicy::SUCCESS(_), res) => res, (LimitPolicy::FAILURE(_), res) => !res, }; if !need_trigger { return Ok(()); } let ip = req.remote_ip(); let cache = get_cache()?; let mut cache = cache.lock().unwrap(); // We search for existing entry for i in 0..cache.len() { if cache[i].ip.eq(&ip) && cache[i].uri.eq(route.uri) { cache[i].count += 1; return Ok(()); } } // Otherwise we must add the entry to the table cache.push(IpInfo { ip, time_start: time(), uri: route.uri.to_string(), count: 1, }); Ok(()) }