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>,
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).copied()
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)| Entry {
139 name: name.clone(),
140 attrs: *attrs,
141 })
142 .collect(),
143 )
144 }
145
146 pub fn put_listing(&self, dir: &str, entries: &[Entry]) {
147 let map = entries
148 .iter()
149 .map(|e| (e.name.clone(), e.attrs))
150 .collect::<HashMap<_, _>>();
151 self.listings
152 .lock()
153 .expect("listing cache poisoned")
154 .insert(
155 dir.to_string(),
156 Listing {
157 entries: map,
158 fetched: Instant::now(),
159 },
160 );
161 }
162
163 pub fn forget_listing(&self, dir: &str) {
169 self.listings
170 .lock()
171 .expect("listing cache poisoned")
172 .remove(dir);
173 }
174
175 pub fn body(&self, path: &str, attrs: &Attrs) -> Option<Bytes> {
176 let key = body_key(path, attrs)?;
177 self.bodies
178 .lock()
179 .expect("body cache poisoned")
180 .map
181 .get(&key)
182 .cloned()
183 }
184
185 pub fn put_body(&self, path: &str, attrs: &Attrs, body: Bytes) {
186 let Some(key) = body_key(path, attrs) else {
189 return;
190 };
191 self.bodies
192 .lock()
193 .expect("body cache poisoned")
194 .insert(key, body);
195 }
196}
197
198impl Default for Cache {
199 fn default() -> Self {
200 Self::new(DEFAULT_TTL, DEFAULT_BODY_CAP)
201 }
202}
203
204fn body_key(path: &str, attrs: &Attrs) -> Option<BodyKey> {
205 Some(BodyKey {
206 path: path.to_string(),
207 mtime: attrs.mtime?,
208 size: attrs.size?,
209 })
210}
211
212pub fn etag(attrs: &Attrs) -> Option<String> {
219 let mtime = attrs.mtime?;
220 let size = attrs.size?;
221 Some(format!("W/\"{mtime:x}-{size:x}\""))
222}
223
224pub fn etag_matches(header: &str, tag: &str) -> bool {
227 let want = normalise(tag);
228 header
229 .split(',')
230 .any(|candidate| candidate.trim() == "*" || normalise(candidate) == want)
231}
232
233fn normalise(tag: &str) -> &str {
234 tag.trim().trim_start_matches("W/").trim_matches('"')
235}
236
237#[cfg(test)]
238mod tests {
239 use super::*;
240
241 fn attrs(mtime: u32, size: u64) -> Attrs {
242 Attrs {
243 mtime: Some(mtime),
244 size: Some(size),
245 ..Attrs::default()
246 }
247 }
248
249 fn entry(name: &str, a: Attrs) -> Entry {
250 Entry {
251 name: name.to_string(),
252 attrs: a,
253 }
254 }
255
256 #[test]
257 fn a_fresh_listing_answers_without_the_remote() {
258 let c = Cache::default();
259 c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
260 assert!(c.has_listing("/srv"));
261 assert_eq!(c.attrs_of("/srv", "a.html").and_then(|a| a.size), Some(7));
262 assert!(c.attrs_of("/srv", "missing.html").is_none());
264 }
265
266 #[test]
267 fn a_stale_listing_is_not_used() {
268 let c = Cache::new(Duration::ZERO, DEFAULT_BODY_CAP);
269 c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
270 assert!(!c.has_listing("/srv"));
271 assert!(c.attrs_of("/srv", "a.html").is_none());
272 }
273
274 #[test]
275 fn forgetting_a_listing_forces_a_refetch() {
276 let c = Cache::default();
277 c.put_listing("/srv", &[entry("a.html", attrs(100, 7))]);
278 c.forget_listing("/srv");
279 assert!(!c.has_listing("/srv"));
280 }
281
282 #[test]
285 fn a_changed_file_misses_the_cache() {
286 let c = Cache::default();
287 let old = attrs(100, 7);
288 c.put_body("/srv/a.html", &old, Bytes::from_static(b"old"));
289 assert_eq!(
290 c.body("/srv/a.html", &old).as_deref(),
291 Some(&b"old"[..]),
292 "the version that was cached is served"
293 );
294
295 let rebuilt_same_size = attrs(200, 7);
296 assert!(
297 c.body("/srv/a.html", &rebuilt_same_size).is_none(),
298 "a new mtime must miss even when the size is unchanged"
299 );
300 let rebuilt_same_mtime = attrs(100, 9);
301 assert!(
302 c.body("/srv/a.html", &rebuilt_same_mtime).is_none(),
303 "a new size must miss even when the mtime is unchanged"
304 );
305 }
306
307 #[test]
308 fn a_body_without_mtime_or_size_is_not_cached() {
309 let c = Cache::default();
310 let bare = Attrs::default();
311 c.put_body("/srv/a.html", &bare, Bytes::from_static(b"x"));
312 assert!(c.body("/srv/a.html", &bare).is_none());
313 }
314
315 #[test]
316 fn the_body_cache_respects_its_budget() {
317 let c = Cache::new(DEFAULT_TTL, 10);
318 for i in 0..5u32 {
319 c.put_body(&format!("/f{i}"), &attrs(i, 4), Bytes::from_static(b"1234"));
320 }
321 let held = c.bodies.lock().expect("body cache").bytes;
322 assert!(held <= 10, "cache holds {held} bytes over a 10 byte budget");
323 assert!(c.body("/f4", &attrs(4, 4)).is_some());
325 assert!(c.body("/f0", &attrs(0, 4)).is_none());
326 }
327
328 #[test]
329 fn a_file_bigger_than_the_budget_is_declined_rather_than_flushing_everything() {
330 let c = Cache::new(DEFAULT_TTL, 4);
331 c.put_body("/small", &attrs(1, 2), Bytes::from_static(b"ab"));
332 c.put_body("/huge", &attrs(2, 99), Bytes::from_static(b"0123456789"));
333 assert!(c.body("/huge", &attrs(2, 99)).is_none());
334 assert!(
335 c.body("/small", &attrs(1, 2)).is_some(),
336 "an oversized insert must not flush the cache"
337 );
338 }
339
340 #[test]
341 fn etags_compare_weakly() {
342 let a = attrs(0x64, 0x7);
343 let tag = etag(&a).expect("attrs carry mtime and size");
344 assert_eq!(tag, "W/\"64-7\"");
345 assert!(etag_matches(&tag, &tag));
346 assert!(etag_matches("\"64-7\"", &tag));
348 assert!(etag_matches("*", &tag));
349 assert!(etag_matches("\"deadbeef\", W/\"64-7\"", &tag));
350 assert!(!etag_matches("W/\"64-8\"", &tag));
351 assert!(!etag_matches("", &tag));
352 }
353
354 #[test]
355 fn an_etag_needs_both_halves() {
356 assert!(etag(&Attrs::default()).is_none());
357 assert!(
358 etag(&Attrs {
359 mtime: Some(1),
360 ..Attrs::default()
361 })
362 .is_none()
363 );
364 }
365}