#[cfg(feature = "parallel-checker")]
pub mod parallel_checker;
#[cfg(feature = "parallel-checker")]
pub use parallel_checker::ParallelBanChecker;
use crate::error::{ConsumeResult, StorageError};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
#[async_trait]
pub trait Storage: Send + Sync {
async fn get(&self, key: &str) -> Result<Option<String>, StorageError>;
async fn set(&self, key: &str, value: &str, ttl: Option<u64>) -> Result<(), StorageError>;
async fn delete(&self, key: &str) -> Result<(), StorageError>;
}
#[derive(Debug, Clone)]
pub struct QuotaInfo {
pub consumed: u64,
pub limit: u64,
pub window_start: DateTime<Utc>,
pub window_end: DateTime<Utc>,
}
#[async_trait]
pub trait QuotaStorage: Send + Sync {
async fn get_quota(
&self,
user_id: &str,
resource: &str,
) -> Result<Option<QuotaInfo>, StorageError>;
async fn consume(
&self,
user_id: &str,
resource: &str,
cost: u64,
limit: u64,
window: Duration,
) -> Result<ConsumeResult, StorageError>;
async fn reset(
&self,
user_id: &str,
resource: &str,
limit: u64,
window: Duration,
) -> Result<(), StorageError>;
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(tag = "type", content = "value")]
pub enum BanTarget {
#[serde(rename = "ip")]
Ip(String),
#[serde(rename = "user")]
UserId(String),
#[serde(rename = "mac")]
Mac(String),
#[serde(rename = "geo")]
Geo { country_code: String },
}
#[derive(Debug, Clone)]
pub struct BanRecord {
pub target: BanTarget,
pub ban_times: u32,
pub duration: Duration,
pub banned_at: DateTime<Utc>,
pub expires_at: DateTime<Utc>,
pub is_manual: bool,
pub reason: String,
}
#[derive(Debug, Clone)]
pub struct BanHistory {
pub ban_times: u32,
pub last_banned_at: DateTime<Utc>,
}
#[async_trait]
pub trait BanStorage: Send + Sync {
async fn is_banned(&self, target: &BanTarget) -> Result<Option<BanRecord>, StorageError>;
async fn get_ban(&self, target: &BanTarget) -> Result<Option<BanRecord>, StorageError> {
self.is_banned(target).await
}
async fn add_ban(&self, record: &BanRecord) -> Result<(), StorageError> {
self.save(record).await
}
async fn save(&self, record: &BanRecord) -> Result<(), StorageError>;
async fn get_history(&self, target: &BanTarget) -> Result<Option<BanHistory>, StorageError>;
async fn increment_ban_times(&self, target: &BanTarget) -> Result<u64, StorageError>;
async fn get_ban_times(&self, target: &BanTarget) -> Result<u64, StorageError>;
async fn remove_ban(&self, target: &BanTarget) -> Result<(), StorageError>;
async fn cleanup_expired_bans(&self) -> Result<u64, StorageError>;
async fn list_bans(
&self,
active_only: bool,
offset: u64,
limit: u64,
) -> Result<Vec<BanRecord>, StorageError>;
fn as_any(&self) -> &dyn std::any::Any;
}
use ahash::AHashMap as HashMap;
use tokio::sync::RwLock;
pub struct MemoryStorage {
data: RwLock<HashMap<String, String>>,
expiration: RwLock<HashMap<String, u64>>,
}
pub struct MemoryBanStorage {
bans: RwLock<HashMap<BanTarget, BanRecord>>,
expiration: RwLock<HashMap<BanTarget, i64>>,
}
mod storage_impl;