74 lines
1.8 KiB
Rust
74 lines
1.8 KiB
Rust
use crate::error::{AppError, ErrorCode};
|
|
use crate::state::AppState;
|
|
|
|
use chrono::{Duration, Utc};
|
|
use std::net::IpAddr;
|
|
|
|
pub async fn consume_anonymous_units(
|
|
state: &AppState,
|
|
session_id: &str,
|
|
ip: IpAddr,
|
|
units: u32,
|
|
) -> Result<(), AppError> {
|
|
if units == 0 {
|
|
return Ok(());
|
|
}
|
|
|
|
let date = utc8_date();
|
|
let session_key = format!("anon_quota:{session_id}:{date}");
|
|
let ip_key = format!("anon_quota_ip:{ip}:{date}");
|
|
|
|
let mut conn = state.redis.clone();
|
|
|
|
let limit = state.config.anon_daily_units as i64;
|
|
let ttl_seconds = 48 * 60 * 60;
|
|
let inc = units as i64;
|
|
|
|
let script = redis::Script::new(
|
|
r#"
|
|
local limit = tonumber(ARGV[1])
|
|
local ttl = tonumber(ARGV[2])
|
|
local inc = tonumber(ARGV[3])
|
|
|
|
local v1 = tonumber(redis.call('GET', KEYS[1]) or '0')
|
|
local v2 = tonumber(redis.call('GET', KEYS[2]) or '0')
|
|
|
|
if v1 + inc > limit or v2 + inc > limit then
|
|
return -1
|
|
end
|
|
|
|
v1 = redis.call('INCRBY', KEYS[1], inc)
|
|
v2 = redis.call('INCRBY', KEYS[2], inc)
|
|
|
|
if v1 == inc then redis.call('EXPIRE', KEYS[1], ttl) end
|
|
if v2 == inc then redis.call('EXPIRE', KEYS[2], ttl) end
|
|
|
|
return v1
|
|
"#,
|
|
);
|
|
|
|
let new_value: i64 = script
|
|
.key(session_key)
|
|
.key(ip_key)
|
|
.arg(limit)
|
|
.arg(ttl_seconds)
|
|
.arg(inc)
|
|
.invoke_async(&mut conn)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "匿名配额检查失败").with_source(err))?;
|
|
|
|
if new_value < 0 {
|
|
return Err(AppError::new(
|
|
ErrorCode::QuotaExceeded,
|
|
"匿名试用次数已用完(每日 10 次)",
|
|
));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn utc8_date() -> String {
|
|
let now = Utc::now() + Duration::hours(8);
|
|
now.format("%Y-%m-%d").to_string()
|
|
}
|