use std::collections::{HashMap, VecDeque};
use std::sync::Mutex;
use std::time::{Duration, Instant};
use bytes::Bytes;
use crate::fs::Entry;
use crate::sftp::wire::Attrs;
pub const DEFAULT_TTL: Duration = Duration::from_secs(2);
pub const DEFAULT_BODY_CAP: usize = 64 * 1024 * 1024;
struct Listing {
entries: HashMap<String, (Attrs, Option<String>)>,
fetched: Instant,
}
#[derive(Clone, PartialEq, Eq, Hash)]
struct BodyKey {
path: String,
mtime: u32,
size: u64,
}
struct Bodies {
map: HashMap<BodyKey, Bytes>,
order: VecDeque<BodyKey>,
bytes: usize,
cap: usize,
}
impl Bodies {
fn insert(&mut self, key: BodyKey, body: Bytes) {
if body.len() > self.cap || self.map.contains_key(&key) {
return;
}
self.bytes += body.len();
self.order.push_back(key.clone());
self.map.insert(key, body);
while self.bytes > self.cap {
let Some(oldest) = self.order.pop_front() else {
break;
};
if let Some(dropped) = self.map.remove(&oldest) {
self.bytes -= dropped.len();
}
}
}
}
pub struct Cache {
listings: Mutex<HashMap<String, Listing>>,
bodies: Mutex<Bodies>,
ttl: Duration,
}
impl Cache {
pub fn new(ttl: Duration, body_cap: usize) -> Self {
Self {
listings: Mutex::new(HashMap::new()),
bodies: Mutex::new(Bodies {
map: HashMap::new(),
order: VecDeque::new(),
bytes: 0,
cap: body_cap,
}),
ttl,
}
}
pub fn has_listing(&self, dir: &str) -> bool {
self.listings
.lock()
.expect("listing cache poisoned")
.get(dir)
.is_some_and(|l| l.fetched.elapsed() < self.ttl)
}
pub fn attrs_of(&self, dir: &str, name: &str) -> Option<Attrs> {
let listings = self.listings.lock().expect("listing cache poisoned");
let listing = listings.get(dir)?;
if listing.fetched.elapsed() >= self.ttl {
return None;
}
listing.entries.get(name).map(|(attrs, _)| *attrs)
}
pub fn listing_entries(&self, dir: &str) -> Option<Vec<Entry>> {
let listings = self.listings.lock().expect("listing cache poisoned");
let listing = listings.get(dir)?;
if listing.fetched.elapsed() >= self.ttl {
return None;
}
Some(
listing
.entries
.iter()
.map(|(name, (attrs, owner))| Entry {
name: name.clone(),
attrs: *attrs,
owner: owner.clone(),
})
.collect(),
)
}
pub fn put_listing(&self, dir: &str, entries: &[Entry]) {
let map = entries
.iter()
.map(|e| (e.name.clone(), (e.attrs, e.owner.clone())))
.collect::<HashMap<_, _>>();
self.listings
.lock()
.expect("listing cache poisoned")
.insert(
dir.to_string(),
Listing {
entries: map,
fetched: Instant::now(),
},
);
}
pub fn forget_listing(&self, dir: &str) {
self.listings
.lock()
.expect("listing cache poisoned")
.remove(dir);
}
pub fn body(&self, path: &str, attrs: &Attrs) -> Option<Bytes> {
let key = body_key(path, attrs)?;
self.bodies
.lock()
.expect("body cache poisoned")
.map
.get(&key)
.cloned()
}
pub fn put_body(&self, path: &str, attrs: &Attrs, body: Bytes) {
let Some(key) = body_key(path, attrs) else {
return;
};
self.bodies
.lock()
.expect("body cache poisoned")
.insert(key, body);
}
}
impl Default for Cache {
fn default() -> Self {
Self::new(DEFAULT_TTL, DEFAULT_BODY_CAP)
}
}
fn body_key(path: &str, attrs: &Attrs) -> Option<BodyKey> {
Some(BodyKey {
path: path.to_string(),
mtime: attrs.mtime?,
size: attrs.size?,
})
}
pub fn etag(attrs: &Attrs) -> Option<String> {
let mtime = attrs.mtime?;
let size = attrs.size?;
Some(format!("W/\"{mtime:x}-{size:x}\""))
}
pub fn etag_matches(header: &str, tag: &str) -> bool {
let want = normalise(tag);
header
.split(',')
.any(|candidate| candidate.trim() == "*" || normalise(candidate) == want)
}
fn normalise(tag: &str) -> &str {
tag.trim().trim_start_matches("W/").trim_matches('"')
}
#[cfg(test)]
mod tests {
use super::*;
fn attrs(mtime: u32, size: u64) -> Attrs {
Attrs {
mtime: Some(mtime),
size: Some(size),
..Attrs::default()
}
}
fn entry(name: &str, a: Attrs) -> Entry {
Entry {
name: name.to_string(),
attrs: a,
owner: Some("souta".to_string()),
}
}
#[test]
fn a_fresh_listing_answers_without_the_remote() {
let c = Cache::default();
c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
assert!(c.has_listing("/srv"));
assert_eq!(c.attrs_of("/srv", "a.html").and_then(|a| a.size), Some(7));
assert!(c.attrs_of("/srv", "missing.html").is_none());
}
#[test]
fn a_stale_listing_is_not_used() {
let c = Cache::new(Duration::ZERO, DEFAULT_BODY_CAP);
c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
assert!(!c.has_listing("/srv"));
assert!(c.attrs_of("/srv", "a.html").is_none());
}
#[test]
fn forgetting_a_listing_forces_a_refetch() {
let c = Cache::default();
c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
c.forget_listing("/srv");
assert!(!c.has_listing("/srv"));
}
#[test]
fn a_changed_file_misses_the_cache() {
let c = Cache::default();
let old = attrs(100, 7);
c.put_body("/srv/a.html", &old, Bytes::from_static(b"old"));
assert_eq!(
c.body("/srv/a.html", &old).as_deref(),
Some(&b"old"[..]),
"the version that was cached is served"
);
let rebuilt_same_size = attrs(200, 7);
assert!(
c.body("/srv/a.html", &rebuilt_same_size).is_none(),
"a new mtime must miss even when the size is unchanged"
);
let rebuilt_same_mtime = attrs(100, 9);
assert!(
c.body("/srv/a.html", &rebuilt_same_mtime).is_none(),
"a new size must miss even when the mtime is unchanged"
);
}
#[test]
fn a_body_without_mtime_or_size_is_not_cached() {
let c = Cache::default();
let bare = Attrs::default();
c.put_body("/srv/a.html", &bare, Bytes::from_static(b"x"));
assert!(c.body("/srv/a.html", &bare).is_none());
}
#[test]
fn the_body_cache_respects_its_budget() {
let c = Cache::new(DEFAULT_TTL, 10);
for i in 0..5u32 {
c.put_body(&format!("/f{i}"), &attrs(i, 4), Bytes::from_static(b"1234"));
}
let held = c.bodies.lock().expect("body cache").bytes;
assert!(held <= 10, "cache holds {held} bytes over a 10 byte budget");
assert!(c.body("/f4", &attrs(4, 4)).is_some());
assert!(c.body("/f0", &attrs(0, 4)).is_none());
}
#[test]
fn a_file_bigger_than_the_budget_is_declined_rather_than_flushing_everything() {
let c = Cache::new(DEFAULT_TTL, 4);
c.put_body("/small", &attrs(1, 2), Bytes::from_static(b"ab"));
c.put_body("/huge", &attrs(2, 99), Bytes::from_static(b"0123456789"));
assert!(c.body("/huge", &attrs(2, 99)).is_none());
assert!(
c.body("/small", &attrs(1, 2)).is_some(),
"an oversized insert must not flush the cache"
);
}
#[test]
fn etags_compare_weakly() {
let a = attrs(0x64, 0x7);
let tag = etag(&a).expect("attrs carry mtime and size");
assert_eq!(tag, "W/\"64-7\"");
assert!(etag_matches(&tag, &tag));
assert!(etag_matches("\"64-7\"", &tag));
assert!(etag_matches("*", &tag));
assert!(etag_matches("\"deadbeef\", W/\"64-7\"", &tag));
assert!(!etag_matches("W/\"64-8\"", &tag));
assert!(!etag_matches("", &tag));
}
#[test]
fn an_etag_needs_both_halves() {
assert!(etag(&Attrs::default()).is_none());
assert!(
etag(&Attrs {
mtime: Some(1),
..Attrs::default()
})
.is_none()
);
}
}