Skip to main content

aube_resolver/
primer.rs

1use aube_manifest::BundledDependencies;
2use aube_registry::{Attestations, Dist, NpmUser, Packument, PeerDepMeta, VersionMetadata};
3use std::collections::BTreeMap;
4use std::io::{Cursor, Write};
5use std::path::{Path, PathBuf};
6use std::sync::OnceLock;
7use std::time::Duration;
8
9#[path = "primer_schema.rs"]
10mod primer_schema;
11
12pub(crate) use primer_schema::Seed;
13use primer_schema::{
14    PrimerBundledDependencies, PrimerDist, PrimerPackument, PrimerPeerDepMeta,
15    PrimerVersionMetadata,
16};
17
18const PRIMER_FORMAT: &str = "rkyv-v1";
19const PRUNE_AGE: Duration = Duration::from_secs(30 * 24 * 60 * 60);
20const AUTO_PRUNE_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
21const AUTO_PRUNE_DENOMINATOR: u8 = 100;
22
23include!(concat!(env!("OUT_DIR"), "/primer_index.rs"));
24
25#[derive(Default)]
26pub struct PruneStats {
27    pub files: u64,
28    pub bytes: u64,
29}
30
31impl Seed {
32    pub(crate) fn packument(&self) -> Packument {
33        self.packument.to_packument()
34    }
35}
36
37impl PrimerPackument {
38    fn to_packument(&self) -> Packument {
39        let mut time = BTreeMap::new();
40        let versions = self
41            .versions
42            .iter()
43            .map(|v| {
44                if let Some(published_at) = v.published_at.as_ref() {
45                    time.insert(v.version.clone(), published_at.clone());
46                }
47                (
48                    v.version.clone(),
49                    v.metadata.to_version_metadata(&self.name, &v.version),
50                )
51            })
52            .collect();
53        Packument {
54            name: self.name.clone(),
55            modified: self.modified.clone(),
56            versions,
57            dist_tags: self.dist_tags.clone(),
58            time,
59        }
60    }
61}
62
63impl PrimerVersionMetadata {
64    fn to_version_metadata(&self, name: &str, version: &str) -> VersionMetadata {
65        VersionMetadata {
66            name: name.to_owned(),
67            version: version.to_owned(),
68            dependencies: self.dependencies.clone(),
69            dev_dependencies: BTreeMap::new(),
70            peer_dependencies: self.peer_dependencies.clone(),
71            peer_dependencies_meta: self
72                .peer_dependencies_meta
73                .iter()
74                .map(|(name, meta)| (name.clone(), meta.to_peer_dep_meta()))
75                .collect(),
76            optional_dependencies: self.optional_dependencies.clone(),
77            bundled_dependencies: self
78                .bundled_dependencies
79                .as_ref()
80                .map(PrimerBundledDependencies::to_bundled_dependencies),
81            dist: self.dist.as_ref().map(PrimerDist::to_dist),
82            os: self.os.clone(),
83            cpu: self.cpu.clone(),
84            libc: self.libc.clone(),
85            engines: self.engines.clone(),
86            license: self.license.clone(),
87            funding_url: self.funding_url.clone(),
88            bin: self.bin.clone(),
89            has_install_script: self.has_install_script,
90            deprecated: self.deprecated.clone(),
91            npm_user: self.trusted_publisher.then(|| NpmUser {
92                trusted_publisher: Some(serde_json::json!({"id": "npm-primer"})),
93            }),
94        }
95    }
96}
97
98impl PrimerPeerDepMeta {
99    fn to_peer_dep_meta(&self) -> PeerDepMeta {
100        PeerDepMeta {
101            optional: self.optional,
102        }
103    }
104}
105
106impl PrimerBundledDependencies {
107    fn to_bundled_dependencies(&self) -> BundledDependencies {
108        match self {
109            Self::List(v) => BundledDependencies::List(v.clone()),
110            Self::All(v) => BundledDependencies::All(*v),
111        }
112    }
113}
114
115impl PrimerDist {
116    fn to_dist(&self) -> Dist {
117        Dist {
118            tarball: self.tarball.clone(),
119            integrity: self.integrity.clone(),
120            shasum: self.shasum.clone(),
121            unpacked_size: None,
122            attestations: self.provenance.then(|| Attestations {
123                provenance: Some(serde_json::json!({
124                    "predicateType": "https://slsa.dev/provenance/v1"
125                })),
126            }),
127        }
128    }
129}
130
131static GENERATED_AT: OnceLock<Option<String>> = OnceLock::new();
132static AUTO_PRUNED: OnceLock<()> = OnceLock::new();
133
134pub(crate) fn get(name: &str) -> Option<Seed> {
135    let (_, offset, len) = PRIMER_INDEX
136        .binary_search_by(|(candidate, _, _)| candidate.cmp(&name))
137        .ok()
138        .and_then(|idx| PRIMER_INDEX.get(idx))?;
139    auto_prune_once();
140    let end = offset.checked_add(*len)?;
141    let compressed = PRIMER_BLOB.get(*offset..end)?;
142    let archived = zstd::stream::decode_all(Cursor::new(compressed)).ok()?;
143    rkyv::from_bytes::<Seed, rkyv::rancor::Error>(&archived).ok()
144}
145
146pub(crate) fn covers_cutoff(cutoff: &str) -> bool {
147    generated_at().is_some_and(|generated_at| generated_at.as_str() >= cutoff)
148}
149
150fn generated_at() -> Option<&'static String> {
151    GENERATED_AT
152        .get_or_init(|| {
153            let secs = option_env!("AUBE_PRIMER_GENERATED_AT")?.parse().ok()?;
154            Some(crate::types::format_iso8601_utc(secs))
155        })
156        .as_ref()
157}
158
159fn auto_prune_once() {
160    AUTO_PRUNED.get_or_init(|| {
161        if let Some(dir) = primer_cache_dir() {
162            auto_prune(&dir);
163        }
164    });
165}
166
167fn auto_prune(dir: &Path) {
168    if !random_byte().is_multiple_of(AUTO_PRUNE_DENOMINATOR) {
169        return;
170    }
171    if let Err(e) = prune_old(dir, PRUNE_AGE, false, Some(AUTO_PRUNE_COOLDOWN)) {
172        tracing::debug!("failed to prune old primer cache files: {e}");
173    }
174}
175
176pub fn prune_cache(dry_run: bool, age: Duration) -> std::io::Result<PruneStats> {
177    let Some(dir) = primer_cache_dir() else {
178        return Ok(PruneStats::default());
179    };
180    prune_old(&dir, age, dry_run, None)
181}
182
183fn prune_old(
184    dir: &Path,
185    age: Duration,
186    dry_run: bool,
187    sentinel_cooldown: Option<Duration>,
188) -> std::io::Result<PruneStats> {
189    let mut stats = PruneStats::default();
190    std::fs::create_dir_all(dir)?;
191    let sentinel = dir.join(".auto_prune");
192    if let Some(cooldown) = sentinel_cooldown
193        && let Ok(modified) = sentinel.metadata().and_then(|m| m.modified())
194        && modified.elapsed().unwrap_or_default() < cooldown
195    {
196        return Ok(stats);
197    }
198    if sentinel_cooldown.is_some() {
199        touch(&sentinel)?;
200    }
201    let entries = std::fs::read_dir(dir)?;
202    for entry in entries {
203        let entry = entry?;
204        let path = entry.path();
205        let Some(name) = path.file_name().and_then(|s| s.to_str()) else {
206            continue;
207        };
208        if !is_primer_cache_file(name) {
209            continue;
210        }
211        let metadata = entry.metadata()?;
212        if metadata.modified()?.elapsed().unwrap_or_default() > age {
213            stats.files += 1;
214            stats.bytes += metadata.len();
215            if !dry_run {
216                std::fs::remove_file(&path)?;
217            }
218        }
219    }
220    Ok(stats)
221}
222
223fn touch(path: &Path) -> std::io::Result<()> {
224    if let Some(parent) = path.parent() {
225        std::fs::create_dir_all(parent)?;
226    }
227    std::fs::OpenOptions::new()
228        .create(true)
229        .write(true)
230        .truncate(true)
231        .open(path)?
232        .write_all(b"\n")
233}
234
235fn is_primer_cache_file(name: &str) -> bool {
236    name.starts_with(&format!("{PRIMER_FORMAT}-")) && name.ends_with(".rkyv")
237}
238
239fn random_byte() -> u8 {
240    let nanos = std::time::SystemTime::now()
241        .duration_since(std::time::UNIX_EPOCH)
242        .map(|d| d.as_nanos())
243        .unwrap_or_default();
244    (nanos as u8) ^ (std::process::id() as u8)
245}
246
247fn primer_cache_dir() -> Option<PathBuf> {
248    if let Some(base) = std::env::var_os("AUBE_CACHE_DIR") {
249        return Some(PathBuf::from(base).join("primer"));
250    }
251    cache_base_dir().map(|p| p.join("aube").join("primer"))
252}
253
254#[cfg(unix)]
255fn cache_base_dir() -> Option<PathBuf> {
256    std::env::var_os("XDG_CACHE_HOME")
257        .map(PathBuf::from)
258        .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".cache")))
259}
260
261#[cfg(windows)]
262fn cache_base_dir() -> Option<PathBuf> {
263    std::env::var_os("LOCALAPPDATA").map(PathBuf::from)
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[test]
271    fn bundled_primer_loads() {
272        let Some((name, _, _)) = PRIMER_INDEX.first() else {
273            return;
274        };
275        assert!(super::get(name).is_some());
276    }
277
278    #[test]
279    fn primer_cache_file_match_is_narrow() {
280        assert!(is_primer_cache_file("rkyv-v1-abc.rkyv"));
281        assert!(!is_primer_cache_file(".auto_prune"));
282        assert!(!is_primer_cache_file("rkyv-v1-abc.tmp"));
283        assert!(!is_primer_cache_file("other-v1-abc.rkyv"));
284    }
285
286    #[test]
287    fn prune_removes_old_extracted_primer_files() {
288        let temp = tempfile::tempdir().unwrap();
289        let dir = temp.path();
290        std::fs::write(dir.join("rkyv-v1-old-0-old.rkyv"), "{}").unwrap();
291        std::fs::write(dir.join("packument.json"), "{}").unwrap();
292        let stats = prune_old(dir, Duration::from_secs(0), false, None).unwrap();
293        assert_eq!(stats.files, 1);
294        assert!(!dir.join("rkyv-v1-old-0-old.rkyv").exists());
295        assert!(dir.join("packument.json").exists());
296    }
297
298    #[test]
299    fn prune_sentinel_uses_own_cooldown() {
300        let temp = tempfile::tempdir().unwrap();
301        let dir = temp.path();
302        let primer_file = dir.join("rkyv-v1-old-0-old.rkyv");
303        std::fs::write(&primer_file, "{}").unwrap();
304        touch(&dir.join(".auto_prune")).unwrap();
305
306        let stats = prune_old(
307            dir,
308            Duration::from_secs(0),
309            false,
310            Some(Duration::from_secs(60)),
311        )
312        .unwrap();
313
314        assert_eq!(stats.files, 0);
315        assert!(primer_file.exists());
316    }
317}