Skip to main content

aqua_registry/
compiled.rs

1use 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 &registry.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
225/// Hashes the canonical package ID with FNV-1a 64-bit to keep compiled cache
226/// filenames deterministic. The sanitized ID prefix is only for readability.
227fn 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}