aqua_registry/
compiled.rs1use crate::codec::{decode_package_rkyv, encode_package_rkyv};
2use crate::types::{AquaPackage, RegistryYaml};
3use crate::{AquaRegistryError, Result};
4use rkyv::rancor::Error as RkyvError;
5use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
6use std::collections::{HashMap, HashSet};
7use std::fs;
8use std::path::{Path, PathBuf};
9
10const INDEX_FILE: &str = "index.rkyv";
11const PACKAGES_DIR: &str = "packages";
12
13#[derive(Debug, Clone)]
14pub struct CompiledRegistry {
15 root: PathBuf,
16 index: CompiledRegistryIndex,
17}
18
19#[derive(Debug, Clone)]
20pub struct ParsedRegistry {
21 packages: HashMap<String, AquaPackage>,
22 aliases: HashMap<String, String>,
23}
24
25#[derive(Debug, Clone, Archive, RkyvDeserialize, RkyvSerialize)]
26struct CompiledRegistryIndex {
27 packages: HashMap<String, String>,
28 aliases: HashMap<String, String>,
29}
30
31impl CompiledRegistry {
32 pub fn load(root: impl AsRef<Path>) -> Result<Self> {
33 let root = root.as_ref().to_path_buf();
34 let index = read_index(&root)?;
35 validate_package_files(&root, &index)?;
36 Ok(Self { root, index })
37 }
38
39 pub fn package(&self, package_id: &str) -> Result<AquaPackage> {
40 let resolved_id = self
41 .index
42 .aliases
43 .get(package_id)
44 .map_or(package_id, String::as_str);
45 let filename = self
46 .index
47 .packages
48 .get(resolved_id)
49 .ok_or_else(|| AquaRegistryError::PackageNotFound(package_id.to_string()))?;
50 let path = self.root.join(PACKAGES_DIR).join(filename);
51 let bytes = fs::read(&path)?;
52 decode_package_rkyv(resolved_id, &bytes)
53 }
54}
55
56impl ParsedRegistry {
57 pub fn parse_yaml(source: &str) -> Result<Self> {
58 let registry_yaml = serde_yaml::from_str::<RegistryYaml>(source)?;
59 Self::from_registry_yaml(registry_yaml)
60 }
61
62 pub fn package(&self, package_id: &str) -> Result<AquaPackage> {
63 let resolved_id = self
64 .aliases
65 .get(package_id)
66 .map_or(package_id, String::as_str);
67 self.packages
68 .get(resolved_id)
69 .cloned()
70 .ok_or_else(|| AquaRegistryError::PackageNotFound(package_id.to_string()))
71 }
72
73 pub fn write_compiled_cache(&self, root: impl AsRef<Path>) -> Result<CompiledRegistry> {
74 let root = root.as_ref().to_path_buf();
75 let index = write_compiled_index(self, &root)?;
76 Ok(CompiledRegistry { root, index })
77 }
78
79 fn from_registry_yaml(registry_yaml: RegistryYaml) -> Result<Self> {
80 let package_entries = registry_yaml
81 .packages
82 .into_iter()
83 .filter_map(|row| canonical_package_id(&row.package).map(|id| (id, row)))
84 .collect::<Vec<_>>();
85
86 if package_entries.is_empty() {
87 return Err(AquaRegistryError::RegistryNotAvailable(
88 "aqua registry contains no packages".to_string(),
89 ));
90 }
91
92 let canonical_ids = package_entries
93 .iter()
94 .map(|(id, _)| id.clone())
95 .collect::<HashSet<_>>();
96 let mut packages = HashMap::new();
97 let mut aliases = HashMap::new();
98
99 for (id, row) in package_entries {
100 for alias in &row.aliases {
101 if alias != &id && !canonical_ids.contains(alias.as_str()) {
102 aliases.insert(alias.clone(), id.clone());
103 }
104 }
105 packages.insert(id, row.package);
106 }
107
108 Ok(Self { packages, aliases })
109 }
110}
111
112fn read_index(root: &Path) -> Result<CompiledRegistryIndex> {
113 let path = root.join(INDEX_FILE);
114 let bytes = fs::read(&path)?;
115 rkyv::from_bytes::<CompiledRegistryIndex, RkyvError>(&bytes).map_err(|err| {
116 AquaRegistryError::RegistryNotAvailable(format!(
117 "failed to decode compiled aqua registry index {} from rkyv: {err}",
118 path.display()
119 ))
120 })
121}
122
123fn validate_package_files(root: &Path, index: &CompiledRegistryIndex) -> Result<()> {
124 let packages_dir = root.join(PACKAGES_DIR);
125 for filename in index.packages.values() {
126 let path = packages_dir.join(filename);
127 if !path.is_file() {
128 return Err(AquaRegistryError::RegistryNotAvailable(format!(
129 "compiled aqua registry package file is missing: {}",
130 path.display()
131 )));
132 }
133 }
134 Ok(())
135}
136
137fn write_index(root: &Path, index: &CompiledRegistryIndex) -> Result<()> {
138 let path = root.join(INDEX_FILE);
139 let bytes = rkyv::to_bytes::<RkyvError>(index)
140 .map(|bytes| bytes.to_vec())
141 .map_err(|err| {
142 AquaRegistryError::RegistryNotAvailable(format!(
143 "failed to encode compiled aqua registry index {} as rkyv: {err}",
144 path.display()
145 ))
146 })?;
147 fs::write(path, bytes)?;
148 Ok(())
149}
150
151fn write_compiled_index(registry: &ParsedRegistry, root: &Path) -> Result<CompiledRegistryIndex> {
152 let packages_dir = root.join(PACKAGES_DIR);
153 fs::create_dir_all(&packages_dir)?;
154
155 let mut used_filenames = HashSet::new();
156 let mut packages = HashMap::new();
157
158 for (id, package) in ®istry.packages {
159 let filename = package_filename(id, &mut used_filenames);
160 let path = packages_dir.join(&filename);
161 let content = encode_package_rkyv(package)?;
162 fs::write(path, content)?;
163 packages.insert(id.clone(), filename);
164 }
165
166 let index = CompiledRegistryIndex {
167 packages,
168 aliases: registry.aliases.clone(),
169 };
170 write_index(root, &index)?;
171 Ok(index)
172}
173
174fn canonical_package_id(package: &AquaPackage) -> Option<String> {
175 package
176 .name
177 .clone()
178 .or_else(|| {
179 if package.repo_owner.is_empty() || package.repo_name.is_empty() {
180 None
181 } else {
182 Some(format!("{}/{}", package.repo_owner, package.repo_name))
183 }
184 })
185 .or_else(|| package.path.clone())
186}
187
188fn package_filename(id: &str, used_filenames: &mut HashSet<String>) -> String {
189 let stem = package_filename_stem(id);
190 let mut filename = format!("{stem}.rkyv");
191 let mut suffix = 2;
192 while !used_filenames.insert(filename.clone()) {
193 filename = format!("{stem}-{suffix}.rkyv");
194 suffix += 1;
195 }
196 filename
197}
198
199fn package_filename_stem(id: &str) -> String {
200 let sanitized = sanitize_filename_prefix(id);
201 let hash = fnv1a64(id);
202 format!("{sanitized}-{hash:016x}")
203}
204
205fn sanitize_filename_prefix(id: &str) -> String {
206 let mut prefix = String::new();
207 for byte in id.bytes() {
208 let c = byte as char;
209 if c.is_ascii_alphanumeric() {
210 prefix.push(c.to_ascii_lowercase());
211 } else {
212 prefix.push('_');
213 }
214 if prefix.len() >= 80 {
215 break;
216 }
217 }
218 if prefix.is_empty() {
219 "package".to_string()
220 } else {
221 prefix
222 }
223}
224
225fn fnv1a64(value: &str) -> u64 {
228 let mut hash = 0xcbf29ce484222325u64;
229 for byte in value.as_bytes() {
230 hash ^= u64::from(*byte);
231 hash = hash.wrapping_mul(0x100000001b3);
232 }
233 hash
234}
235
236#[cfg(test)]
237mod tests {
238 use super::*;
239 use crate::AquaPackageType;
240 use std::time::{SystemTime, UNIX_EPOCH};
241
242 #[test]
243 fn compiles_flat_registry_cache_and_resolves_aliases() {
244 let root = temp_cache_dir("compiled-aqua-registry");
245 let source = r#"
246packages:
247 - type: http
248 name: example/canonical-tool
249 repo_owner: example
250 repo_name: tool
251 url: https://example.com/tool
252 aliases:
253 - name: example/tool-alias
254 version_overrides:
255 - aliases:
256 - name: example/nested-alias
257"#;
258
259 let registry = compile_registry(source, &root);
260 let package = registry.package("example/tool-alias").unwrap();
261
262 assert_eq!(package.name.as_deref(), Some("example/canonical-tool"));
263 assert_eq!(package.repo_owner, "example");
264 assert_eq!(package.repo_name, "tool");
265 assert!(registry.package("example/canonical-tool").is_ok());
266 assert!(matches!(
267 registry.package("example/tool"),
268 Err(AquaRegistryError::PackageNotFound(_))
269 ));
270 assert!(matches!(
271 registry.package("example/nested-alias"),
272 Err(AquaRegistryError::PackageNotFound(_))
273 ));
274 assert!(root.join(INDEX_FILE).exists());
275
276 let packages_dir = root.join(PACKAGES_DIR);
277 let files = fs::read_dir(&packages_dir)
278 .unwrap()
279 .collect::<std::result::Result<Vec<_>, _>>()
280 .unwrap();
281 assert_eq!(files.len(), 1);
282 assert!(files[0].file_type().unwrap().is_file());
283
284 fs::remove_dir_all(root).unwrap();
285 }
286
287 #[test]
288 fn parsed_registry_resolves_packages_before_cache_is_written() {
289 let source = r#"
290packages:
291 - type: http
292 name: example/canonical-tool
293 url: https://example.com/tool
294 aliases:
295 - name: example/tool-alias
296"#;
297
298 let registry = ParsedRegistry::parse_yaml(source).unwrap();
299 let package = registry.package("example/tool-alias").unwrap();
300
301 assert_eq!(package.name.as_deref(), Some("example/canonical-tool"));
302 assert!(matches!(
303 registry.package("example/missing"),
304 Err(AquaRegistryError::PackageNotFound(_))
305 ));
306 }
307
308 #[test]
309 fn loads_compiled_registry_without_reparsing_yaml() {
310 let root = temp_cache_dir("compiled-aqua-registry-load");
311 let source = r#"
312packages:
313 - type: http
314 name: example/named-tool
315 url: https://example.com/tool
316"#;
317
318 compile_registry(source, &root);
319 let registry = CompiledRegistry::load(&root).unwrap();
320 let package = registry.package("example/named-tool").unwrap();
321
322 assert_eq!(package.name.as_deref(), Some("example/named-tool"));
323
324 fs::remove_dir_all(root).unwrap();
325 }
326
327 #[test]
328 fn compiled_registry_preserves_omitted_and_explicit_package_types() {
329 let root = temp_cache_dir("compiled-aqua-registry-package-types");
330 let source = r#"
331packages:
332 - name: example/default-type
333 - name: example/explicit-type
334 type: github_release
335"#;
336
337 compile_registry(source, &root);
338 let registry = CompiledRegistry::load(&root).unwrap();
339 let default_type = registry.package("example/default-type").unwrap();
340 let explicit_type = registry.package("example/explicit-type").unwrap();
341
342 assert_eq!(default_type.r#type, None);
343 assert_eq!(default_type.package_type(), AquaPackageType::GithubRelease);
344 assert_eq!(explicit_type.r#type, Some(AquaPackageType::GithubRelease));
345 assert_eq!(explicit_type.package_type(), AquaPackageType::GithubRelease);
346
347 fs::remove_dir_all(root).unwrap();
348 }
349
350 #[test]
351 fn load_rejects_missing_package_blob() {
352 let root = temp_cache_dir("compiled-aqua-registry-missing-package");
353 let source = r#"
354packages:
355 - type: http
356 name: example/missing-package
357 url: https://example.com/tool
358"#;
359
360 compile_registry(source, &root);
361 let packages_dir = root.join(PACKAGES_DIR);
362 let package_file = fs::read_dir(&packages_dir)
363 .unwrap()
364 .next()
365 .unwrap()
366 .unwrap()
367 .path();
368 fs::remove_file(package_file).unwrap();
369
370 let err = CompiledRegistry::load(&root).unwrap_err();
371 assert!(matches!(err, AquaRegistryError::RegistryNotAvailable(_)));
372
373 fs::remove_dir_all(root).unwrap();
374 }
375
376 fn compile_registry(source: &str, root: &Path) -> CompiledRegistry {
377 ParsedRegistry::parse_yaml(source)
378 .unwrap()
379 .write_compiled_cache(root)
380 .unwrap()
381 }
382
383 fn temp_cache_dir(name: &str) -> PathBuf {
384 let nanos = SystemTime::now()
385 .duration_since(UNIX_EPOCH)
386 .unwrap()
387 .as_nanos();
388 std::env::temp_dir().join(format!("{name}-{nanos}"))
389 }
390}