use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use time::OffsetDateTime;
use crate::Clock;
pub const DEFAULT_MAX_URL_LEN: usize = 2048;
pub const DEFAULT_MAX_BATCH: usize = 500;
pub const DEFAULT_CACHE_CAPACITY: usize = 4096;
pub const DEFAULT_NEGATIVE_TTL: Duration = Duration::from_secs(300);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum ThreatType {
Malware,
SocialEngineering,
UnwantedSoftware,
PotentiallyHarmfulApplication,
}
impl ThreatType {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Malware => "MALWARE",
Self::SocialEngineering => "SOCIAL_ENGINEERING",
Self::UnwantedSoftware => "UNWANTED_SOFTWARE",
Self::PotentiallyHarmfulApplication => "POTENTIALLY_HARMFUL_APPLICATION",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum UrlSafetySource {
GoogleWebRisk,
GoogleSafeBrowsing,
Cache,
}
impl UrlSafetySource {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::GoogleWebRisk => "google-webrisk",
Self::GoogleSafeBrowsing => "google-safebrowsing",
Self::Cache => "cache",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UrlVerdict {
pub url: String,
pub threats: Vec<ThreatType>,
pub source: UrlSafetySource,
pub checked_at: OffsetDateTime,
pub cache_ttl: Option<Duration>,
pub cached: bool,
}
impl UrlVerdict {
#[must_use]
pub fn is_threat(&self) -> bool {
!self.threats.is_empty()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum UrlSafetyLimit {
UrlLength,
BatchSize,
}
impl UrlSafetyLimit {
#[must_use]
pub const fn name(self) -> &'static str {
match self {
Self::UrlLength => "url length",
Self::BatchSize => "batch size",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UrlSafetyLimits {
pub max_url_len: usize,
pub max_batch: usize,
}
impl Default for UrlSafetyLimits {
fn default() -> Self {
Self {
max_url_len: DEFAULT_MAX_URL_LEN,
max_batch: DEFAULT_MAX_BATCH,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum UrlSafetyError {
InvalidUrl,
LimitExceeded {
limit: UrlSafetyLimit,
actual: usize,
max: usize,
},
Rejected { status: u16 },
Transient { retry_after: Option<Duration> },
Provider(String),
}
impl std::fmt::Display for UrlSafetyError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let scrub = crate::logging::scrub_text;
match self {
Self::InvalidUrl => f.write_str("not an http(s) URL; nothing was checked"),
Self::LimitExceeded { limit, actual, max } => {
write!(f, "over the {} ceiling: {actual} > {max}", limit.name())
}
Self::Rejected { status } => {
write!(
f,
"url safety provider rejected the request: status {status}"
)
}
Self::Transient { retry_after } => {
write!(
f,
"url safety call failed, retryable: retry after {retry_after:?}"
)
}
Self::Provider(message) => write!(f, "url safety provider failed: {}", scrub(message)),
}
}
}
impl std::error::Error for UrlSafetyError {}
impl UrlSafetyError {
#[must_use]
pub fn retry_after(&self) -> Option<Duration> {
match self {
Self::Transient { retry_after } => *retry_after,
_ => None,
}
}
}
pub fn check_url(url: &str, limits: &UrlSafetyLimits) -> Result<(), UrlSafetyError> {
if url.len() > limits.max_url_len {
return Err(UrlSafetyError::LimitExceeded {
limit: UrlSafetyLimit::UrlLength,
actual: url.len(),
max: limits.max_url_len,
});
}
let parsed = url::Url::parse(url).map_err(|_| UrlSafetyError::InvalidUrl)?;
if parsed.scheme() != "http" && parsed.scheme() != "https" {
return Err(UrlSafetyError::InvalidUrl);
}
Ok(())
}
pub fn check_batch(urls: &[String], limits: &UrlSafetyLimits) -> Result<(), UrlSafetyError> {
if urls.len() > limits.max_batch {
return Err(UrlSafetyError::LimitExceeded {
limit: UrlSafetyLimit::BatchSize,
actual: urls.len(),
max: limits.max_batch,
});
}
for url in urls {
check_url(url, limits)?;
}
Ok(())
}
#[async_trait]
pub trait UrlSafety: Send + Sync {
async fn check(&self, url: &str) -> Result<UrlVerdict, UrlSafetyError>;
async fn check_many(&self, urls: &[String]) -> Result<Vec<UrlVerdict>, UrlSafetyError> {
check_batch(urls, &UrlSafetyLimits::default())?;
let mut verdicts = Vec::with_capacity(urls.len());
for url in urls {
verdicts.push(self.check(url).await?);
}
Ok(verdicts)
}
}
#[allow(clippy::disallowed_types)]
type Guarded<T> = std::sync::Mutex<T>;
struct Entry {
verdict: UrlVerdict,
expires_at: OffsetDateTime,
}
impl Entry {
fn fresh_at(&self, now: OffsetDateTime) -> bool {
now < self.expires_at
}
}
pub struct CachedUrlSafety {
inner: Arc<dyn UrlSafety>,
clock: Arc<dyn Clock>,
#[allow(clippy::disallowed_types)]
state: Guarded<HashMap<String, Entry>>,
capacity: usize,
negative_ttl: Duration,
}
impl CachedUrlSafety {
#[must_use]
pub fn new(inner: Arc<dyn UrlSafety>, clock: Arc<dyn Clock>) -> Self {
Self {
inner,
clock,
state: Guarded::new(HashMap::new()),
capacity: DEFAULT_CACHE_CAPACITY,
negative_ttl: DEFAULT_NEGATIVE_TTL,
}
}
#[must_use]
pub fn with_capacity(mut self, capacity: usize) -> Self {
self.capacity = if capacity == 0 {
DEFAULT_CACHE_CAPACITY
} else {
capacity
};
self
}
#[must_use]
pub fn with_negative_ttl(mut self, ttl: Duration) -> Self {
self.negative_ttl = ttl;
self
}
fn fresh(&self, url: &str) -> Option<UrlVerdict> {
let now = self.clock.now();
let state = self.state.lock().expect("url safety cache");
state
.get(url)
.filter(|entry| entry.fresh_at(now))
.map(Self::served)
}
fn served(entry: &Entry) -> UrlVerdict {
let mut verdict = entry.verdict.clone();
verdict.cached = true;
verdict
}
fn store(&self, verdicts: &[UrlVerdict]) {
if verdicts.is_empty() {
return;
}
let now = self.clock.now();
let mut state = self.state.lock().expect("url safety cache");
for verdict in verdicts {
Self::insert(&mut state, verdict, now, self.capacity, self.negative_ttl);
}
}
fn insert(
state: &mut HashMap<String, Entry>,
verdict: &UrlVerdict,
now: OffsetDateTime,
capacity: usize,
negative_ttl: Duration,
) {
state.retain(|_, entry| entry.fresh_at(now));
while state.len() >= capacity {
let Some(oldest) = state
.iter()
.min_by_key(|(_, entry)| entry.expires_at)
.map(|(url, _)| url.clone())
else {
break;
};
state.remove(&oldest);
}
let ttl = verdict.cache_ttl.unwrap_or(negative_ttl);
let expires_at = verdict
.checked_at
.saturating_add(time::Duration::try_from(ttl).unwrap_or(time::Duration::MAX));
state.insert(
verdict.url.clone(),
Entry {
verdict: verdict.clone(),
expires_at,
},
);
}
fn partition(&self, urls: &[String]) -> (Vec<Option<UrlVerdict>>, Vec<String>) {
let now = self.clock.now();
let state = self.state.lock().expect("url safety cache");
let mut slots = Vec::with_capacity(urls.len());
let mut misses = Vec::new();
for url in urls {
let hit = state
.get(url)
.filter(|entry| entry.fresh_at(now))
.map(Self::served);
if hit.is_none() {
misses.push(url.clone());
}
slots.push(hit);
}
drop(state);
(slots, misses)
}
}
#[async_trait]
impl UrlSafety for CachedUrlSafety {
async fn check(&self, url: &str) -> Result<UrlVerdict, UrlSafetyError> {
if let Some(verdict) = self.fresh(url) {
return Ok(verdict);
}
let verdict = self.inner.check(url).await?;
self.store(std::slice::from_ref(&verdict));
Ok(verdict)
}
async fn check_many(&self, urls: &[String]) -> Result<Vec<UrlVerdict>, UrlSafetyError> {
check_batch(urls, &UrlSafetyLimits::default())?;
let (mut slots, misses) = self.partition(urls);
if !misses.is_empty() {
let fetched = self.inner.check_many(&misses).await?;
self.store(&fetched);
let mut fetched = fetched.into_iter();
for slot in &mut slots {
if slot.is_none() {
*slot = fetched.next();
}
}
}
Ok(slots.into_iter().flatten().collect())
}
}
impl std::fmt::Debug for CachedUrlSafety {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let entries = self
.state
.lock()
.map_or(0, |state: std::sync::MutexGuard<'_, _>| state.len());
f.debug_struct("CachedUrlSafety")
.field("capacity", &self.capacity)
.field("negative_ttl", &self.negative_ttl)
.field("entries", &entries)
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering};
fn limits() -> UrlSafetyLimits {
UrlSafetyLimits {
max_url_len: 32,
max_batch: 2,
}
}
#[test]
fn threat_types_and_sources_name_themselves() {
assert_eq!(ThreatType::Malware.as_str(), "MALWARE");
assert_eq!(ThreatType::SocialEngineering.as_str(), "SOCIAL_ENGINEERING");
assert_eq!(ThreatType::UnwantedSoftware.as_str(), "UNWANTED_SOFTWARE");
assert_eq!(
ThreatType::PotentiallyHarmfulApplication.as_str(),
"POTENTIALLY_HARMFUL_APPLICATION"
);
assert_eq!(UrlSafetySource::GoogleWebRisk.as_str(), "google-webrisk");
assert_eq!(
UrlSafetySource::GoogleSafeBrowsing.as_str(),
"google-safebrowsing"
);
assert_eq!(UrlSafetySource::Cache.as_str(), "cache");
}
#[test]
fn check_url_accepts_http_and_https_and_refuses_the_rest() {
let limits = UrlSafetyLimits::default();
assert_eq!(check_url("https://example.com/", &limits), Ok(()));
assert_eq!(check_url("http://example.com/a?b=c", &limits), Ok(()));
assert_eq!(
check_url("ftp://example.com/file", &limits),
Err(UrlSafetyError::InvalidUrl)
);
assert_eq!(
check_url("example.com/path", &limits),
Err(UrlSafetyError::InvalidUrl)
);
assert_eq!(check_url("", &limits), Err(UrlSafetyError::InvalidUrl));
assert_eq!(
check_url("mailto:someone@example.com", &limits),
Err(UrlSafetyError::InvalidUrl)
);
}
#[test]
fn check_url_refuses_an_over_long_url_as_a_limit() {
let limits = limits();
let fits = format!("https://example.com/{}", "a".repeat(12));
assert_eq!(fits.len(), 32);
assert_eq!(check_url(&fits, &limits), Ok(()));
let long = format!("https://example.com/{}", "a".repeat(13));
assert_eq!(
check_url(&long, &limits),
Err(UrlSafetyError::LimitExceeded {
limit: UrlSafetyLimit::UrlLength,
actual: 33,
max: 32,
})
);
}
#[test]
fn check_batch_bounds_the_size_and_still_checks_each_url() {
let limits = limits();
assert_eq!(check_batch(&[], &limits), Ok(()), "an empty batch is fine");
let ok = vec![
"https://a.example".to_owned(),
"https://b.example".to_owned(),
];
assert_eq!(check_batch(&ok, &limits), Ok(()));
let oversized = vec![
"https://a.example".to_owned(),
"https://b.example".to_owned(),
"https://c.example".to_owned(),
];
assert_eq!(
check_batch(&oversized, &limits),
Err(UrlSafetyError::LimitExceeded {
limit: UrlSafetyLimit::BatchSize,
actual: 3,
max: 2,
})
);
let one_bad = vec!["https://a.example".to_owned(), "ftp://b.example".to_owned()];
assert_eq!(
check_batch(&one_bad, &limits),
Err(UrlSafetyError::InvalidUrl)
);
}
#[test]
fn provider_text_is_scrubbed_in_display_but_not_debug() {
let error =
UrlSafetyError::Provider("400 for https://api.example.test?token=secret".to_owned());
let text = error.to_string();
assert!(!text.contains("secret"), "{text}");
assert!(format!("{error:?}").contains("secret"));
assert_eq!(
UrlSafetyError::Rejected { status: 429 }.to_string(),
"url safety provider rejected the request: status 429"
);
assert_eq!(
UrlSafetyError::InvalidUrl.to_string(),
"not an http(s) URL; nothing was checked"
);
}
#[test]
fn transient_carries_an_optional_retry_after() {
let throttled = UrlSafetyError::Transient {
retry_after: Some(Duration::from_secs(30)),
};
assert_eq!(throttled.retry_after(), Some(Duration::from_secs(30)));
assert_eq!(
UrlSafetyError::Transient { retry_after: None }.retry_after(),
None
);
assert_eq!(UrlSafetyError::InvalidUrl.retry_after(), None);
}
struct FakeClock(AtomicI64);
impl FakeClock {
fn at(secs: i64) -> Arc<Self> {
Arc::new(Self(AtomicI64::new(secs)))
}
fn advance(&self, secs: i64) {
self.0.fetch_add(secs, Ordering::Relaxed);
}
}
impl Clock for FakeClock {
fn now(&self) -> OffsetDateTime {
OffsetDateTime::from_unix_timestamp(self.0.load(Ordering::Relaxed)).expect("in range")
}
}
struct FakeInner {
clock: Arc<FakeClock>,
checks: AtomicUsize,
batches: AtomicUsize,
fail: AtomicBool,
ttl: Option<Duration>,
}
impl FakeInner {
fn checks(&self) -> usize {
self.checks.load(Ordering::Relaxed)
}
fn batches(&self) -> usize {
self.batches.load(Ordering::Relaxed)
}
fn verdict(&self, url: &str) -> UrlVerdict {
let threats = if url.contains("malware") {
vec![ThreatType::Malware]
} else if url.contains("phishing") {
vec![ThreatType::SocialEngineering]
} else {
Vec::new()
};
UrlVerdict {
url: url.to_owned(),
threats,
source: UrlSafetySource::GoogleWebRisk,
checked_at: self.clock.now(),
cache_ttl: self.ttl,
cached: false,
}
}
}
#[async_trait]
impl UrlSafety for FakeInner {
async fn check(&self, url: &str) -> Result<UrlVerdict, UrlSafetyError> {
self.checks.fetch_add(1, Ordering::Relaxed);
if self.fail.load(Ordering::Relaxed) {
return Err(UrlSafetyError::Provider(format!("boom for {url}")));
}
Ok(self.verdict(url))
}
async fn check_many(&self, urls: &[String]) -> Result<Vec<UrlVerdict>, UrlSafetyError> {
self.batches.fetch_add(1, Ordering::Relaxed);
let mut verdicts = Vec::with_capacity(urls.len());
for url in urls {
if self.fail.load(Ordering::Relaxed) {
return Err(UrlSafetyError::Provider(format!("boom for {url}")));
}
verdicts.push(self.verdict(url));
}
Ok(verdicts)
}
}
type Parts = (Arc<FakeClock>, Arc<FakeInner>, CachedUrlSafety);
fn cache_with(ttl: Option<Duration>, fail: bool) -> Parts {
let clock = FakeClock::at(1_000);
let inner = Arc::new(FakeInner {
clock: clock.clone(),
checks: AtomicUsize::new(0),
batches: AtomicUsize::new(0),
fail: AtomicBool::new(fail),
ttl,
});
let cache = CachedUrlSafety::new(inner.clone(), clock.clone());
(clock, inner, cache)
}
#[test]
fn a_hit_within_the_ttl_serves_the_cache_and_a_miss_goes_to_the_provider() {
let (clock, inner, cache) = cache_with(Some(Duration::from_secs(60)), false);
let url = "https://example.test/page";
let first = pollster::block_on(cache.check(url)).expect("checked");
assert_eq!(inner.checks(), 1);
assert!(!first.cached);
assert_eq!(
first.url, url,
"the verdict's url equals the input, verbatim"
);
assert!(!first.is_threat());
clock.advance(59);
let second = pollster::block_on(cache.check(url)).expect("checked");
assert_eq!(inner.checks(), 1, "served from the cache");
assert!(second.cached);
assert_eq!(second.threats, first.threats);
assert_eq!(second.checked_at, first.checked_at);
clock.advance(2); pollster::block_on(cache.check(url)).expect("checked");
assert_eq!(inner.checks(), 2, "past the TTL the provider answers again");
}
#[test]
fn a_verdict_without_a_provider_ttl_uses_the_negative_ttl() {
let (clock, inner, cache) = cache_with(None, false);
let cache = cache.with_negative_ttl(Duration::from_secs(30));
pollster::block_on(cache.check("https://example.test")).expect("checked");
clock.advance(29);
pollster::block_on(cache.check("https://example.test")).expect("checked");
assert_eq!(inner.checks(), 1, "still inside the negative TTL");
clock.advance(2); pollster::block_on(cache.check("https://example.test")).expect("checked");
assert_eq!(inner.checks(), 2, "past the negative TTL");
}
#[test]
fn check_many_merges_in_input_order_with_one_batch_call_for_the_misses() {
let (_, inner, cache) = cache_with(Some(Duration::from_secs(300)), false);
pollster::block_on(cache.check("https://a.example")).expect("warmed");
assert_eq!(inner.checks(), 1);
let urls = vec![
"https://malware.example".to_owned(),
"https://a.example".to_owned(),
"https://phishing.example".to_owned(),
];
let verdicts = pollster::block_on(cache.check_many(&urls)).expect("checked");
assert_eq!(verdicts.len(), urls.len(), "one verdict per input");
assert_eq!(inner.batches(), 1, "one batch call covered every miss");
assert_eq!(inner.checks(), 1, "the warm URL was not re-checked");
assert_eq!(verdicts[0].url, urls[0]);
assert_eq!(verdicts[0].threats, vec![ThreatType::Malware]);
assert!(!verdicts[0].cached);
assert_eq!(verdicts[1].url, urls[1], "the hit kept its input position");
assert!(verdicts[1].cached);
assert_eq!(verdicts[2].url, urls[2]);
assert_eq!(verdicts[2].threats, vec![ThreatType::SocialEngineering]);
assert!(!verdicts[2].cached);
let again = pollster::block_on(cache.check_many(&urls)).expect("checked");
assert_eq!(inner.batches(), 1, "nothing new reached the provider");
assert!(again.iter().all(|verdict| verdict.cached));
}
#[test]
fn the_cache_checks_limits_before_the_provider() {
let (_, inner, cache) = cache_with(Some(Duration::from_secs(60)), false);
let urls = (0..501)
.map(|n| format!("https://{n}.example"))
.collect::<Vec<_>>();
let error = pollster::block_on(cache.check_many(&urls)).expect_err("over the ceiling");
assert!(matches!(
error,
UrlSafetyError::LimitExceeded {
limit: UrlSafetyLimit::BatchSize,
..
}
));
assert_eq!(inner.batches(), 0, "the provider was never asked");
}
#[test]
fn a_full_cache_evicts_the_earliest_expiring_entry() {
let (_, inner, cache) = cache_with(Some(Duration::from_secs(60)), false);
let cache = cache.with_capacity(1);
pollster::block_on(cache.check("https://first.example")).expect("checked");
pollster::block_on(cache.check("https://second.example")).expect("checked");
assert_eq!(inner.checks(), 2);
pollster::block_on(cache.check("https://first.example")).expect("checked");
assert_eq!(inner.checks(), 3);
pollster::block_on(cache.check("https://second.example")).expect("checked");
assert_eq!(inner.checks(), 4);
}
#[test]
fn an_error_propagates_and_is_never_cached() {
let (_, inner, cache) = cache_with(Some(Duration::from_secs(60)), true);
let error = pollster::block_on(cache.check("https://example.test")).expect_err("failed");
assert!(matches!(error, UrlSafetyError::Provider(_)));
pollster::block_on(cache.check("https://example.test")).expect_err("failed again");
assert_eq!(inner.checks(), 2, "a failed check must not be cached");
}
}