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            attestations: self.provenance.then(|| Attestations {
122                provenance: Some(serde_json::json!({
123                    "predicateType": "https://slsa.dev/provenance/v1"
124                })),
125            }),
126        }
127    }
128}
129
130static GENERATED_AT: OnceLock<Option<String>> = OnceLock::new();
131static AUTO_PRUNED: OnceLock<()> = OnceLock::new();
132
133pub(crate) fn get(name: &str) -> Option<Seed> {
134    let (_, offset, len) = PRIMER_INDEX
135        .binary_search_by(|(candidate, _, _)| candidate.cmp(&name))
136        .ok()
137        .and_then(|idx| PRIMER_INDEX.get(idx))?;
138    auto_prune_once();
139    let end = offset.checked_add(*len)?;
140    let compressed = PRIMER_BLOB.get(*offset..end)?;
141    let archived = zstd::stream::decode_all(Cursor::new(compressed)).ok()?;
142    rkyv::from_bytes::<Seed, rkyv::rancor::Error>(&archived).ok()
143}
144
145pub(crate) fn covers_cutoff(cutoff: &str) -> bool {
146    generated_at().is_some_and(|generated_at| generated_at.as_str() >= cutoff)
147}
148
149fn generated_at() -> Option<&'static String> {
150    GENERATED_AT
151        .get_or_init(|| {
152            let secs = option_env!("AUBE_PRIMER_GENERATED_AT")?.parse().ok()?;
153            Some(crate::types::format_iso8601_utc(secs))
154        })
155        .as_ref()
156}
157
158fn auto_prune_once() {
159    AUTO_PRUNED.get_or_init(|| {
160        if let Some(dir) = primer_cache_dir() {
161            auto_prune(&dir);
162        }
163    });
164}
165
166fn auto_prune(dir: &Path) {
167    if !random_byte().is_multiple_of(AUTO_PRUNE_DENOMINATOR) {
168        return;
169    }
170    if let Err(e) = prune_old(dir, PRUNE_AGE, false, Some(AUTO_PRUNE_COOLDOWN)) {
171        tracing::debug!("failed to prune old primer cache files: {e}");
172    }
173}
174
175pub fn prune_cache(dry_run: bool, age: Duration) -> std::io::Result<PruneStats> {
176    let Some(dir) = primer_cache_dir() else {
177        return Ok(PruneStats::default());
178    };
179    prune_old(&dir, age, dry_run, None)
180}
181
182fn prune_old(
183    dir: &Path,
184    age: Duration,
185    dry_run: bool,
186    sentinel_cooldown: Option<Duration>,
187) -> std::io::Result<PruneStats> {
188    let mut stats = PruneStats::default();
189    std::fs::create_dir_all(dir)?;
190    let sentinel = dir.join(".auto_prune");
191    if let Some(cooldown) = sentinel_cooldown
192        && let Ok(modified) = sentinel.metadata().and_then(|m| m.modified())
193        && modified.elapsed().unwrap_or_default() < cooldown
194    {
195        return Ok(stats);
196    }
197    if sentinel_cooldown.is_some() {
198        touch(&sentinel)?;
199    }
200    let entries = std::fs::read_dir(dir)?;
201    for entry in entries {
202        let entry = entry?;
203        let path = entry.path();
204        let Some(name) = path.file_name().and_then(|s| s.to_str()) else {
205            continue;
206        };
207        if !is_primer_cache_file(name) {
208            continue;
209        }
210        let metadata = entry.metadata()?;
211        if metadata.modified()?.elapsed().unwrap_or_default() > age {
212            stats.files += 1;
213            stats.bytes += metadata.len();
214            if !dry_run {
215                std::fs::remove_file(&path)?;
216            }
217        }
218    }
219    Ok(stats)
220}
221
222fn touch(path: &Path) -> std::io::Result<()> {
223    if let Some(parent) = path.parent() {
224        std::fs::create_dir_all(parent)?;
225    }
226    std::fs::OpenOptions::new()
227        .create(true)
228        .write(true)
229        .truncate(true)
230        .open(path)?
231        .write_all(b"\n")
232}
233
234fn is_primer_cache_file(name: &str) -> bool {
235    name.starts_with(&format!("{PRIMER_FORMAT}-")) && name.ends_with(".rkyv")
236}
237
238fn random_byte() -> u8 {
239    let nanos = std::time::SystemTime::now()
240        .duration_since(std::time::UNIX_EPOCH)
241        .map(|d| d.as_nanos())
242        .unwrap_or_default();
243    (nanos as u8) ^ (std::process::id() as u8)
244}
245
246fn primer_cache_dir() -> Option<PathBuf> {
247    if let Some(base) = std::env::var_os("AUBE_CACHE_DIR") {
248        return Some(PathBuf::from(base).join("primer"));
249    }
250    cache_base_dir().map(|p| p.join("aube").join("primer"))
251}
252
253#[cfg(unix)]
254fn cache_base_dir() -> Option<PathBuf> {
255    std::env::var_os("XDG_CACHE_HOME")
256        .map(PathBuf::from)
257        .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".cache")))
258}
259
260#[cfg(windows)]
261fn cache_base_dir() -> Option<PathBuf> {
262    std::env::var_os("LOCALAPPDATA").map(PathBuf::from)
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    fn bundled_primer_loads() {
271        let Some((name, _, _)) = PRIMER_INDEX.first() else {
272            return;
273        };
274        assert!(super::get(name).is_some());
275    }
276
277    #[test]
278    fn primer_cache_file_match_is_narrow() {
279        assert!(is_primer_cache_file("rkyv-v1-abc.rkyv"));
280        assert!(!is_primer_cache_file(".auto_prune"));
281        assert!(!is_primer_cache_file("rkyv-v1-abc.tmp"));
282        assert!(!is_primer_cache_file("other-v1-abc.rkyv"));
283    }
284
285    #[test]
286    fn prune_removes_old_extracted_primer_files() {
287        let temp = tempfile::tempdir().unwrap();
288        let dir = temp.path();
289        std::fs::write(dir.join("rkyv-v1-old-0-old.rkyv"), "{}").unwrap();
290        std::fs::write(dir.join("packument.json"), "{}").unwrap();
291        let stats = prune_old(dir, Duration::from_secs(0), false, None).unwrap();
292        assert_eq!(stats.files, 1);
293        assert!(!dir.join("rkyv-v1-old-0-old.rkyv").exists());
294        assert!(dir.join("packument.json").exists());
295    }
296
297    #[test]
298    fn prune_sentinel_uses_own_cooldown() {
299        let temp = tempfile::tempdir().unwrap();
300        let dir = temp.path();
301        let primer_file = dir.join("rkyv-v1-old-0-old.rkyv");
302        std::fs::write(&primer_file, "{}").unwrap();
303        touch(&dir.join(".auto_prune")).unwrap();
304
305        let stats = prune_old(
306            dir,
307            Duration::from_secs(0),
308            false,
309            Some(Duration::from_secs(60)),
310        )
311        .unwrap();
312
313        assert_eq!(stats.files, 0);
314        assert!(primer_file.exists());
315    }
316}