1use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9
10use a3s_box_core::error::{BoxError, Result};
11use a3s_box_core::StoredImage;
12
13use super::image::OciImage;
14use super::reference::ImageReference;
15use super::registry::{RegistryAuth, RegistryPuller};
16use super::store::ImageStore;
17
18type PullProgressFn = Arc<dyn Fn(usize, usize, &str, i64) + Send + Sync>;
20
21static PULL_TMP_SEQ: AtomicU64 = AtomicU64::new(0);
23
24pub struct ImagePuller {
26 store: Arc<ImageStore>,
27 puller: RegistryPuller,
28 metrics: Option<crate::prom::RuntimeMetrics>,
29 mirrors: std::collections::HashMap<String, String>,
35}
36
37fn parse_registry_mirrors() -> std::collections::HashMap<String, String> {
39 std::env::var("A3S_REGISTRY_MIRRORS")
40 .ok()
41 .map(|raw| {
42 raw.split(',')
43 .filter_map(|pair| pair.split_once('='))
44 .map(|(host, mirror)| (host.trim().to_string(), mirror.trim().to_string()))
45 .filter(|(host, mirror)| !host.is_empty() && !mirror.is_empty())
46 .collect()
47 })
48 .unwrap_or_default()
49}
50
51impl ImagePuller {
52 pub fn new(store: Arc<ImageStore>, auth: RegistryAuth) -> Self {
54 Self::with_platform(store, auth, None)
55 }
56
57 #[cfg(test)]
58 pub(crate) fn with_registry_puller(store: Arc<ImageStore>, puller: RegistryPuller) -> Self {
59 Self {
60 store,
61 puller,
62 metrics: None,
63 mirrors: std::collections::HashMap::new(),
64 }
65 }
66
67 pub fn with_platform(
71 store: Arc<ImageStore>,
72 auth: RegistryAuth,
73 platform: Option<String>,
74 ) -> Self {
75 Self {
76 store,
77 puller: RegistryPuller::with_auth_and_platform(auth, platform),
78 metrics: None,
79 mirrors: parse_registry_mirrors(),
80 }
81 }
82
83 fn mirror_reference(&self, reference: &ImageReference) -> ImageReference {
87 match self.mirrors.get(&reference.registry) {
88 Some(mirror) if !mirror.is_empty() => {
89 tracing::info!(
90 registry = %reference.registry,
91 mirror = %mirror,
92 "Pulling via configured registry mirror"
93 );
94 let mut mirrored = reference.clone();
95 mirror.clone_into(&mut mirrored.registry);
96 mirrored
97 }
98 _ => reference.clone(),
99 }
100 }
101
102 pub fn set_metrics(mut self, metrics: crate::prom::RuntimeMetrics) -> Self {
104 self.metrics = Some(metrics);
105 self
106 }
107
108 pub fn with_signature_policy(mut self, policy: super::signing::SignaturePolicy) -> Self {
110 self.puller = self.puller.with_signature_policy(policy);
111 self
112 }
113
114 pub fn with_progress_fn(mut self, f: PullProgressFn) -> Self {
116 self.puller = self.puller.with_progress_fn(f);
117 self
118 }
119
120 pub async fn pull(&self, reference: &str) -> Result<OciImage> {
124 self.pull_resolved(reference).await.map(|(image, _)| image)
125 }
126
127 async fn pull_resolved(&self, reference: &str) -> Result<(OciImage, String)> {
128 let reference = reference.trim();
129 if is_digest_reference(reference) {
130 let Some((matched_reference, stored)) = self.cached_digest_image(reference).await?
131 else {
132 return Err(BoxError::OciImageError(format!(
133 "Image digest not found in local cache: {reference}"
134 )));
135 };
136 tracing::info!(
137 requested_reference = %reference,
138 matched_reference = %matched_reference,
139 digest = %stored.digest,
140 "Using cached image by digest"
141 );
142 return Ok((OciImage::from_path(&stored.path)?, matched_reference));
143 }
144
145 let parsed = ImageReference::parse(reference)?;
146
147 if let Some((matched_reference, stored)) = self.cached_image(reference, &parsed).await? {
148 tracing::info!(
149 requested_reference = %reference,
150 matched_reference = %matched_reference,
151 digest = %stored.digest,
152 "Using cached image"
153 );
154 return Ok((OciImage::from_path(&stored.path)?, matched_reference));
155 }
156
157 Ok((self.pull_and_store(&parsed).await?, parsed.full_reference()))
158 }
159
160 pub async fn force_pull(&self, reference: &str) -> Result<OciImage> {
162 let reference = reference.trim();
163 if is_digest_reference(reference) {
164 return Err(BoxError::OciImageError(format!(
165 "Cannot force-pull digest-only reference {reference}; use a tagged registry reference"
166 )));
167 }
168
169 let parsed = ImageReference::parse(reference)?;
170
171 for candidate in cache_reference_candidates(reference, &parsed) {
172 if self.store.get(&candidate).await.is_some() {
173 let _ = self.store.remove(&candidate).await;
174 }
175 }
176
177 self.pull_and_store(&parsed).await
178 }
179
180 pub async fn is_cached(&self, reference: &str) -> bool {
182 let reference = reference.trim();
183 if is_digest_reference(reference) {
184 return matches!(self.cached_digest_image(reference).await, Ok(Some(_)));
185 }
186
187 let parsed = match ImageReference::parse(reference) {
188 Ok(p) => p,
189 Err(_) => return false,
190 };
191 matches!(self.cached_image(reference, &parsed).await, Ok(Some(_)))
192 }
193
194 pub async fn remove_cached(&self, reference: &str) -> Result<bool> {
196 let reference = reference.trim();
197 if is_digest_reference(reference) {
198 if let Some((matched_reference, _)) = self.cached_digest_image(reference).await? {
199 self.store.remove(&matched_reference).await?;
200 return Ok(true);
201 }
202 return Ok(false);
203 }
204
205 let parsed = ImageReference::parse(reference)?;
206 if let Some((matched_reference, _)) = self.cached_image(reference, &parsed).await? {
207 self.store.remove(&matched_reference).await?;
208 Ok(true)
209 } else {
210 Ok(false)
211 }
212 }
213
214 pub async fn list_cached(&self) -> Result<Vec<String>> {
216 Ok(self
217 .store
218 .list()
219 .await
220 .into_iter()
221 .map(|img| img.reference)
222 .collect())
223 }
224
225 async fn pull_and_store(&self, reference: &ImageReference) -> Result<OciImage> {
227 let full_ref = reference.full_reference();
228
229 let fetch = self.mirror_reference(reference);
232
233 let digest = self.puller.pull_manifest_digest(&fetch).await?;
239 super::registry::validated_digest_hex(&digest)?;
240
241 if let Some(stored) = self.store.get_by_digest(&digest).await {
243 tracing::info!(
244 reference = %full_ref,
245 digest = %digest,
246 "Image content already cached under different reference"
247 );
248 self.store.put(&full_ref, &digest, &stored.path).await?;
250 return OciImage::from_path(&stored.path);
251 }
252
253 let digest_hex = digest.strip_prefix("sha256:").unwrap_or(&digest);
257 let tmp_dir = unique_pull_tmp_dir(self.store.store_dir(), digest_hex);
258
259 let pull_start = std::time::Instant::now();
260 if let Err(e) = self.puller.pull(&fetch, &tmp_dir).await {
265 let _ = std::fs::remove_dir_all(&tmp_dir);
266 return Err(e);
267 }
268 if let Some(ref m) = self.metrics {
269 m.image_pull_total.inc();
270 m.image_pull_duration
271 .observe(pull_start.elapsed().as_secs_f64());
272 }
273
274 let stored = match self.store.put(&full_ref, &digest, &tmp_dir).await {
276 Ok(stored) => stored,
277 Err(e) => {
278 let _ = std::fs::remove_dir_all(&tmp_dir);
279 return Err(e);
280 }
281 };
282
283 if let Err(e) = std::fs::remove_dir_all(&tmp_dir) {
285 tracing::warn!(path = %tmp_dir.display(), error = %e, "Failed to remove temp dir after pull");
286 }
287
288 let evicted = self.store.evict().await?;
290 if !evicted.is_empty() {
291 tracing::info!(
292 count = evicted.len(),
293 references = ?evicted,
294 "Evicted images from cache"
295 );
296 }
297
298 OciImage::from_path(&stored.path)
299 }
300
301 async fn cached_image(
302 &self,
303 reference: &str,
304 parsed: &ImageReference,
305 ) -> Result<Option<(String, StoredImage)>> {
306 for candidate in cache_reference_candidates(reference, parsed) {
307 if let Some(stored) = self.store.get(&candidate).await {
308 return Ok(Some((candidate, stored)));
309 }
310 }
311 if let Some(digest) = parsed.digest.as_deref() {
312 return self.cached_digest_image(digest).await;
313 }
314 Ok(None)
315 }
316
317 async fn cached_digest_image(&self, digest: &str) -> Result<Option<(String, StoredImage)>> {
318 let images = self.store.list().await;
319 let mut matches = Vec::new();
320
321 for image in images {
322 let already_matched = matches
323 .iter()
324 .any(|matched: &StoredImage| matched.reference == image.reference);
325 if (digest_matches(&image.digest, digest) || image.reference == digest)
326 && !already_matched
327 {
328 matches.push(image);
329 }
330 }
331
332 match matches.len() {
333 0 => Ok(None),
334 1 => {
335 let stored = matches.pop().expect("checked one match");
336 let stored = self.store.get(&stored.reference).await.unwrap_or(stored);
337 Ok(Some((stored.reference.clone(), stored)))
338 }
339 _ => Err(BoxError::OciImageError(ambiguous_digest_error(
340 digest, &matches,
341 ))),
342 }
343 }
344}
345
346fn ambiguous_digest_error(query: &str, matches: &[StoredImage]) -> String {
347 let mut references: Vec<_> = matches
348 .iter()
349 .map(|image| image.reference.as_str())
350 .collect();
351 references.sort_unstable();
352 format!(
353 "Image digest '{query}' is ambiguous; it matches: {}",
354 references.join(", ")
355 )
356}
357
358fn is_digest_reference(reference: &str) -> bool {
359 reference.starts_with("sha256:")
360}
361
362fn digest_matches(stored_digest: &str, query: &str) -> bool {
363 if stored_digest == query {
364 return true;
365 }
366 let Some(query_hex) = query.strip_prefix("sha256:") else {
367 return false;
368 };
369 !query_hex.is_empty() && stored_digest.starts_with(query)
370}
371
372fn cache_reference_candidates(reference: &str, parsed: &ImageReference) -> Vec<String> {
373 let mut candidates = Vec::new();
374 push_unique(&mut candidates, reference.trim().to_string());
375 push_unique(&mut candidates, parsed.full_reference());
376
377 if parsed.registry == "docker.io" {
378 let repository = parsed
379 .repository
380 .strip_prefix("library/")
381 .unwrap_or(&parsed.repository);
382 push_unique(
383 &mut candidates,
384 reference_from_repository(repository, parsed.tag.as_deref(), parsed.digest.as_deref()),
385 );
386
387 if parsed.digest.is_none() && parsed.tag.as_deref() == Some("latest") {
388 push_unique(&mut candidates, repository.to_string());
389 }
390 }
391
392 if let Some(digest) = &parsed.digest {
393 push_unique(&mut candidates, digest.clone());
394 }
395
396 candidates
397}
398
399fn reference_from_repository(repository: &str, tag: Option<&str>, digest: Option<&str>) -> String {
400 let mut reference = repository.to_string();
401 if let Some(tag) = tag {
402 reference.push(':');
403 reference.push_str(tag);
404 }
405 if let Some(digest) = digest {
406 reference.push('@');
407 reference.push_str(digest);
408 }
409 reference
410}
411
412fn unique_pull_tmp_dir(store_dir: &std::path::Path, digest_hex: &str) -> std::path::PathBuf {
413 let seq = PULL_TMP_SEQ.fetch_add(1, Ordering::Relaxed);
414 store_dir
415 .join("tmp")
416 .join(format!("pull-{digest_hex}-{}-{seq}", std::process::id()))
417}
418
419fn push_unique(values: &mut Vec<String>, value: String) {
420 if !value.is_empty() && !values.iter().any(|existing| existing == &value) {
421 values.push(value);
422 }
423}
424
425#[async_trait::async_trait]
426impl a3s_box_core::traits::ImageRegistry for ImagePuller {
427 async fn pull(&self, reference: &str) -> Result<a3s_box_core::traits::PulledImage> {
428 let (image, resolved_reference) = self.pull_resolved(reference).await?;
429 Ok(a3s_box_core::traits::PulledImage {
430 path: image.root_dir().to_path_buf(),
431 digest: image.manifest_digest().to_string(),
432 reference: resolved_reference,
433 })
434 }
435
436 async fn force_pull(&self, reference: &str) -> Result<a3s_box_core::traits::PulledImage> {
437 let image = self.force_pull(reference).await?;
438 let parsed = ImageReference::parse(reference)?;
439 Ok(a3s_box_core::traits::PulledImage {
440 path: image.root_dir().to_path_buf(),
441 digest: image.manifest_digest().to_string(),
442 reference: parsed.full_reference(),
443 })
444 }
445
446 async fn is_cached(&self, reference: &str) -> bool {
447 self.is_cached(reference).await
448 }
449
450 async fn remove(&self, reference: &str) -> Result<bool> {
451 self.remove_cached(reference).await
452 }
453
454 async fn list_cached(&self) -> Result<Vec<String>> {
455 self.list_cached().await
456 }
457}
458
459#[cfg(test)]
460mod tests {
461 use super::*;
462 use crate::oci::store::ImageStore;
463 use std::path::Path;
464 use tempfile::TempDir;
465
466 #[test]
467 fn test_image_puller_creation() {
468 let tmp = TempDir::new().unwrap();
469 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
470 let _puller = ImagePuller::new(store, RegistryAuth::anonymous());
471 }
472
473 #[tokio::test]
474 async fn test_is_cached_empty_store() {
475 let tmp = TempDir::new().unwrap();
476 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
477 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
478 assert!(!puller.is_cached("nginx:latest").await);
479 }
480
481 #[tokio::test]
482 async fn test_is_cached_invalid_reference() {
483 let tmp = TempDir::new().unwrap();
484 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
485 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
486 assert!(!puller.is_cached("").await);
487 }
488
489 #[tokio::test]
490 async fn test_is_cached_matches_docker_hub_aliases() {
491 let tmp = TempDir::new().unwrap();
492 let source = TempDir::new().unwrap();
493 create_complete_oci_image(source.path());
494 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
495 store
496 .put("alpine:latest", "sha256:manifestxyz789", source.path())
497 .await
498 .unwrap();
499
500 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
501
502 assert!(puller.is_cached("alpine:latest").await);
503 assert!(puller.is_cached("docker.io/library/alpine:latest").await);
504 assert!(puller.is_cached("alpine").await);
505 }
506
507 #[tokio::test]
508 async fn test_pull_uses_cached_short_alias_for_full_reference() {
509 let tmp = TempDir::new().unwrap();
510 let source = TempDir::new().unwrap();
511 create_complete_oci_image(source.path());
512 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
513 store
514 .put("alpine:latest", "sha256:manifestxyz789", source.path())
515 .await
516 .unwrap();
517 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
518
519 let image = puller
520 .pull("docker.io/library/alpine:latest")
521 .await
522 .unwrap();
523
524 assert_eq!(image.manifest_digest(), "sha256:manifestxyz789");
525 }
526
527 #[tokio::test]
528 async fn test_pull_uses_cached_digest_reference_without_registry_parse() {
529 let tmp = TempDir::new().unwrap();
530 let source = TempDir::new().unwrap();
531 create_complete_oci_image(source.path());
532 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
533 store
534 .put("alpine:latest", "sha256:manifestxyz789", source.path())
535 .await
536 .unwrap();
537 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
538
539 let image = puller.pull("sha256:manifestxyz789").await.unwrap();
540
541 assert_eq!(image.manifest_digest(), "sha256:manifestxyz789");
542 }
543
544 #[tokio::test]
545 async fn test_is_cached_matches_digest_prefix() {
546 let tmp = TempDir::new().unwrap();
547 let source = TempDir::new().unwrap();
548 create_complete_oci_image(source.path());
549 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
550 store
551 .put("alpine:latest", "sha256:manifestxyz789", source.path())
552 .await
553 .unwrap();
554 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
555
556 assert!(puller.is_cached("sha256:manifest").await);
557 }
558
559 #[tokio::test]
560 async fn test_pull_reports_missing_digest_reference_as_local_cache_miss() {
561 let tmp = TempDir::new().unwrap();
562 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
563 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
564
565 let error = puller.pull("sha256:notfound").await.unwrap_err();
566
567 assert!(error.to_string().contains("not found in local cache"));
568 }
569
570 #[tokio::test]
571 async fn test_pull_reports_ambiguous_digest_prefix() {
572 let tmp = TempDir::new().unwrap();
573 let source = TempDir::new().unwrap();
574 create_complete_oci_image(source.path());
575 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
576 store
577 .put("alpine:latest", "sha256:manifestaaa", source.path())
578 .await
579 .unwrap();
580 store
581 .put("busybox:latest", "sha256:manifestbbb", source.path())
582 .await
583 .unwrap();
584 let puller = ImagePuller::new(store, RegistryAuth::anonymous());
585
586 let error = puller.pull("sha256:manifest").await.unwrap_err();
587
588 assert!(error.to_string().contains("ambiguous"));
589 assert!(error.to_string().contains("alpine:latest"));
590 assert!(error.to_string().contains("busybox:latest"));
591 }
592
593 #[tokio::test]
594 async fn test_remove_cached_matches_docker_hub_alias() {
595 let tmp = TempDir::new().unwrap();
596 let source = TempDir::new().unwrap();
597 create_complete_oci_image(source.path());
598 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
599 store
600 .put("alpine:latest", "sha256:manifestxyz789", source.path())
601 .await
602 .unwrap();
603 let puller = ImagePuller::new(store.clone(), RegistryAuth::anonymous());
604
605 assert!(puller
606 .remove_cached("docker.io/library/alpine:latest")
607 .await
608 .unwrap());
609 assert!(store.get("alpine:latest").await.is_none());
610 }
611
612 #[test]
613 fn test_set_metrics_attaches_to_puller() {
614 let tmp = TempDir::new().unwrap();
615 let store = Arc::new(ImageStore::new(tmp.path(), 10 * 1024 * 1024).unwrap());
616 let metrics = crate::prom::RuntimeMetrics::new();
617 let puller =
619 ImagePuller::new(store, RegistryAuth::anonymous()).set_metrics(metrics.clone());
620 assert!(puller.metrics.is_some());
621 assert_eq!(metrics.image_pull_total.get(), 0);
622 assert_eq!(metrics.image_pull_duration.get_sample_count(), 0);
623 }
624
625 #[test]
626 fn test_cache_reference_candidates_include_short_docker_hub_aliases() {
627 let parsed = ImageReference::parse("docker.io/library/alpine:latest").unwrap();
628 let candidates = cache_reference_candidates("docker.io/library/alpine:latest", &parsed);
629
630 assert_eq!(
631 candidates,
632 vec![
633 "docker.io/library/alpine:latest".to_string(),
634 "alpine:latest".to_string(),
635 "alpine".to_string(),
636 ]
637 );
638 }
639
640 #[test]
641 fn test_digest_matches_exact_and_prefix_queries() {
642 assert!(digest_matches("sha256:abcdef123456", "sha256:abcdef123456"));
643 assert!(digest_matches("sha256:abcdef123456", "sha256:abcdef"));
644 assert!(!digest_matches("sha256:abcdef123456", "sha256:"));
645 assert!(!digest_matches("sha256:abcdef123456", "abcdef"));
646 }
647
648 #[test]
649 fn test_unique_pull_tmp_dir_does_not_reuse_digest_path() {
650 let tmp = TempDir::new().unwrap();
651 let first = unique_pull_tmp_dir(tmp.path(), "abc123");
652 let second = unique_pull_tmp_dir(tmp.path(), "abc123");
653
654 assert_ne!(first, second);
655 assert_eq!(first.parent().unwrap(), tmp.path().join("tmp"));
656 assert_eq!(second.parent().unwrap(), tmp.path().join("tmp"));
657 assert!(first
658 .file_name()
659 .unwrap()
660 .to_string_lossy()
661 .contains("abc123"));
662 }
663
664 fn create_complete_oci_image(path: &Path) {
665 std::fs::create_dir_all(path.join("blobs/sha256")).unwrap();
666 std::fs::write(path.join("oci-layout"), r#"{"imageLayoutVersion":"1.0.0"}"#).unwrap();
667
668 let config_content = r#"{
669 "architecture": "amd64",
670 "os": "linux",
671 "config": {
672 "Entrypoint": ["/bin/sh"],
673 "Cmd": ["-c", "true"],
674 "Env": ["PATH=/usr/bin:/bin"],
675 "WorkingDir": "/"
676 },
677 "rootfs": {
678 "type": "layers",
679 "diff_ids": ["sha256:layerdiff"]
680 },
681 "history": []
682 }"#;
683 let config_hash = "configabc123";
684 std::fs::write(path.join("blobs/sha256").join(config_hash), config_content).unwrap();
685
686 let layer_hash = "layerdef456";
687 std::fs::write(path.join("blobs/sha256").join(layer_hash), b"layer").unwrap();
688
689 let manifest_content = format!(
690 r#"{{
691 "schemaVersion": 2,
692 "mediaType": "application/vnd.oci.image.manifest.v1+json",
693 "config": {{
694 "mediaType": "application/vnd.oci.image.config.v1+json",
695 "digest": "sha256:{}",
696 "size": {}
697 }},
698 "layers": [
699 {{
700 "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
701 "digest": "sha256:{}",
702 "size": 5
703 }}
704 ]
705 }}"#,
706 config_hash,
707 config_content.len(),
708 layer_hash
709 );
710 let manifest_hash = "manifestxyz789";
711 std::fs::write(
712 path.join("blobs/sha256").join(manifest_hash),
713 &manifest_content,
714 )
715 .unwrap();
716
717 let index_content = format!(
718 r#"{{
719 "schemaVersion": 2,
720 "mediaType": "application/vnd.oci.image.index.v1+json",
721 "manifests": [
722 {{
723 "mediaType": "application/vnd.oci.image.manifest.v1+json",
724 "digest": "sha256:{}",
725 "size": {}
726 }}
727 ]
728 }}"#,
729 manifest_hash,
730 manifest_content.len()
731 );
732 std::fs::write(path.join("index.json"), index_content).unwrap();
733 }
734}