Make limit system lives

src/helpers/requests_limit_helper.rs

@@ -2,26 +2,32 @@
 //!
 //! Handle the limitation of requests, depending on threshold criterias
 
+
 use std::sync::{Arc, Mutex};
 
 use crate::constants::LIMIT_COUNTER_LIFETIME;
-use crate::data::error::ResultBoxError;
+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 = Arc<dashmap::DashMap<String, IpInfo>>;
+type Cache = Vec<IpInfo>;
 
 static mut LIMITS_CACHE: Option<Arc<Mutex<Cache>>> = None;
 
 /// Initialize limit cache storage
 pub fn init() {
-    let limits_cache = Arc::new(dashmap::DashMap::new());
+    let limits_cache = Vec::new();
     let limits_cache = Some(Arc::new(Mutex::new(limits_cache)));
 
     unsafe {
@@ -30,12 +36,11 @@ pub fn init() {
 }
 
 /// Get access to the cache. This resource must absolutely be released as quickly as possible
-fn get_cache() -> ResultBoxError<Cache> {
-    let cache: Cache;
+fn get_cache() -> ResultBoxError<Arc<Mutex<Cache>>> {
+    let cache;
 
     unsafe {
-        let guard = LIMITS_CACHE.as_ref().unwrap().lock();
-        cache = guard.unwrap().clone();
+        cache = LIMITS_CACHE.as_ref().unwrap().clone();
     }
 
     Ok(cache)
@@ -45,16 +50,82 @@ fn get_cache() -> ResultBoxError<Cache> {
 pub fn clean_cache() -> ResultBoxError {
     let time = date_utils::time();
 
-    let cache = get_cache()?;
-    let obsolete_ips: Vec<String> = cache
-        .iter()
-        .filter(|k| k.time_start + LIMIT_COUNTER_LIFETIME < time)
-        .map(|k| k.key().to_string())
-        .collect();
+    let mut i = 0;
 
-    for obsolete_ip in obsolete_ips {
-        cache.remove(&obsolete_ip);
+    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(())
 }