GeneIT/geneit_backend/src/services/rate_limiter_service.rs

use crate::connections::redis_connection;
use crate::utils::time_utils::time;
use std::net::IpAddr;
use std::time::Duration;

#[derive(Debug, Copy, Clone)]
pub enum RatedAction {
    CreateAccount,
    CheckResetPasswordTokenFailed,
    RequestNewPasswordResetLink,
    FailedPasswordLogin,
    StartOpenIDLogin,
    RequestReplacePasswordSignedIn,
    RequestDeleteAccount,
    JoinFamily,
}

impl RatedAction {
    fn id(&self) -> &'static str {
        match self {
            RatedAction::CreateAccount => "create-account",
            RatedAction::CheckResetPasswordTokenFailed => "check-reset-password-token",
            RatedAction::RequestNewPasswordResetLink => "req-pwd-reset-lnk",
            RatedAction::FailedPasswordLogin => "failed-login",
            RatedAction::StartOpenIDLogin => "start-oidc-login",
            RatedAction::RequestReplacePasswordSignedIn => "req-pwd-signed-in",
            RatedAction::RequestDeleteAccount => "req-del-acct",
            RatedAction::JoinFamily => "join-family",
        }
    }

    fn limit(&self) -> usize {
        match self {
            RatedAction::CreateAccount => 5,
            RatedAction::CheckResetPasswordTokenFailed => 100,
            RatedAction::RequestNewPasswordResetLink => 5,
            RatedAction::FailedPasswordLogin => 15,
            RatedAction::StartOpenIDLogin => 30,
            RatedAction::RequestReplacePasswordSignedIn => 5,
            RatedAction::RequestDeleteAccount => 5,
            RatedAction::JoinFamily => 10,
        }
    }

    fn keep_seconds(&self) -> u64 {
        3600
    }

    fn key(&self, ip: IpAddr) -> String {
        format!("rate-{}-{}", self.id(), ip)
    }
}

/// Keep track of the time the action was executed by the user
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
struct ActionRecord(Vec<u64>);

impl ActionRecord {
    pub fn clean(&mut self, action: RatedAction) {
        self.0.retain(|e| e + action.keep_seconds() > time());
    }
}

/// Record a new action of the user
pub async fn record_action(ip: IpAddr, action: RatedAction) -> anyhow::Result<()> {
    let key = action.key(ip);

    let mut record = redis_connection::get_value::<ActionRecord>(&key)
        .await?
        .unwrap_or_default();
    record.clean(action);
    record.0.push(time());
    redis_connection::set_value(&key, &record, Duration::from_secs(action.keep_seconds())).await?;

    Ok(())
}

/// Check whether an action should be blocked, due to too much attempts from the user
pub async fn should_block_action(ip: IpAddr, action: RatedAction) -> anyhow::Result<bool> {
    let mut record = redis_connection::get_value::<ActionRecord>(&action.key(ip))
        .await?
        .unwrap_or_default();

    record.clean(action);

    Ok(record.0.len() >= action.limit())
}