Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion src/cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,20 @@ struct CacheInner {
// Small in-memory hot index for quick negatives + stats. Bodies live on disk.
known: HashMap<String, ()>,
disk: disk::DiskStore,

// BAN patterns; if request URL contains any pattern, treat as miss.
bans: Vec<String>,
}

impl Cache {
pub fn new(cache_dir: &str) -> Result<Self, String> {
let disk = disk::DiskStore::open(cache_dir)?;
let bans = disk.list_bans().unwrap_or_default();
Ok(Self {
inner: Arc::new(RwLock::new(CacheInner {
known: HashMap::new(),
disk,
bans,
})),
})
}
Expand All @@ -41,6 +46,19 @@ impl Cache {
let inner = self.inner.read();
inner.disk.remove_by_url(normalized_url)
}

pub fn add_ban(&self, pattern: &str) -> Result<u64, String> {
let mut inner = self.inner.write();
let id = inner.disk.add_ban(pattern)?;
inner.bans.push(pattern.to_string());
Ok(id)
}

pub fn is_banned(&self, normalized_url: &Uri) -> bool {
let inner = self.inner.read();
let u = normalized_url.to_string();
inner.bans.iter().any(|p| !p.is_empty() && u.contains(p))
}
}

pub async fn handle_cached(
Expand All @@ -53,7 +71,10 @@ pub async fn handle_cached(

let cache_key = build_cache_key(&norm_uri, &parts.headers);

if let Some(resp) = lookup(&state.cache, &cache_key, &parts.headers, &norm_uri)? {
// BAN logic: if URL is banned, skip cache.
if state.cache.is_banned(&norm_uri) {
// fall through to upstream
} else if let Some(resp) = lookup(&state.cache, &cache_key, &parts.headers, &norm_uri)? {
return Ok(resp);
}

Expand Down
22 changes: 22 additions & 0 deletions src/disk.rs
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,28 @@ impl DiskStore {
}
Ok(gone)
}

pub fn add_ban(&self, pattern: &str) -> Result<u64, String> {
let _g = self.lock.lock();
let id = now_ms();
let k = format!("ban:{id}");
self.db
.insert(k.as_bytes(), pattern.as_bytes())
.map_err(|e| format!("sled insert: {e}"))?;
self.db.flush().map_err(|e| format!("sled flush: {e}"))?;
Ok(id)
}

pub fn list_bans(&self) -> Result<Vec<String>, String> {
let mut out = Vec::new();
for item in self.db.scan_prefix(b"ban:") {
let (_k, v) = item.map_err(|e| format!("sled iter: {e}"))?;
if let Ok(s) = String::from_utf8(v.to_vec()) {
out.push(s);
}
}
Ok(out)
}
}

pub fn headers_to_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
Expand Down
47 changes: 19 additions & 28 deletions src/purge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,48 +30,39 @@ pub fn handle_purge(
}
}

pub fn handle_ban(_cache: std::sync::Arc<Cache>, uri: &Uri) -> (StatusCode, String) {
// Varnish BAN is pattern-based. We can implement later with a ban list evaluated at lookup.
// For now, return 501 so users know it's not implemented.
(
StatusCode::NOT_IMPLEMENTED,
format!(
"BAN not implemented (requested ban on URLs containing {})",
uri.path()
pub fn handle_ban(cache: std::sync::Arc<Cache>, uri: &Uri) -> (StatusCode, String) {
// Minimal Varnish-like BAN: store a substring pattern and check it during cache lookup.
// We store the path component as the pattern by default.
let pattern = uri.path();
match cache.add_ban(pattern) {
Ok(id) => (
StatusCode::OK,
format!("Added ban #{id} for URLs containing '{pattern}'"),
),
)
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e),
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn purge_with_xkey_is_ok_even_if_no_objects_match() {
fn ban_marks_matching_urls_as_banned() {
let dir = tempfile::tempdir().unwrap();
let cache = std::sync::Arc::new(Cache::new(dir.path().to_str().unwrap()).unwrap());

let uri: Uri = "/foo".parse().unwrap();
let mut headers = HeaderMap::new();
headers.insert("xkey", http::HeaderValue::from_static("a b c"));

let (status, body) = handle_purge(cache, &uri, &headers);
let (status, body) = handle_ban(cache.clone(), &uri);
assert_eq!(status, StatusCode::OK);
assert!(body.contains("Invalidated"));
}
assert!(body.contains("Added ban"));

#[test]
fn purge_by_url_is_ok_even_when_nothing_matches() {
let dir = tempfile::tempdir().unwrap();
let cache = std::sync::Arc::new(Cache::new(dir.path().to_str().unwrap()).unwrap());
let norm: Uri = "/foo?a=1".parse().unwrap();
let norm = crate::normalize::normalize_uri(&norm);
assert!(cache.is_banned(&norm));

// Includes tracking params; purge should normalize first.
let uri: Uri = "/foo?utm_source=x&a=1".parse().unwrap();
let headers = HeaderMap::new();

let (status, body) = handle_purge(cache, &uri, &headers);
assert_eq!(status, StatusCode::OK);
assert!(body.contains("Invalidated"));
assert!(body.contains("/foo?a=1"));
let other: Uri = "/bar".parse().unwrap();
let other = crate::normalize::normalize_uri(&other);
assert!(!cache.is_banned(&other));
}
}