use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use sha2::{Digest, Sha256};
use super::{InvitePolicyConfig, RateLimitPolicyError, RateLimitScope, SourceFacts};
const HIGH_RISK_TARGET_DOMAINS: &[&str] = &[
"qq.com",
"proton.me",
"protonmail.com",
"163.com",
"126.com",
"outlook.com",
"hotmail.com",
];
pub(super) fn normalize_domain(domain: &str) -> String {
domain.trim().trim_end_matches('.').to_ascii_lowercase()
}
pub(super) fn email_domain(email: &str) -> Result<String, RateLimitPolicyError> {
let email = email.trim();
if email.matches('@').count() != 1 {
return Err(RateLimitPolicyError::InvalidEmail);
}
let (local, domain) = email.split_once('@').ok_or(RateLimitPolicyError::InvalidEmail)?;
if local.is_empty()
|| local.starts_with('.')
|| local.ends_with('.')
|| local.contains("..")
|| !local
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'+' | b'-'))
{
return Err(RateLimitPolicyError::InvalidEmail);
}
let raw_domain = domain.trim();
if raw_domain.ends_with('.') {
return Err(RateLimitPolicyError::InvalidEmail);
}
let domain = normalize_domain(raw_domain);
if domain.is_empty()
|| !domain.contains('.')
|| domain.split('.').any(|label| {
label.is_empty()
|| !label
.chars()
.all(|character| character.is_ascii_alphanumeric() || character == '-')
|| label.starts_with('-')
|| label.ends_with('-')
})
{
return Err(RateLimitPolicyError::InvalidEmail);
}
Ok(domain)
}
pub(super) fn short_hash(value: &str) -> String {
hex::encode(Sha256::digest(value.as_bytes()))[..24].to_string()
}
pub(super) fn subject_hash(email: &str, config: &InvitePolicyConfig<'_>) -> String {
let normalized = email.trim().to_ascii_lowercase();
let hash = Sha256::digest(format!("{}:{normalized}", config.subject_hash_salt).as_bytes());
format!("actor_email_sha256:v1:{}", hex::encode(hash))
}
pub(super) fn workspace_subject_key(workspace_id: &str) -> String {
format!("workspace:v1:{}", short_hash(workspace_id))
}
pub(super) fn source_cohort_subject_key(prefix: &str, domain: &str) -> String {
format!(
"source_prefix_domain_sha256:v1:{}",
short_hash(&format!("{}:{}", prefix, normalize_domain(domain)))
)
}
pub(super) fn source_prefix(source: Option<&SourceFacts>) -> Option<String> {
let source = source?;
if !source.trusted {
return None;
}
let ip = source.ip.as_ref()?.parse::<IpAddr>().ok()?;
match ip {
IpAddr::V4(ip) => {
let octets = ip.octets();
Some(Ipv4Addr::new(octets[0], octets[1], octets[2], 0).to_string() + "/24")
}
IpAddr::V6(ip) => {
let segments = ip.segments();
Some(Ipv6Addr::new(segments[0], segments[1], segments[2], 0, 0, 0, 0, 0).to_string() + "/48")
}
}
}
pub(super) fn high_risk_domain(domain: &str) -> bool {
let domain = normalize_domain(domain);
HIGH_RISK_TARGET_DOMAINS.contains(&domain.as_str())
}
pub(super) fn scope(scope_key: String, window_seconds: i32, limit: i32, requested: i32) -> RateLimitScope {
RateLimitScope {
scope_key,
window_seconds,
bucket_seconds: bucket_seconds(window_seconds),
limit,
requested,
}
}
pub fn auth_refresh_scope(selector: &str) -> RateLimitScope {
scope(
format!("auth:session_refresh:{selector}"),
crate::auth::AUTH_REFRESH_WINDOW_SECONDS,
crate::auth::AUTH_REFRESH_LIMIT,
1,
)
}
fn bucket_seconds(window_seconds: i32) -> i64 {
match window_seconds {
60 => 10,
3600 => 5 * 60,
86_400 => 60 * 60,
604_800 => 6 * 60 * 60,
_ => 60,
}
}
#[cfg(test)]
#[path = "../tests/rate_limit/scope/tests.rs"]
mod tests;