use crate::{ReadResponseExt, request::RequestExt};
use publicsuffix::Psl as _;
use std::{
error::Error,
sync::{
LazyLock, RwLock,
atomic::{AtomicBool, Ordering},
},
thread,
time::{Duration, SystemTime},
};
static TTL: Duration = Duration::from_secs(24 * 60 * 60);
static CACHE: LazyLock<RwLock<ListCache>> = LazyLock::new(Default::default);
struct ListCache {
list: Option<publicsuffix::List>,
last_refreshed: Option<SystemTime>,
last_updated: Option<SystemTime>,
}
impl Default for ListCache {
fn default() -> Self {
Self {
list: None,
last_refreshed: None,
last_updated: None,
}
}
}
impl ListCache {
fn is_public_suffix(&self, domain: &[u8]) -> Option<bool> {
Some(
self.list
.as_ref()?
.suffix(domain)
.filter(publicsuffix::Suffix::is_known)
.filter(|suffix| suffix == &domain)
.is_some(),
)
}
fn needs_refreshed(&self) -> bool {
match self.last_refreshed {
Some(last_refreshed) => match last_refreshed.elapsed() {
Ok(elapsed) => elapsed > TTL,
Err(_) => false,
},
None => true,
}
}
fn refresh(&mut self) -> Result<(), Box<dyn Error>> {
let result = self.try_refresh();
self.last_refreshed = Some(SystemTime::now());
result
}
fn try_refresh(&mut self) -> Result<(), Box<dyn Error>> {
let mut request = http::Request::get(publicsuffix::LIST_URL);
if let Some(last_updated) = self.last_updated {
request = request.header(
http::header::IF_MODIFIED_SINCE,
httpdate::fmt_http_date(last_updated),
);
}
let mut response = request.body(())?.send()?;
match response.status() {
http::StatusCode::OK => {
self.list = Some(response.text()?.parse()?);
self.last_updated = Some(SystemTime::now());
tracing::debug!("public suffix list updated");
}
http::StatusCode::NOT_MODIFIED => {
self.last_updated = Some(SystemTime::now());
}
status => tracing::warn!(
"could not update public suffix list, got status code {}",
status,
),
}
Ok(())
}
}
pub(crate) fn is_public_suffix(domain: impl AsRef<str>) -> bool {
let domain = domain.as_ref().as_bytes();
let cache = CACHE.read().unwrap();
if cache.needs_refreshed() {
refresh_in_background();
}
if let Some(v) = cache.is_public_suffix(domain) {
return v;
}
drop(cache);
psl::suffix(domain)
.filter(psl::Suffix::is_known)
.filter(|suffix| suffix == &domain)
.is_some()
}
fn refresh_in_background() {
static IS_REFRESHING: AtomicBool = AtomicBool::new(false);
if IS_REFRESHING
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_ok()
{
thread::spawn(|| {
let mut cache = CACHE.write().unwrap();
if let Err(error) = cache.refresh() {
tracing::warn!(?error, "could not refresh public suffix list");
}
IS_REFRESHING.store(false, Ordering::SeqCst);
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn basic_is_public_suffix() {
assert!(is_public_suffix("co.jp"));
assert!(!is_public_suffix("google.com"));
}
#[test]
fn refresh_cache() {
let mut cache = ListCache::default();
assert!(cache.last_refreshed.is_none());
assert!(cache.last_updated.is_none());
assert!(cache.needs_refreshed());
cache.refresh().unwrap();
assert!(cache.last_refreshed.is_some());
assert!(cache.last_updated.is_some());
assert!(!cache.needs_refreshed());
let last_refreshed = cache.last_refreshed.unwrap();
let last_updated = cache.last_updated.unwrap();
cache.refresh().unwrap();
assert!(cache.last_refreshed.unwrap() > last_refreshed);
assert!(cache.last_updated.unwrap() > last_updated);
}
}