use ahash::AHashMap as HashMap;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Identifier {
UserId(String),
Ip(String),
Mac(String),
ApiKey(String),
DeviceId(String),
}
impl Identifier {
pub fn as_str(&self) -> &str {
match self {
Identifier::UserId(s) => s,
Identifier::Ip(s) => s,
Identifier::Mac(s) => s,
Identifier::ApiKey(s) => s,
Identifier::DeviceId(s) => s,
}
}
pub fn type_name(&self) -> &'static str {
match self {
Identifier::UserId(_) => "user_id",
Identifier::Ip(_) => "ip",
Identifier::Mac(_) => "mac",
Identifier::ApiKey(_) => "api_key",
Identifier::DeviceId(_) => "device_id",
}
}
pub fn key(&self) -> String {
format!("{}:{}", self.type_name(), self.as_str())
}
#[cfg(feature = "ban-manager")]
pub fn to_ban_target(&self) -> Option<crate::storage::BanTarget> {
match self {
Identifier::UserId(id) => Some(crate::storage::BanTarget::UserId(id.clone())),
Identifier::Ip(ip) => Some(crate::storage::BanTarget::Ip(ip.clone())),
Identifier::Mac(mac) => Some(crate::storage::BanTarget::Mac(mac.clone())),
Identifier::ApiKey(_) | Identifier::DeviceId(_) => None,
}
}
}
#[derive(Clone)]
pub struct RequestContext {
pub user_id: Option<String>,
pub ip: Option<String>,
pub mac: Option<String>,
pub device_id: Option<String>,
pub api_key: Option<String>,
pub headers: HashMap<String, String>,
pub path: String,
pub method: String,
pub client_ip: Option<String>,
pub query_params: HashMap<String, String>,
}
impl std::fmt::Debug for RequestContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut debug = f.debug_struct("RequestContext");
debug
.field("user_id", &self.user_id)
.field("ip", &self.ip)
.field("mac", &self.mac)
.field("device_id", &self.device_id)
.field("api_key", &self.api_key.as_ref().map(|_| "***"));
let headers: HashMap<String, String> = self
.headers
.iter()
.map(|(k, v)| {
let v = if k.to_lowercase().contains("auth")
|| k.to_lowercase().contains("cookie")
|| k.to_lowercase().contains("key")
{
"***".to_string()
} else {
v.clone()
};
(k.clone(), v)
})
.collect();
debug.field("headers", &headers);
debug
.field("path", &self.path)
.field("method", &self.method)
.field("client_ip", &self.client_ip);
let query_params: HashMap<String, String> = self
.query_params
.iter()
.map(|(k, v)| {
let v = if k.to_lowercase().contains("token")
|| k.to_lowercase().contains("key")
|| k.to_lowercase().contains("secret")
{
"***".to_string()
} else {
v.clone()
};
(k.clone(), v)
})
.collect();
debug.field("query_params", &query_params);
debug.finish()
}
}
impl RequestContext {
pub fn new() -> Self {
Self {
user_id: None,
ip: None,
mac: None,
device_id: None,
api_key: None,
headers: HashMap::new(),
path: String::new(),
method: String::new(),
client_ip: None,
query_params: HashMap::new(),
}
}
pub fn with_header(mut self, key: &str, value: &str) -> Self {
self.headers.insert(key.to_lowercase(), value.to_string());
self
}
pub fn with_client_ip(mut self, ip: &str) -> Self {
self.client_ip = Some(ip.to_string());
self
}
pub fn with_query_param(mut self, key: &str, value: &str) -> Self {
self.query_params.insert(key.to_string(), value.to_string());
self
}
pub fn with_path(mut self, path: &str) -> Self {
self.path = path.to_string();
self
}
pub fn with_method(mut self, method: &str) -> Self {
self.method = method.to_string();
self
}
pub fn get_header(&self, key: &str) -> Option<&String> {
self.headers.get(&key.to_lowercase())
}
}
impl Default for RequestContext {
fn default() -> Self {
Self::new()
}
}
pub trait IdentifierExtractor: Send + Sync {
fn extract(&self, context: &RequestContext) -> Option<Identifier>;
fn name(&self) -> &str;
}
#[cfg(test)]
mod tests {
#[allow(unused_imports)]
use super::*;
#[cfg(feature = "ban-manager")]
#[test]
fn test_identifier_mac_to_ban_target() {
let id = Identifier::Mac("AA:BB:CC:DD:EE:FF".to_string());
let target = id.to_ban_target();
assert!(target.is_some());
match target.unwrap() {
crate::storage::BanTarget::Mac(mac) => {
assert_eq!(mac, "AA:BB:CC:DD:EE:FF");
}
other => panic!("expected Mac, got: {:?}", other),
}
}
#[cfg(feature = "ban-manager")]
#[test]
fn test_identifier_apikey_to_ban_target_none() {
let id = Identifier::ApiKey("key123".to_string());
assert!(id.to_ban_target().is_none());
}
}