1use std::collections::{HashMap, VecDeque};
11use std::sync::Mutex;
12use std::time::{Duration, Instant};
13
14use bytes::Bytes;
15
16use crate::fs::Entry;
17use crate::sftp::wire::Attrs;
18
19pub const DEFAULT_TTL: Duration = Duration::from_secs(2);
23
24pub const DEFAULT_BODY_CAP: usize = 64 * 1024 * 1024;
26
27struct Listing {
28 entries: HashMap<String, (Attrs, Option<String>)>,
35 fetched: Instant,
36}
37
38#[derive(Clone, PartialEq, Eq, Hash)]
39struct BodyKey {
40 path: String,
41 mtime: u32,
42 size: u64,
43}
44
45struct Bodies {
52 map: HashMap<BodyKey, Bytes>,
53 order: VecDeque<BodyKey>,
54 bytes: usize,
55 cap: usize,
56}
57
58impl Bodies {
59 fn insert(&mut self, key: BodyKey, body: Bytes) {
60 if body.len() > self.cap || self.map.contains_key(&key) {
63 return;
64 }
65 self.bytes += body.len();
66 self.order.push_back(key.clone());
67 self.map.insert(key, body);
68
69 while self.bytes > self.cap {
70 let Some(oldest) = self.order.pop_front() else {
71 break;
72 };
73 if let Some(dropped) = self.map.remove(&oldest) {
74 self.bytes -= dropped.len();
75 }
76 }
77 }
78}
79
80pub struct Cache {
81 listings: Mutex<HashMap<String, Listing>>,
82 bodies: Mutex<Bodies>,
83 ttl: Duration,
84}
85
86impl Cache {
87 pub fn new(ttl: Duration, body_cap: usize) -> Self {
88 Self {
89 listings: Mutex::new(HashMap::new()),
90 bodies: Mutex::new(Bodies {
91 map: HashMap::new(),
92 order: VecDeque::new(),
93 bytes: 0,
94 cap: body_cap,
95 }),
96 ttl,
97 }
98 }
99
100 pub fn has_listing(&self, dir: &str) -> bool {
106 self.listings
107 .lock()
108 .expect("listing cache poisoned")
109 .get(dir)
110 .is_some_and(|l| l.fetched.elapsed() < self.ttl)
111 }
112
113 pub fn attrs_of(&self, dir: &str, name: &str) -> Option<Attrs> {
116 let listings = self.listings.lock().expect("listing cache poisoned");
117 let listing = listings.get(dir)?;
118 if listing.fetched.elapsed() >= self.ttl {
119 return None;
120 }
121 listing.entries.get(name).map(|(attrs, _)| *attrs)
122 }
123
124 pub fn listing_entries(&self, dir: &str) -> Option<Vec<Entry>> {
129 let listings = self.listings.lock().expect("listing cache poisoned");
130 let listing = listings.get(dir)?;
131 if listing.fetched.elapsed() >= self.ttl {
132 return None;
133 }
134 Some(
135 listing
136 .entries
137 .iter()
138 .map(|(name, (attrs, owner))| Entry {
139 name: name.clone(),
140 attrs: *attrs,
141 owner: owner.clone(),
142 })
143 .collect(),
144 )
145 }
146
147 pub fn put_listing(&self, dir: &str, entries: &[Entry]) {
148 let map = entries
149 .iter()
150 .map(|e| (e.name.clone(), (e.attrs, e.owner.clone())))
151 .collect::<HashMap<_, _>>();
152 self.listings
153 .lock()
154 .expect("listing cache poisoned")
155 .insert(
156 dir.to_string(),
157 Listing {
158 entries: map,
159 fetched: Instant::now(),
160 },
161 );
162 }
163
164 pub fn forget_listing(&self, dir: &str) {
170 self.listings
171 .lock()
172 .expect("listing cache poisoned")
173 .remove(dir);
174 }
175
176 pub fn body(&self, path: &str, attrs: &Attrs) -> Option<Bytes> {
177 let key = body_key(path, attrs)?;
178 self.bodies
179 .lock()
180 .expect("body cache poisoned")
181 .map
182 .get(&key)
183 .cloned()
184 }
185
186 pub fn put_body(&self, path: &str, attrs: &Attrs, body: Bytes) {
187 let Some(key) = body_key(path, attrs) else {
190 return;
191 };
192 self.bodies
193 .lock()
194 .expect("body cache poisoned")
195 .insert(key, body);
196 }
197}
198
199impl Default for Cache {
200 fn default() -> Self {
201 Self::new(DEFAULT_TTL, DEFAULT_BODY_CAP)
202 }
203}
204
205fn body_key(path: &str, attrs: &Attrs) -> Option<BodyKey> {
206 Some(BodyKey {
207 path: path.to_string(),
208 mtime: attrs.mtime?,
209 size: attrs.size?,
210 })
211}
212
213pub fn etag(attrs: &Attrs) -> Option<String> {
220 let mtime = attrs.mtime?;
221 let size = attrs.size?;
222 Some(format!("W/\"{mtime:x}-{size:x}\""))
223}
224
225pub fn etag_matches(header: &str, tag: &str) -> bool {
228 let want = normalise(tag);
229 header
230 .split(',')
231 .any(|candidate| candidate.trim() == "*" || normalise(candidate) == want)
232}
233
234fn normalise(tag: &str) -> &str {
235 tag.trim().trim_start_matches("W/").trim_matches('"')
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241
242 fn attrs(mtime: u32, size: u64) -> Attrs {
243 Attrs {
244 mtime: Some(mtime),
245 size: Some(size),
246 ..Attrs::default()
247 }
248 }
249
250 fn entry(name: &str, a: Attrs) -> Entry {
251 Entry {
252 name: name.to_string(),
253 attrs: a,
254 owner: Some("souta".to_string()),
255 }
256 }
257
258 #[test]
259 fn a_fresh_listing_answers_without_the_remote() {
260 let c = Cache::default();
261 c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
262 assert!(c.has_listing("/srv"));
263 assert_eq!(c.attrs_of("/srv", "a.html").and_then(|a| a.size), Some(7));
264 assert!(c.attrs_of("/srv", "missing.html").is_none());
266 }
267
268 #[test]
269 fn a_stale_listing_is_not_used() {
270 let c = Cache::new(Duration::ZERO, DEFAULT_BODY_CAP);
271 c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
272 assert!(!c.has_listing("/srv"));
273 assert!(c.attrs_of("/srv", "a.html").is_none());
274 }
275
276 #[test]
277 fn forgetting_a_listing_forces_a_refetch() {
278 let c = Cache::default();
279 c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
280 c.forget_listing("/srv");
281 assert!(!c.has_listing("/srv"));
282 }
283
284 #[test]
287 fn a_changed_file_misses_the_cache() {
288 let c = Cache::default();
289 let old = attrs(100, 7);
290 c.put_body("/srv/a.html", &old, Bytes::from_static(b"old"));
291 assert_eq!(
292 c.body("/srv/a.html", &old).as_deref(),
293 Some(&b"old"[..]),
294 "the version that was cached is served"
295 );
296
297 let rebuilt_same_size = attrs(200, 7);
298 assert!(
299 c.body("/srv/a.html", &rebuilt_same_size).is_none(),
300 "a new mtime must miss even when the size is unchanged"
301 );
302 let rebuilt_same_mtime = attrs(100, 9);
303 assert!(
304 c.body("/srv/a.html", &rebuilt_same_mtime).is_none(),
305 "a new size must miss even when the mtime is unchanged"
306 );
307 }
308
309 #[test]
310 fn a_body_without_mtime_or_size_is_not_cached() {
311 let c = Cache::default();
312 let bare = Attrs::default();
313 c.put_body("/srv/a.html", &bare, Bytes::from_static(b"x"));
314 assert!(c.body("/srv/a.html", &bare).is_none());
315 }
316
317 #[test]
318 fn the_body_cache_respects_its_budget() {
319 let c = Cache::new(DEFAULT_TTL, 10);
320 for i in 0..5u32 {
321 c.put_body(&format!("/f{i}"), &attrs(i, 4), Bytes::from_static(b"1234"));
322 }
323 let held = c.bodies.lock().expect("body cache").bytes;
324 assert!(held <= 10, "cache holds {held} bytes over a 10 byte budget");
325 assert!(c.body("/f4", &attrs(4, 4)).is_some());
327 assert!(c.body("/f0", &attrs(0, 4)).is_none());
328 }
329
330 #[test]
331 fn a_file_bigger_than_the_budget_is_declined_rather_than_flushing_everything() {
332 let c = Cache::new(DEFAULT_TTL, 4);
333 c.put_body("/small", &attrs(1, 2), Bytes::from_static(b"ab"));
334 c.put_body("/huge", &attrs(2, 99), Bytes::from_static(b"0123456789"));
335 assert!(c.body("/huge", &attrs(2, 99)).is_none());
336 assert!(
337 c.body("/small", &attrs(1, 2)).is_some(),
338 "an oversized insert must not flush the cache"
339 );
340 }
341
342 #[test]
343 fn etags_compare_weakly() {
344 let a = attrs(0x64, 0x7);
345 let tag = etag(&a).expect("attrs carry mtime and size");
346 assert_eq!(tag, "W/\"64-7\"");
347 assert!(etag_matches(&tag, &tag));
348 assert!(etag_matches("\"64-7\"", &tag));
350 assert!(etag_matches("*", &tag));
351 assert!(etag_matches("\"deadbeef\", W/\"64-7\"", &tag));
352 assert!(!etag_matches("W/\"64-8\"", &tag));
353 assert!(!etag_matches("", &tag));
354 }
355
356 #[test]
357 fn an_etag_needs_both_halves() {
358 assert!(etag(&Attrs::default()).is_none());
359 assert!(
360 etag(&Attrs {
361 mtime: Some(1),
362 ..Attrs::default()
363 })
364 .is_none()
365 );
366 }
367}