1use std::fs;
2use std::path::PathBuf;
3use std::sync::atomic::{AtomicU32, Ordering};
4use std::time::{Duration, SystemTime, UNIX_EPOCH};
5
6use serde::{Deserialize, Serialize};
7
8const SWEEP_INTERVAL: u32 = 10;
9const SWEEP_MAX_AGE: Duration = Duration::from_secs(3600);
10
11#[derive(Serialize, Deserialize)]
12struct CacheEntry {
13 cached_at: u64,
14 body: String,
15}
16
17pub struct Cache {
18 dir: PathBuf,
19 enabled: bool,
20 write_count: AtomicU32,
21}
22
23impl Cache {
24 pub fn new(dir: PathBuf, enabled: bool) -> Self {
25 if enabled {
26 if let Err(e) = fs::create_dir_all(&dir) {
27 tracing::debug!("cache dir creation failed, caching disabled: {e}");
28 return Self {
29 dir,
30 enabled: false,
31 write_count: AtomicU32::new(0),
32 };
33 }
34 }
35 Self {
36 dir,
37 enabled,
38 write_count: AtomicU32::new(0),
39 }
40 }
41
42 pub fn get(&self, url: &str, max_age: Duration) -> Option<String> {
43 if !self.enabled {
44 return None;
45 }
46 let path = self.path_for(url);
47 let data = fs::read_to_string(&path).ok()?;
48 let entry: CacheEntry = match serde_json::from_str(&data) {
49 Ok(e) => e,
50 Err(_) => {
51 tracing::debug!("corrupt cache file, removing: {}", path.display());
52 let _ = fs::remove_file(&path);
53 return None;
54 }
55 };
56 let now = now_secs();
57 if entry.cached_at > now {
58 tracing::debug!("cache entry has future timestamp, treating as expired");
59 let _ = fs::remove_file(&path);
60 return None;
61 }
62 if now - entry.cached_at > max_age.as_secs() {
63 return None;
64 }
65 tracing::debug!("cache hit: {url}");
66 Some(entry.body)
67 }
68
69 pub fn set(&self, url: &str, body: &str) {
70 if !self.enabled {
71 return;
72 }
73 let path = self.path_for(url);
74 let entry = CacheEntry {
75 cached_at: now_secs(),
76 body: body.to_string(),
77 };
78 let Ok(data) = serde_json::to_string(&entry) else {
79 return;
80 };
81
82 let tmp = path.with_extension(format!("{}.tmp", std::process::id()));
83 if write_atomic(&tmp, &path, data.as_bytes()).is_err() {
84 tracing::debug!("cache write failed: {}", path.display());
85 }
86
87 let count = self.write_count.fetch_add(1, Ordering::Relaxed);
88 if count > 0 && count.is_multiple_of(SWEEP_INTERVAL) {
89 self.sweep();
90 }
91 }
92
93 pub fn invalidate_pattern(&self, prefix: &str) {
94 if !self.enabled {
95 return;
96 }
97 let entries = match fs::read_dir(&self.dir) {
98 Ok(e) => e,
99 Err(_) => return,
100 };
101 for entry in entries.flatten() {
102 let name = entry.file_name();
103 let name = name.to_string_lossy();
104 if name.contains(prefix) && name.ends_with(".json") {
105 let _ = fs::remove_file(entry.path());
106 }
107 }
108 }
109
110 pub fn clear(&self) {
111 let entries = match fs::read_dir(&self.dir) {
112 Ok(e) => e,
113 Err(_) => return,
114 };
115 for entry in entries.flatten() {
116 let path = entry.path();
117 let ext = path.extension().and_then(|e| e.to_str());
118 if ext == Some("json") || ext == Some("tmp") {
119 let _ = fs::remove_file(path);
120 }
121 }
122 }
123
124 fn path_for(&self, url: &str) -> PathBuf {
125 self.dir.join(format!("{}.json", sanitize_url(url)))
126 }
127
128 fn sweep(&self) {
129 let entries = match fs::read_dir(&self.dir) {
130 Ok(e) => e,
131 Err(_) => return,
132 };
133 let now = now_secs();
134 for entry in entries.flatten() {
135 let path = entry.path();
136 if path.extension().and_then(|e| e.to_str()) != Some("json") {
137 continue;
138 }
139 let data = match fs::read_to_string(&path) {
140 Ok(d) => d,
141 Err(_) => continue,
142 };
143 if let Ok(entry) = serde_json::from_str::<CacheEntry>(&data) {
144 if now.saturating_sub(entry.cached_at) > SWEEP_MAX_AGE.as_secs() {
145 let _ = fs::remove_file(&path);
146 }
147 } else {
148 let _ = fs::remove_file(&path);
149 }
150 }
151 }
152}
153
154fn sanitize_url(url: &str) -> String {
155 let without_scheme = url
156 .strip_prefix("https://")
157 .or_else(|| url.strip_prefix("http://"))
158 .unwrap_or(url);
159 without_scheme
160 .chars()
161 .map(|c| {
162 if c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.' {
163 c
164 } else {
165 '_'
166 }
167 })
168 .collect()
169}
170
171fn now_secs() -> u64 {
172 SystemTime::now()
173 .duration_since(UNIX_EPOCH)
174 .unwrap_or_default()
175 .as_secs()
176}
177
178fn write_atomic(tmp: &std::path::Path, dest: &std::path::Path, data: &[u8]) -> std::io::Result<()> {
179 #[cfg(unix)]
180 {
181 use std::io::Write;
182 use std::os::unix::fs::OpenOptionsExt;
183 let mut f = fs::OpenOptions::new()
184 .write(true)
185 .create(true)
186 .truncate(true)
187 .mode(0o600)
188 .open(tmp)?;
189 f.write_all(data)?;
190 }
191 #[cfg(not(unix))]
192 {
193 fs::write(tmp, data)?;
194 }
195 fs::rename(tmp, dest)
196}
197
198#[cfg(test)]
199mod tests {
200 use super::*;
201
202 fn temp_cache(name: &str) -> (Cache, PathBuf) {
203 let dir =
204 std::env::temp_dir().join(format!("htb-cache-test-{name}-{}", std::process::id()));
205 let _ = fs::remove_dir_all(&dir);
206 let cache = Cache::new(dir.clone(), true);
207 (cache, dir)
208 }
209
210 #[test]
211 fn cache_hit_and_miss() {
212 let (cache, _dir) = temp_cache("hit-miss");
213 let url = "https://labs.hackthebox.com/api/v4/machine/profile/Test";
214
215 assert!(cache.get(url, Duration::from_secs(120)).is_none());
216
217 cache.set(url, r#"{"info":{"id":1}}"#);
218
219 let hit = cache.get(url, Duration::from_secs(120));
220 assert_eq!(hit.as_deref(), Some(r#"{"info":{"id":1}}"#));
221 }
222
223 #[test]
224 fn cache_expiry() {
225 let (cache, _dir) = temp_cache("expiry");
226 let url = "https://labs.hackthebox.com/api/v4/machine/profile/Old";
227
228 let path = cache.path_for(url);
230 let old_entry = CacheEntry {
231 cached_at: now_secs() - 300,
232 body: r#"{"stale":true}"#.to_string(),
233 };
234 fs::write(&path, serde_json::to_string(&old_entry).unwrap()).unwrap();
235
236 assert!(cache.get(url, Duration::from_secs(120)).is_none());
237 }
238
239 #[test]
240 fn cache_future_timestamp_treated_as_expired() {
241 let (cache, _dir) = temp_cache("future-ts");
242 let url = "https://labs.hackthebox.com/api/v4/test";
243
244 let path = cache.path_for(url);
245 let future_entry = CacheEntry {
246 cached_at: now_secs() + 9999,
247 body: r#"{"future":true}"#.to_string(),
248 };
249 fs::write(&path, serde_json::to_string(&future_entry).unwrap()).unwrap();
250
251 assert!(cache.get(url, Duration::from_secs(120)).is_none());
252 assert!(!path.exists());
253 }
254
255 #[test]
256 fn corrupt_file_treated_as_miss() {
257 let (cache, _dir) = temp_cache("corrupt");
258 let url = "https://labs.hackthebox.com/api/v4/corrupt";
259
260 let path = cache.path_for(url);
261 fs::write(&path, "not valid json{{{").unwrap();
262
263 assert!(cache.get(url, Duration::from_secs(120)).is_none());
264 assert!(!path.exists());
265 }
266
267 #[test]
268 fn invalidate_pattern() {
269 let (cache, _dir) = temp_cache("invalidate");
270 let url1 = "https://labs.hackthebox.com/api/v4/machine/profile/A";
271 let url2 = "https://labs.hackthebox.com/api/v4/machine/profile/B";
272 let url3 = "https://labs.hackthebox.com/api/v4/challenge/info/C";
273
274 cache.set(url1, "a");
275 cache.set(url2, "b");
276 cache.set(url3, "c");
277
278 cache.invalidate_pattern("api_v4_machine");
279
280 assert!(
281 cache.get(url1, Duration::from_secs(300)).is_none(),
282 "url1 should be invalidated"
283 );
284 assert!(
285 cache.get(url2, Duration::from_secs(300)).is_none(),
286 "url2 should be invalidated"
287 );
288 assert!(
289 cache.get(url3, Duration::from_secs(300)).is_some(),
290 "url3 should survive"
291 );
292 }
293
294 #[test]
295 fn clear_removes_all() {
296 let (cache, _dir) = temp_cache("clear");
297 cache.set("https://example.com/api/v4/a", "1");
298 cache.set("https://example.com/api/v4/b", "2");
299
300 cache.clear();
301
302 assert!(cache
303 .get("https://example.com/api/v4/a", Duration::from_secs(300))
304 .is_none());
305 assert!(cache
306 .get("https://example.com/api/v4/b", Duration::from_secs(300))
307 .is_none());
308 }
309
310 #[test]
311 fn sanitize_url_separates_labs_and_ctf() {
312 let labs = sanitize_url("https://labs.hackthebox.com/api/users/profile");
313 let ctf = sanitize_url("https://ctf.hackthebox.com/api/users/profile");
314 assert_ne!(labs, ctf);
315 assert!(labs.starts_with("labs.hackthebox.com_"));
316 assert!(ctf.starts_with("ctf.hackthebox.com_"));
317 }
318
319 #[test]
320 fn disabled_cache_is_noop() {
321 let dir = std::env::temp_dir().join(format!("htb-cache-disabled-{}", std::process::id()));
322 let cache = Cache::new(dir, false);
323 let url = "https://example.com/api/v4/test";
324
325 cache.set(url, "data");
326 assert!(cache.get(url, Duration::from_secs(300)).is_none());
327 }
328
329 #[test]
330 fn sanitize_url_produces_safe_filenames() {
331 assert_eq!(
332 sanitize_url("https://labs.hackthebox.com/api/v4/machine/profile/Bedside"),
333 "labs.hackthebox.com_api_v4_machine_profile_Bedside"
334 );
335 assert_eq!(
336 sanitize_url("https://labs.hackthebox.com/api/v5/machines?per_page=100&page=1"),
337 "labs.hackthebox.com_api_v5_machines_per_page_100_page_1"
338 );
339 }
340}