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(|d| d.to_dist(name, version)),
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, name: &str, version: &str) -> Dist {
117        Dist {
118            tarball: self
119                .tarball
120                .clone()
121                .unwrap_or_else(|| deterministic_tarball_url(name, version)),
122            integrity: self.integrity.clone(),
123            shasum: None,
124            unpacked_size: None,
125            attestations: self.provenance.then(|| Attestations {
126                provenance: Some(serde_json::json!({
127                    "predicateType": "https://slsa.dev/provenance/v1"
128                })),
129            }),
130        }
131    }
132}
133
134/// Reconstruct the npmjs tarball URL when the primer omitted it
135/// (the common case — see PrimerDist::tarball docs). Mirrors
136/// `RegistryClient::tarball_url`'s format for `registry.npmjs.org`.
137/// In force-metadata-primer mode the URL is rewritten to the active
138/// registry by the resolver, so this default is only consulted on
139/// the default-registry path.
140fn deterministic_tarball_url(name: &str, version: &str) -> String {
141    let unscoped = name
142        .strip_prefix('@')
143        .and_then(|rest| rest.split('/').nth(1))
144        .unwrap_or(name);
145    format!("https://registry.npmjs.org/{name}/-/{unscoped}-{version}.tgz")
146}
147
148static GENERATED_AT: OnceLock<Option<String>> = OnceLock::new();
149static AUTO_PRUNED: OnceLock<()> = OnceLock::new();
150
151pub(crate) fn get(name: &str) -> Option<Seed> {
152    let (_, offset, len) = PRIMER_INDEX
153        .binary_search_by(|(candidate, _, _)| candidate.cmp(&name))
154        .ok()
155        .and_then(|idx| PRIMER_INDEX.get(idx))?;
156    auto_prune_once();
157    let end = offset.checked_add(*len)?;
158    let compressed = PRIMER_BLOB.get(*offset..end)?;
159    let archived = zstd::stream::decode_all(Cursor::new(compressed)).ok()?;
160    rkyv::from_bytes::<Seed, rkyv::rancor::Error>(&archived).ok()
161}
162
163pub(crate) fn covers_cutoff(cutoff: &str) -> bool {
164    generated_at().is_some_and(|generated_at| generated_at.as_str() >= cutoff)
165}
166
167fn generated_at() -> Option<&'static String> {
168    GENERATED_AT
169        .get_or_init(|| {
170            let secs = option_env!("AUBE_PRIMER_GENERATED_AT")?.parse().ok()?;
171            Some(crate::types::format_iso8601_utc(secs))
172        })
173        .as_ref()
174}
175
176fn auto_prune_once() {
177    AUTO_PRUNED.get_or_init(|| {
178        if let Some(dir) = primer_cache_dir() {
179            auto_prune(&dir);
180        }
181    });
182}
183
184fn auto_prune(dir: &Path) {
185    if !random_byte().is_multiple_of(AUTO_PRUNE_DENOMINATOR) {
186        return;
187    }
188    if let Err(e) = prune_old(dir, PRUNE_AGE, false, Some(AUTO_PRUNE_COOLDOWN)) {
189        tracing::debug!("failed to prune old primer cache files: {e}");
190    }
191}
192
193pub fn prune_cache(dry_run: bool, age: Duration) -> std::io::Result<PruneStats> {
194    let Some(dir) = primer_cache_dir() else {
195        return Ok(PruneStats::default());
196    };
197    prune_old(&dir, age, dry_run, None)
198}
199
200fn prune_old(
201    dir: &Path,
202    age: Duration,
203    dry_run: bool,
204    sentinel_cooldown: Option<Duration>,
205) -> std::io::Result<PruneStats> {
206    let mut stats = PruneStats::default();
207    std::fs::create_dir_all(dir)?;
208    let sentinel = dir.join(".auto_prune");
209    if let Some(cooldown) = sentinel_cooldown
210        && let Ok(modified) = sentinel.metadata().and_then(|m| m.modified())
211        && modified.elapsed().unwrap_or_default() < cooldown
212    {
213        return Ok(stats);
214    }
215    if sentinel_cooldown.is_some() {
216        touch(&sentinel)?;
217    }
218    let entries = std::fs::read_dir(dir)?;
219    for entry in entries {
220        let entry = entry?;
221        let path = entry.path();
222        let Some(name) = path.file_name().and_then(|s| s.to_str()) else {
223            continue;
224        };
225        if !is_primer_cache_file(name) {
226            continue;
227        }
228        let metadata = entry.metadata()?;
229        if metadata.modified()?.elapsed().unwrap_or_default() > age {
230            stats.files += 1;
231            stats.bytes += metadata.len();
232            if !dry_run {
233                std::fs::remove_file(&path)?;
234            }
235        }
236    }
237    Ok(stats)
238}
239
240fn touch(path: &Path) -> std::io::Result<()> {
241    if let Some(parent) = path.parent() {
242        std::fs::create_dir_all(parent)?;
243    }
244    std::fs::OpenOptions::new()
245        .create(true)
246        .write(true)
247        .truncate(true)
248        .open(path)?
249        .write_all(b"\n")
250}
251
252fn is_primer_cache_file(name: &str) -> bool {
253    name.starts_with(&format!("{PRIMER_FORMAT}-")) && name.ends_with(".rkyv")
254}
255
256fn random_byte() -> u8 {
257    let nanos = std::time::SystemTime::now()
258        .duration_since(std::time::UNIX_EPOCH)
259        .map(|d| d.as_nanos())
260        .unwrap_or_default();
261    (nanos as u8) ^ (std::process::id() as u8)
262}
263
264fn primer_cache_dir() -> Option<PathBuf> {
265    if let Some(base) = std::env::var_os("AUBE_CACHE_DIR") {
266        return Some(PathBuf::from(base).join("primer"));
267    }
268    cache_base_dir().map(|p| p.join("aube").join("primer"))
269}
270
271#[cfg(unix)]
272fn cache_base_dir() -> Option<PathBuf> {
273    std::env::var_os("XDG_CACHE_HOME")
274        .map(PathBuf::from)
275        .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".cache")))
276}
277
278#[cfg(windows)]
279fn cache_base_dir() -> Option<PathBuf> {
280    std::env::var_os("LOCALAPPDATA").map(PathBuf::from)
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    #[test]
288    fn bundled_primer_loads() {
289        let Some((name, _, _)) = PRIMER_INDEX.first() else {
290            return;
291        };
292        assert!(super::get(name).is_some());
293    }
294
295    #[test]
296    fn bundled_primer_synthesizes_tarball_urls() {
297        // The generator omits the tarball URL when it matches the
298        // deterministic `{registry}/{name}/-/{unscoped}-{version}.tgz`
299        // pattern. Verify the runtime fills it in correctly: every
300        // dist must surface a tarball URL whose path segments match
301        // the package name + version we asked for, so a synthesis bug
302        // that drops or swaps either field can't pass silently.
303        let Some((name, _, _)) = PRIMER_INDEX.first() else {
304            return;
305        };
306        let packument = super::get(name).expect("primer hit").packument();
307        let (version, meta) = packument
308            .versions
309            .iter()
310            .find(|(_, v)| v.dist.is_some())
311            .expect("packument has at least one version with dist metadata");
312        let dist = meta.dist.as_ref().unwrap();
313        assert!(
314            dist.tarball.starts_with("https://"),
315            "tarball: {}",
316            dist.tarball
317        );
318        assert!(dist.tarball.ends_with(".tgz"), "tarball: {}", dist.tarball);
319        assert!(
320            dist.tarball.contains(*name),
321            "tarball {} missing package name {name}",
322            dist.tarball,
323        );
324        assert!(
325            dist.tarball.contains(version),
326            "tarball {} missing version {version}",
327            dist.tarball,
328        );
329    }
330
331    #[test]
332    fn deterministic_tarball_url_handles_scoped_names() {
333        assert_eq!(
334            deterministic_tarball_url("react", "18.2.0"),
335            "https://registry.npmjs.org/react/-/react-18.2.0.tgz"
336        );
337        assert_eq!(
338            deterministic_tarball_url("@types/node", "20.10.0"),
339            "https://registry.npmjs.org/@types/node/-/node-20.10.0.tgz"
340        );
341    }
342
343    #[test]
344    fn primer_cache_file_match_is_narrow() {
345        assert!(is_primer_cache_file("rkyv-v1-abc.rkyv"));
346        assert!(!is_primer_cache_file(".auto_prune"));
347        assert!(!is_primer_cache_file("rkyv-v1-abc.tmp"));
348        assert!(!is_primer_cache_file("other-v1-abc.rkyv"));
349    }
350
351    #[test]
352    fn prune_removes_old_extracted_primer_files() {
353        let temp = tempfile::tempdir().unwrap();
354        let dir = temp.path();
355        std::fs::write(dir.join("rkyv-v1-old-0-old.rkyv"), "{}").unwrap();
356        std::fs::write(dir.join("packument.json"), "{}").unwrap();
357        let stats = prune_old(dir, Duration::from_secs(0), false, None).unwrap();
358        assert_eq!(stats.files, 1);
359        assert!(!dir.join("rkyv-v1-old-0-old.rkyv").exists());
360        assert!(dir.join("packument.json").exists());
361    }
362
363    #[test]
364    fn prune_sentinel_uses_own_cooldown() {
365        let temp = tempfile::tempdir().unwrap();
366        let dir = temp.path();
367        let primer_file = dir.join("rkyv-v1-old-0-old.rkyv");
368        std::fs::write(&primer_file, "{}").unwrap();
369        touch(&dir.join(".auto_prune")).unwrap();
370
371        let stats = prune_old(
372            dir,
373            Duration::from_secs(0),
374            false,
375            Some(Duration::from_secs(60)),
376        )
377        .unwrap();
378
379        assert_eq!(stats.files, 0);
380        assert!(primer_file.exists());
381    }
382}