tina-core 0.0.2

//! 限流数据
use crate::tina::data::i18n_string::I18nString;
use crate::tina::data::map::{ExpandToListMap, InsertOrAddToMap};
use crate::tina::data::AppResult;
use crate::tina::server::session::Session;
use crate::tina::util::string::StringExt;
use indexmap::IndexMap;
use serde_json::Value;
use std::time::Duration;

use super::date_time::LocalDateTime;

/// 限流配置
#[derive(Debug, Clone)]
pub struct ThrottleConfig {
    /// 限流类别
    pub(crate) category: &'static str,
    /// 限流策略
    pub(crate) policy: ThrottlePolicy,
    /// 限流触发策略
    pub(crate) fire_policy: ThrottleFirePolicy,
    /// 触发限流和阈值(次数)
    pub(crate) threshold: u32,
    /// 阈值的有效间隔
    pub(crate) duration: Duration,
    /// 相关参数
    pub(crate) param_name: Option<Vec<&'static str>>,
    /// 自定义提示消息
    pub(crate) message: Option<I18nString>,
}

/// 限流策略
#[derive(Serialize, Deserialize, Debug, Copy, Clone)]
#[serde(crate = "crate::serde")]
pub enum ThrottlePolicy {
    /// IP地址
    Ip,
    /// 用户
    User,
    /// 请求参数
    RequestParam,
    /// IP+用户
    IpUser,
    /// IP+请求参数
    IpRequestParam,
    /// 用户+请求参数
    UserRequestParam,
    /// IP+用户+请求参数
    IpUserRequestParam,
}

impl Default for ThrottlePolicy {
    fn default() -> Self {
        Self::Ip
    }
}

/// 限流触发策略
#[derive(Serialize, Deserialize, Copy, Debug, Clone)]
#[serde(crate = "crate::serde")]
pub enum ThrottleFirePolicy {
    /// 总是触发
    Always,
    /// 无异常时触发
    WithoutError,
    /// 发生异常时触发
    WithError,
}

impl Default for ThrottleFirePolicy {
    fn default() -> Self {
        Self::Always
    }
}

/// 限流参数
#[derive(Debug, Clone, Default, Serialize)]
#[serde(crate = "crate::serde")]
pub struct ThrottleParam {
    /// 限流的唯一key
    pub key: ThrottleUniqueKey,
    /// 阈值(次数)
    pub threshold: u32,
    /// 触发策略
    pub fire_policy: ThrottleFirePolicy,
    /// 参数信息
    pub param_map: IndexMap<String, String>,
    /// 自定义提示消息
    pub message: Option<I18nString>,
}

/// 新建限流参数
#[derive(Debug, Default, Serialize)]
#[serde(crate = "crate::serde")]
pub struct ThrottleNewParam {
    /// 当前日期时间
    pub current_date_time: LocalDateTime,
    /// 已存在的Throttle
    pub exist: Option<Throttle>,
    /// key
    pub key: String,
    /// 阈值
    pub threshold: u32,
    /// 间隔
    pub duration_millis: u64,
    /// 触发标记
    pub fire_flag: bool,
    /// 参数
    pub param_map: IndexMap<String, String>,
    /// 消息
    pub message: Option<I18nString>,
}

/// 限流的唯一Key
#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialOrd, PartialEq, Ord, Eq, Hash)]
#[serde(crate = "crate::serde")]
pub struct ThrottleUniqueKey {
    /// 限流的key
    pub key: String,
    /// 阈值有效间隔(毫秒)
    pub duration_millis: u64,
}

/// 限流数据
#[derive(Serialize, Deserialize, Debug, Default)]
#[serde(crate = "crate::serde")]
pub struct Throttle {
    /// 限流的唯一key
    pub key: ThrottleUniqueKey,
    /// 已触发次数
    pub fires: Option<u32>,
    /// 阈值(次数)
    pub threshold: Option<u32>,
    /// 创建时间
    pub create_time: Option<LocalDateTime>,
    /// 更新时间
    pub update_time: Option<LocalDateTime>,
    /// 过期时间
    pub expire_at: Option<LocalDateTime>,
    /// 参数信息
    #[serde(default)]
    pub param_map: IndexMap<String, String>,
    /// 自定义提示消息
    pub message: Option<I18nString>,
}

impl ThrottleConfig {
    fn new(
        category: &'static str,
        policy: ThrottlePolicy,
        fire_policy: ThrottleFirePolicy,
        threshold: u32,
        duration: Duration,
        param_name: Option<Vec<&'static str>>,
    ) -> Self {
        Self {
            category,
            policy,
            fire_policy,
            threshold,
            duration,
            param_name,
            message: None,
        }
    }
    /// 根据Ip构建(总是触发)
    pub fn fire_by_ip_always(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::Ip, ThrottleFirePolicy::Always, threshold, duration, None)
    }
    /// 根据Ip构建(无异常时触发)
    pub fn fire_by_ip_without_error(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::Ip, ThrottleFirePolicy::WithoutError, threshold, duration, None)
    }
    /// 根据Ip构建(有异常时触发)
    pub fn fire_by_ip_with_error(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::Ip, ThrottleFirePolicy::WithError, threshold, duration, None)
    }
    /// 根据用户构建(总是触发)
    pub fn fire_by_user_always(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::User, ThrottleFirePolicy::Always, threshold, duration, None)
    }
    /// 根据用户构建(无异常时触发)
    pub fn fire_by_user_without_error(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::User, ThrottleFirePolicy::WithoutError, threshold, duration, None)
    }
    /// 根据用户构建(有异常时触发)
    pub fn fire_by_user_with_error(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::User, ThrottleFirePolicy::WithError, threshold, duration, None)
    }
    /// 根据请求参数构建(总是触发)
    pub fn fire_by_req_param_always(category: &'static str, params: Vec<&'static str>, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::RequestParam, ThrottleFirePolicy::Always, threshold, duration, Some(params))
    }
    /// 根据请求参数构建(无异常时触发)
    pub fn fire_by_req_param_without_error(category: &'static str, params: Vec<&'static str>, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::RequestParam, ThrottleFirePolicy::WithoutError, threshold, duration, Some(params))
    }
    /// 根据请求参数构建(有异常时触发)
    pub fn fire_by_req_param_with_error(category: &'static str, params: Vec<&'static str>, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::RequestParam, ThrottleFirePolicy::WithError, threshold, duration, Some(params))
    }
    /// 根据Ip+用户构建(总是触发)
    pub fn fire_by_ip_user_always(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::IpUser, ThrottleFirePolicy::Always, threshold, duration, None)
    }
    /// 根据Ip+用户构建(无异常时触发)
    pub fn fire_by_ip_user_without_error(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::IpUser, ThrottleFirePolicy::WithoutError, threshold, duration, None)
    }
    /// 根据Ip+用户构建(有异常时触发)
    pub fn fire_by_ip_user_with_error(category: &'static str, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::IpUser, ThrottleFirePolicy::WithError, threshold, duration, None)
    }
    /// 根据Ip+请求参数构建(总是触发)
    pub fn fire_by_ip_req_param_always(category: &'static str, params: Vec<&'static str>, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::IpRequestParam, ThrottleFirePolicy::Always, threshold, duration, Some(params))
    }
    /// 根据Ip+请求参数构建(无异常时触发)
    pub fn fire_by_ip_req_param_without_error(
        category: &'static str,
        params: Vec<&'static str>,
        threshold: u32,
        duration: Duration,
    ) -> Self {
        Self::new(category, ThrottlePolicy::IpRequestParam, ThrottleFirePolicy::WithoutError, threshold, duration, Some(params))
    }
    /// 根据Ip+请求参数构建(有异常时触发)
    pub fn fire_by_ip_req_param_with_error(category: &'static str, params: Vec<&'static str>, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::IpRequestParam, ThrottleFirePolicy::WithError, threshold, duration, Some(params))
    }
    /// 根据用户+请求参数构建(总是触发)
    pub fn fire_by_user_req_param_always(category: &'static str, params: Vec<&'static str>, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::UserRequestParam, ThrottleFirePolicy::Always, threshold, duration, Some(params))
    }
    /// 根据用户+请求参数构建(无异常时触发)
    pub fn fire_by_user_req_param_without_error(
        category: &'static str,
        params: Vec<&'static str>,
        threshold: u32,
        duration: Duration,
    ) -> Self {
        Self::new(category, ThrottlePolicy::UserRequestParam, ThrottleFirePolicy::WithoutError, threshold, duration, Some(params))
    }
    /// 根据用户+请求参数构建(有异常时触发)
    pub fn fire_by_user_req_param_with_error(
        category: &'static str,
        params: Vec<&'static str>,
        threshold: u32,
        duration: Duration,
    ) -> Self {
        Self::new(category, ThrottlePolicy::UserRequestParam, ThrottleFirePolicy::WithError, threshold, duration, Some(params))
    }
    /// 根据Ip+用户+请求参数构建(总是触发)
    pub fn fire_by_ip_user_req_param_always(category: &'static str, params: Vec<&'static str>, threshold: u32, duration: Duration) -> Self {
        Self::new(category, ThrottlePolicy::IpUserRequestParam, ThrottleFirePolicy::Always, threshold, duration, Some(params))
    }
    /// 根据Ip+用户+请求参数构建(无异常时触发)
    pub fn fire_by_ip_user_req_param_without_error(
        category: &'static str,
        params: Vec<&'static str>,
        threshold: u32,
        duration: Duration,
    ) -> Self {
        Self::new(category, ThrottlePolicy::IpUserRequestParam, ThrottleFirePolicy::WithoutError, threshold, duration, Some(params))
    }
    /// 根据Ip+用户+请求参数构建(有异常时触发)
    pub fn fire_by_ip_user_req_param_with_error(
        category: &'static str,
        params: Vec<&'static str>,
        threshold: u32,
        duration: Duration,
    ) -> Self {
        Self::new(category, ThrottlePolicy::IpUserRequestParam, ThrottleFirePolicy::WithError, threshold, duration, Some(params))
    }
}

impl ThrottleConfig {
    /// 转换成限流参数
    pub fn to_throttle_param(&self, remote_ip_address: &str, session: &Session, param_value: &Value) -> AppResult<Vec<ThrottleParam>> {
        let mut params = Vec::new();
        match self.policy {
            ThrottlePolicy::Ip | ThrottlePolicy::User | ThrottlePolicy::IpUser => {
                let mut param = ThrottleParam {
                    key: Default::default(),
                    threshold: self.threshold,
                    fire_policy: self.fire_policy,
                    param_map: Default::default(),
                    message: self.message.clone(),
                };
                setup_throttle_key(&mut param, self.policy, self.category, "", self.duration, remote_ip_address, session)?;
                params.push(param);
            }
            ThrottlePolicy::RequestParam
            | ThrottlePolicy::IpRequestParam
            | ThrottlePolicy::UserRequestParam
            | ThrottlePolicy::IpUserRequestParam => {
                if let Some(request_params) = self.param_name.as_ref() {
                    let mut expression_map = IndexMap::new();
                    add_to_expression_map(param_value, &mut expression_map, "");

                    if !expression_map.is_empty() {
                        let list_map = expression_map.to_list_map();
                        for expression_map in list_map.into_iter() {
                            let mut param_map = IndexMap::new();
                            let resolved_params = request_params
                                .iter()
                                .map(|&request_param| {
                                    let resolved_param =
                                        expression_map.get(request_param).map(|s| s.to_string()).unwrap_or_else(|| "".to_owned());
                                    param_map.insert(request_param.to_string(), resolved_param.to_string());
                                    resolved_param
                                })
                                .reduce(|mut pre, cur| {
                                    if !pre.is_empty() {
                                        pre.push_str("_##_");
                                    }
                                    pre.push_str(cur.as_str());
                                    pre
                                })
                                .unwrap_or_else(|| "".to_owned());
                            let mut param = ThrottleParam {
                                key: Default::default(),
                                threshold: self.threshold,
                                fire_policy: self.fire_policy,
                                param_map,
                                message: self.message.clone(),
                            };
                            setup_throttle_key(
                                &mut param,
                                self.policy,
                                self.category,
                                resolved_params.to_md5_lowercase().as_str(),
                                self.duration,
                                remote_ip_address,
                                session,
                            )?;
                            params.push(param);
                        }
                    }
                }
            }
        }
        Ok(params)
    }
    /// 设置自定义消息
    pub fn message(mut self, message: I18nString) -> Self {
        self.message = Some(message);
        self
    }
}

fn setup_throttle_key(
    param: &mut ThrottleParam,
    policy: ThrottlePolicy,
    category: &str,
    resolved_param_expression: &str,
    duration: Duration,
    remote_ip_address: &str,
    session: &Session,
) -> AppResult<()> {
    let prefix = match category.is_empty() {
        true => "".to_owned(),
        false => format!("{}:", category),
    };
    let key = match policy {
        ThrottlePolicy::Ip => format!("{}{}:{}", prefix, "ip", remote_ip_address),
        ThrottlePolicy::User => {
            format!("{}{}:{}", prefix, "user", session.get_user_id().ok_or(crate::app_system_error!("no user_id found in session"))?)
        }
        ThrottlePolicy::IpUser => {
            format!(
                "{}{}:{}:{}",
                prefix,
                "ip_user",
                remote_ip_address,
                session.get_user_id().ok_or(crate::app_system_error!("no user_id found in session"))?
            )
        }
        ThrottlePolicy::RequestParam => format!("{}{}:{}", prefix, "param", resolved_param_expression),
        ThrottlePolicy::IpRequestParam => {
            format!("{}{}:{}:{}", prefix, "ip_param", remote_ip_address, resolved_param_expression)
        }
        ThrottlePolicy::UserRequestParam => {
            format!(
                "{}{}:{}:{}",
                prefix,
                "user_param",
                session.get_user_id().ok_or(crate::app_system_error!("no user_id found in session"))?,
                resolved_param_expression
            )
        }
        ThrottlePolicy::IpUserRequestParam => {
            format!(
                "{}{}:{}:{}:{}",
                prefix,
                "ip_user_param",
                remote_ip_address,
                session.get_user_id().ok_or(crate::app_system_error!("no user_id found in session"))?,
                resolved_param_expression
            )
        }
    };
    param.key = ThrottleUniqueKey {
        key,
        duration_millis: duration.as_millis() as u64,
    };
    Ok(())
}

fn add_to_expression_map(param_value: &Value, expression_map: &mut IndexMap<String, Vec<String>>, prefix: &str) {
    let key_prefix = match prefix.is_empty() {
        true => "".to_owned(),
        false => format!("{}.", prefix),
    };
    match param_value {
        Value::Array(arr) => {
            for item in arr.iter() {
                add_to_expression_map(item, expression_map, key_prefix.as_str());
            }
        }
        Value::Object(map) => {
            for (key, value) in map.iter() {
                if key.is_empty() {
                    continue;
                }
                let key = match key_prefix.is_empty() {
                    true => key.to_string(),
                    false => format!("{}.{}", key_prefix, key),
                };
                match value {
                    Value::Null => {
                        expression_map.insert_or_add_to_value(key, Value::Null.to_string());
                    }
                    Value::Bool(v) => {
                        expression_map.insert_or_add_to_value(key, v.to_string());
                    }
                    Value::Number(v) => {
                        expression_map.insert_or_add_to_value(key, v.to_string());
                    }
                    Value::String(v) => {
                        expression_map.insert_or_add_to_value(key, v.to_string());
                    }
                    Value::Array(arr) => {
                        for item in arr.iter() {
                            add_to_expression_map(item, expression_map, key.as_str());
                        }
                    }
                    Value::Object(v) => {
                        add_to_expression_map(&Value::from(v.clone()), expression_map, key.as_str());
                    }
                }
            }
        }
        _ => {}
    }
}

#[cfg(feature = "redis")]
mod redis {
    use crate::tina::data::throttle::Throttle;
    use crate::tina::util::json::JsonUtil;
    use redis::{FromRedisValue, RedisError, RedisResult, RedisWrite, ToRedisArgs, Value};

    impl ToRedisArgs for &Throttle {
        fn write_redis_args<W>(&self, out: &mut W)
        where
            W: ?Sized + RedisWrite,
        {
            let str = JsonUtil::to_json_string(&self);
            out.write_arg(str.as_bytes())
        }
    }

    impl FromRedisValue for Throttle {
        fn from_redis_value(v: &Value) -> RedisResult<Self> {
            if let Value::Data(data) = v {
                let str =
                    String::from_utf8(data.clone()).map_err(|_| RedisError::from((redis::ErrorKind::TypeError, "invalid value format")))?;
                let cache: Throttle = JsonUtil::parse_json_string(str.as_str())
                    .map_err(|_| RedisError::from((redis::ErrorKind::TypeError, "invalid value format")))?;
                return Ok(cache);
            }
            Err(RedisError::from((redis::ErrorKind::TypeError, "invalid value format")))
        }
    }
}