1use crate::{AquaRegistryError, CompiledRegistry, ParsedRegistry, Result};
2use blake3::Hasher as Blake3Hasher;
3use siphasher::sip::SipHasher;
4use std::fs;
5use std::hash::{Hash, Hasher};
6use std::io::Write;
7use std::path::{Path, PathBuf};
8use std::time::{Duration, SystemTime};
9
10const COMPILED_REGISTRY_CACHE_VERSION: &str = "v9";
11
12#[derive(Debug, Clone)]
13pub struct RegistryCache {
14 root: PathBuf,
15}
16
17impl RegistryCache {
18 pub fn new(root: impl Into<PathBuf>) -> Self {
19 Self { root: root.into() }
20 }
21
22 pub fn source_path(&self, registry_url: &str) -> PathBuf {
23 self.root
24 .join("sources")
25 .join(format!("{}.yaml", registry_url_hash(registry_url)))
26 }
27
28 pub fn read_source(&self, registry_url: &str) -> Result<Option<String>> {
29 let path = self.source_path(registry_url);
30 read_optional_to_string(&path)
31 }
32
33 pub fn read_fresh_source(
34 &self,
35 registry_url: &str,
36 max_age: Duration,
37 ) -> Result<Option<String>> {
38 let path = self.source_path(registry_url);
39 if !path_is_fresh(&path, max_age)? {
40 return Ok(None);
41 }
42 read_optional_to_string(&path)
43 }
44
45 pub fn write_source(&self, registry_url: &str, source: &str) -> Result<()> {
46 let path = self.source_path(registry_url);
47 let Some(parent) = path.parent() else {
48 return Err(AquaRegistryError::RegistryNotAvailable(format!(
49 "cached aqua registry source path has no parent: {}",
50 path.display()
51 )));
52 };
53 fs::create_dir_all(parent)?;
54
55 let mut tmp = tempfile::NamedTempFile::with_prefix_in("registry-source.", parent)?;
56 tmp.write_all(source.as_bytes())?;
57 tmp.persist(&path).map_err(|err| err.error)?;
58 Ok(())
59 }
60
61 pub fn source_hash(source: &str) -> String {
62 source_hash(source)
63 }
64
65 pub fn compiled_dir(&self, registry_url: &str, source_hash: &str) -> PathBuf {
66 self.root
67 .join("compiled")
68 .join(registry_url_hash(registry_url))
69 .join(COMPILED_REGISTRY_CACHE_VERSION)
70 .join(source_hash)
71 }
72
73 pub fn load_compiled(&self, registry_url: &str, source_hash: &str) -> Result<CompiledRegistry> {
74 CompiledRegistry::load(self.compiled_dir(registry_url, source_hash))
75 }
76
77 pub fn write_compiled(
78 &self,
79 registry_url: &str,
80 source_hash: &str,
81 registry: &ParsedRegistry,
82 ) -> Result<CompiledRegistry> {
83 let compiled_dir = self.compiled_dir(registry_url, source_hash);
84 if let Ok(existing) = CompiledRegistry::load(&compiled_dir) {
85 self.prune_stale_compiled(registry_url, source_hash);
86 return Ok(existing);
87 }
88
89 let Some(parent) = compiled_dir.parent() else {
90 return Err(AquaRegistryError::RegistryNotAvailable(format!(
91 "compiled aqua registry cache path has no parent: {}",
92 compiled_dir.display()
93 )));
94 };
95 fs::create_dir_all(parent)?;
96
97 let tmp_dir = tempfile::Builder::new()
98 .prefix(&format!("{source_hash}.tmp-"))
99 .tempdir_in(parent)?;
100 let tmp_path = tmp_dir.path().to_path_buf();
101
102 registry.write_compiled_cache(&tmp_path)?;
103 let tmp_path = tmp_dir.keep();
104
105 if let Ok(existing) = CompiledRegistry::load(&compiled_dir) {
106 cleanup_tmp_dir_for_existing_compiled_cache(&tmp_path, &compiled_dir)?;
107 self.prune_stale_compiled(registry_url, source_hash);
108 return Ok(existing);
109 }
110
111 if compiled_dir.exists() {
112 remove_dir_all_if_exists(&compiled_dir)?;
113 }
114
115 if let Err(err) = fs::rename(&tmp_path, &compiled_dir) {
116 if let Ok(existing) = CompiledRegistry::load(&compiled_dir) {
117 cleanup_tmp_dir_for_existing_compiled_cache(&tmp_path, &compiled_dir)?;
118 self.prune_stale_compiled(registry_url, source_hash);
119 return Ok(existing);
120 }
121 let _ = remove_dir_all_if_exists(&tmp_path);
122 return Err(err.into());
123 }
124
125 let compiled = CompiledRegistry::load(&compiled_dir)?;
126 self.prune_stale_compiled(registry_url, source_hash);
127 Ok(compiled)
128 }
129
130 pub fn prune_stale_compiled(&self, registry_url: &str, source_hash: &str) {
131 let current_dir = self.compiled_dir(registry_url, source_hash);
132 prune_stale_compiled_registries(¤t_dir);
133 }
134}
135
136fn registry_url_hash(registry_url: &str) -> String {
137 hash_to_str(®istry_url)
138}
139
140fn source_hash(source: &str) -> String {
141 let mut hasher = Blake3Hasher::new();
142 hasher.update(source.as_bytes());
143 hasher.finalize().to_hex().to_string()
144}
145
146fn hash_to_str<T: Hash>(t: &T) -> String {
147 let mut s = SipHasher::new();
148 t.hash(&mut s);
149 format!("{:x}", s.finish())
150}
151
152fn read_optional_to_string(path: &Path) -> Result<Option<String>> {
153 match fs::read_to_string(path) {
154 Ok(source) => Ok(Some(source)),
155 Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
156 Err(err) => Err(err.into()),
157 }
158}
159
160fn path_is_fresh(path: &Path, max_age: Duration) -> Result<bool> {
161 let Some(age) = path_age(path)? else {
162 return Ok(false);
163 };
164 Ok(age < max_age)
165}
166
167fn path_age(path: &Path) -> Result<Option<Duration>> {
168 let metadata = match fs::metadata(path) {
169 Ok(metadata) => metadata,
170 Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
171 Err(err) => return Err(err.into()),
172 };
173 let modified = metadata.modified()?;
174 Ok(Some(
175 SystemTime::now()
176 .duration_since(modified)
177 .unwrap_or_default(),
178 ))
179}
180
181fn prune_stale_compiled_registries(current_dir: &Path) {
182 let Some(parent) = current_dir.parent() else {
183 return;
184 };
185 let Ok(entries) = fs::read_dir(parent) else {
186 return;
187 };
188
189 for entry in entries.flatten() {
190 let path = entry.path();
191 if path == current_dir {
192 continue;
193 }
194 if entry.file_type().is_ok_and(|file_type| file_type.is_dir())
195 && is_compiled_source_hash_dir(&path)
196 && let Err(err) = fs::remove_dir_all(&path)
197 {
198 log::debug!(
199 "failed to prune stale compiled aqua registry cache {}: {err}",
200 path.display()
201 );
202 }
203 }
204}
205
206fn is_compiled_source_hash_dir(path: &Path) -> bool {
207 path.file_name()
208 .and_then(|name| name.to_str())
209 .is_some_and(|name| name.len() == 64 && name.bytes().all(|b| b.is_ascii_hexdigit()))
210}
211
212fn cleanup_tmp_dir_for_existing_compiled_cache(tmp_dir: &Path, compiled_dir: &Path) -> Result<()> {
213 match fs::remove_dir_all(tmp_dir) {
214 Ok(()) => Ok(()),
215 Err(err)
216 if err.kind() == std::io::ErrorKind::NotFound
217 && CompiledRegistry::load(compiled_dir).is_ok() =>
218 {
219 Ok(())
220 }
221 Err(err) => Err(err.into()),
222 }
223}
224
225fn remove_dir_all_if_exists(path: &Path) -> std::io::Result<()> {
226 match fs::remove_dir_all(path) {
227 Ok(()) => Ok(()),
228 Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
229 Err(err) => Err(err),
230 }
231}
232
233#[cfg(test)]
234mod tests {
235 use super::*;
236
237 fn registry_source(package_id: &str) -> String {
238 format!("packages:\n - name: {package_id}\n url: https://example.com/tool\n")
239 }
240
241 #[test]
242 fn source_cache_reads_fresh_sources_and_skips_stale_sources() {
243 let temp = tempfile::tempdir().unwrap();
244 let cache = RegistryCache::new(temp.path());
245 let registry_url = "https://example.com/aqua-registry";
246
247 cache.write_source(registry_url, "packages: []").unwrap();
248
249 assert_eq!(
250 cache
251 .read_fresh_source(registry_url, Duration::from_secs(60))
252 .unwrap()
253 .as_deref(),
254 Some("packages: []")
255 );
256 assert!(
257 cache
258 .read_fresh_source(registry_url, Duration::ZERO)
259 .unwrap()
260 .is_none()
261 );
262 }
263
264 #[test]
265 fn source_cache_writes_atomically_and_overwrites_existing_source() {
266 let temp = tempfile::tempdir().unwrap();
267 let cache = RegistryCache::new(temp.path());
268 let registry_url = "https://example.com/aqua-registry";
269
270 cache.write_source(registry_url, "first").unwrap();
271 cache.write_source(registry_url, "second").unwrap();
272
273 assert_eq!(
274 cache.read_source(registry_url).unwrap().as_deref(),
275 Some("second")
276 );
277 assert!(cache.source_path(registry_url).is_file());
278 }
279
280 #[test]
281 fn compiled_cache_is_scoped_by_registry_url() {
282 let cache = RegistryCache::new("/cache");
283 let source_hash = RegistryCache::source_hash("packages: []");
284 let first = cache.compiled_dir("https://example.com/one", &source_hash);
285 let second = cache.compiled_dir("https://example.com/two", &source_hash);
286
287 assert_ne!(first.parent(), second.parent());
288 assert_eq!(
289 first.file_name().and_then(|name| name.to_str()),
290 Some(source_hash.as_str())
291 );
292 }
293
294 #[test]
295 fn compiled_cache_writes_loads_and_prunes_stale_source_hash_siblings() {
296 let temp = tempfile::tempdir().unwrap();
297 let cache = RegistryCache::new(temp.path());
298 let registry_url = "https://example.com/aqua-registry";
299 let first_source = registry_source("example/first");
300 let second_source = registry_source("example/second");
301 let first_hash = RegistryCache::source_hash(&first_source);
302 let second_hash = RegistryCache::source_hash(&second_source);
303 let first_registry = ParsedRegistry::parse_yaml(&first_source).unwrap();
304 let second_registry = ParsedRegistry::parse_yaml(&second_source).unwrap();
305
306 cache
307 .write_compiled(registry_url, &first_hash, &first_registry)
308 .unwrap();
309 let first_dir = cache.compiled_dir(registry_url, &first_hash);
310 assert!(first_dir.is_dir());
311
312 cache
313 .write_compiled(registry_url, &second_hash, &second_registry)
314 .unwrap();
315 let second_dir = cache.compiled_dir(registry_url, &second_hash);
316 let loaded = cache.load_compiled(registry_url, &second_hash).unwrap();
317
318 assert!(second_dir.is_dir());
319 assert!(!first_dir.exists());
320 assert!(loaded.package("example/second").is_ok());
321 }
322
323 #[test]
324 fn compiled_cache_prune_skips_temp_directories() {
325 let temp = tempfile::tempdir().unwrap();
326 let cache = RegistryCache::new(temp.path());
327 let registry_url = "https://example.com/aqua-registry";
328 let current_hash = RegistryCache::source_hash(®istry_source("example/current"));
329 let stale_hash = RegistryCache::source_hash(®istry_source("example/stale"));
330 let current_dir = cache.compiled_dir(registry_url, ¤t_hash);
331 let stale_dir = cache.compiled_dir(registry_url, &stale_hash);
332 let tmp_dir = current_dir
333 .parent()
334 .unwrap()
335 .join(format!("{current_hash}.tmp-in-progress"));
336
337 fs::create_dir_all(¤t_dir).unwrap();
338 fs::create_dir_all(&stale_dir).unwrap();
339 fs::create_dir_all(&tmp_dir).unwrap();
340
341 cache.prune_stale_compiled(registry_url, ¤t_hash);
342
343 assert!(current_dir.is_dir());
344 assert!(!stale_dir.exists());
345 assert!(tmp_dir.is_dir());
346 }
347
348 #[test]
349 fn compiled_temp_cleanup_treats_missing_temp_as_success_when_final_cache_exists() {
350 let temp = tempfile::tempdir().unwrap();
351 let cache = RegistryCache::new(temp.path());
352 let registry_url = "https://example.com/aqua-registry";
353 let source = registry_source("example/tool");
354 let source_hash = RegistryCache::source_hash(&source);
355 let registry = ParsedRegistry::parse_yaml(&source).unwrap();
356 let compiled_dir = cache.compiled_dir(registry_url, &source_hash);
357 let missing_tmp_dir = compiled_dir.with_file_name(format!("{source_hash}.tmp-missing"));
358
359 registry.write_compiled_cache(&compiled_dir).unwrap();
360
361 cleanup_tmp_dir_for_existing_compiled_cache(&missing_tmp_dir, &compiled_dir).unwrap();
362 }
363
364 #[test]
365 fn registry_url_hash_matches_existing_cache_layout() {
366 assert_eq!(registry_url_hash("foo"), "e1b19adfb2e348a2");
367 }
368}