affine_core 0.0.2

AFFiNE primitive core.
Documentation
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;