Skip to main content

oci_zero/
pull.rs

1//! Callback-driven, allocation-free OCI graph traversal.
2
3use crate::{
4    digest::Digest,
5    metadata::{Descriptor, Document, DocumentKind, ImageConfig, ImageIndex, ImageManifest},
6    reference::{Reference, Selector},
7};
8
9pub const OCI_IMAGE_CONFIG: &str = "application/vnd.oci.image.config.v1+json";
10pub const DOCKER_IMAGE_CONFIG: &str = "application/vnd.docker.container.image.v1+json";
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq)]
13pub enum ManifestReference<'a> {
14    Tag(&'a str),
15    Digest(Digest),
16}
17
18#[derive(Clone, Copy, Debug, Eq, PartialEq)]
19pub enum Selection {
20    Skip,
21    Pull,
22}
23
24#[derive(Clone, Copy, Debug, Eq, PartialEq)]
25pub enum BlobAction {
26    Skip,
27    Fetch,
28}
29
30#[derive(Clone, Copy, Debug, Eq, PartialEq)]
31pub enum BlobKind {
32    Config,
33    Layer { diff_id: Option<Digest> },
34}
35
36/// Receives streamed blob bytes from a [`Fetcher`].
37pub trait BlobSink {
38    fn chunk(&mut self, bytes: &[u8]);
39
40    fn cancelled(&self) -> bool {
41        false
42    }
43}
44
45/// Supplies verified registry objects to the transport-independent pull walk.
46///
47/// Implementations must verify manifest digest references and each blob
48/// descriptor's size and digest before returning success.
49pub trait Fetcher {
50    type Error;
51
52    async fn manifest(
53        &mut self,
54        reference: ManifestReference<'_>,
55        destination: &mut [u8],
56    ) -> Result<usize, Self::Error>;
57
58    async fn blob<S: BlobSink>(
59        &mut self,
60        descriptor: Descriptor<'_>,
61        sink: &mut S,
62    ) -> Result<(), Self::Error>;
63}
64
65/// Selects manifests and consumes pull events without retaining graph lists.
66pub trait PullVisitor {
67    type Error;
68
69    fn index(&mut self, _index: ImageIndex<'_>) -> Result<(), Self::Error> {
70        Ok(())
71    }
72
73    fn select_manifest(&mut self, _descriptor: Descriptor<'_>) -> Result<Selection, Self::Error> {
74        Ok(Selection::Pull)
75    }
76
77    fn manifest(&mut self, _manifest: ImageManifest<'_>) -> Result<(), Self::Error> {
78        Ok(())
79    }
80
81    fn subject(&mut self, _descriptor: Descriptor<'_>) -> Result<(), Self::Error> {
82        Ok(())
83    }
84
85    fn begin_blob(
86        &mut self,
87        _kind: BlobKind,
88        _descriptor: Descriptor<'_>,
89    ) -> Result<BlobAction, Self::Error> {
90        Ok(BlobAction::Fetch)
91    }
92
93    fn blob_data(&mut self, _kind: BlobKind, _bytes: &[u8]) -> Result<(), Self::Error> {
94        Ok(())
95    }
96
97    fn end_blob(&mut self, _kind: BlobKind) -> Result<(), Self::Error> {
98        Ok(())
99    }
100}
101
102pub struct PullBuffers<'a> {
103    pub root_manifest: &'a mut [u8],
104    pub child_manifest: &'a mut [u8],
105    pub config: &'a mut [u8],
106}
107
108/// Pulls a root manifest, selects index children, fetches configs, and streams
109/// selected layers in manifest order.
110pub async fn pull<F, V>(
111    fetcher: &mut F,
112    reference: Reference<'_>,
113    buffers: PullBuffers<'_>,
114    visitor: &mut V,
115) -> Result<(), PullError<F::Error, V::Error>>
116where
117    F: Fetcher,
118    V: PullVisitor,
119{
120    let root_reference = match reference.selector() {
121        Selector::Tag(tag) => ManifestReference::Tag(tag),
122        Selector::Digest(digest) => ManifestReference::Digest(digest),
123    };
124    let root_length = fetcher
125        .manifest(root_reference, buffers.root_manifest)
126        .await
127        .map_err(PullError::Fetch)?;
128    if root_length > buffers.root_manifest.len() {
129        return Err(PullError::BufferContract);
130    }
131    let root =
132        Document::parse(&buffers.root_manifest[..root_length]).map_err(PullError::Metadata)?;
133    match root.kind() {
134        DocumentKind::Manifest => {
135            pull_manifest(
136                fetcher,
137                root.manifest().map_err(PullError::Metadata)?,
138                buffers.config,
139                visitor,
140            )
141            .await
142        }
143        DocumentKind::Index => {
144            let index = root.index().map_err(PullError::Metadata)?;
145            visitor.index(index).map_err(PullError::Visitor)?;
146            for descriptor in index.manifests().map_err(PullError::Metadata)? {
147                let descriptor = descriptor.map_err(PullError::Metadata)?;
148                if visitor
149                    .select_manifest(descriptor)
150                    .map_err(PullError::Visitor)?
151                    == Selection::Skip
152                {
153                    continue;
154                }
155                let length = fetcher
156                    .manifest(
157                        ManifestReference::Digest(descriptor.digest()),
158                        buffers.child_manifest,
159                    )
160                    .await
161                    .map_err(PullError::Fetch)?;
162                if length > buffers.child_manifest.len() {
163                    return Err(PullError::BufferContract);
164                }
165                let child = Document::parse(&buffers.child_manifest[..length])
166                    .map_err(PullError::Metadata)?;
167                if child.kind() != DocumentKind::Manifest {
168                    return Err(PullError::NestedIndex);
169                }
170                pull_manifest(
171                    fetcher,
172                    child.manifest().map_err(PullError::Metadata)?,
173                    buffers.config,
174                    visitor,
175                )
176                .await?;
177            }
178            Ok(())
179        }
180    }
181}
182
183async fn pull_manifest<F, V>(
184    fetcher: &mut F,
185    manifest: ImageManifest<'_>,
186    config_buffer: &mut [u8],
187    visitor: &mut V,
188) -> Result<(), PullError<F::Error, V::Error>>
189where
190    F: Fetcher,
191    V: PullVisitor,
192{
193    visitor.manifest(manifest).map_err(PullError::Visitor)?;
194    if let Some(subject) = manifest.subject().map_err(PullError::Metadata)? {
195        visitor.subject(subject).map_err(PullError::Visitor)?;
196    }
197
198    let config = manifest.config().map_err(PullError::Metadata)?;
199    let mut config_sink = BufferSink::new(config_buffer);
200    fetcher
201        .blob(config, &mut config_sink)
202        .await
203        .map_err(PullError::Fetch)?;
204    let config_length = config_sink
205        .finish()
206        .map_err(|_| PullError::ConfigTooLarge)?;
207    let action = visitor
208        .begin_blob(BlobKind::Config, config)
209        .map_err(PullError::Visitor)?;
210    if action == BlobAction::Fetch {
211        visitor
212            .blob_data(BlobKind::Config, &config_buffer[..config_length])
213            .map_err(PullError::Visitor)?;
214        visitor
215            .end_blob(BlobKind::Config)
216            .map_err(PullError::Visitor)?;
217    }
218
219    let config_media_type = config.media_type();
220    let image_config = config_media_type.decoded_eq_ascii(OCI_IMAGE_CONFIG)
221        || config_media_type.decoded_eq_ascii(DOCKER_IMAGE_CONFIG);
222    let mut diff_ids = if image_config {
223        let config_document =
224            ImageConfig::parse(&config_buffer[..config_length]).map_err(PullError::Metadata)?;
225        let diff_ids = config_document.diff_ids().map_err(PullError::Metadata)?;
226        if diff_ids.is_none() {
227            return Err(PullError::MissingDiffIds);
228        }
229        diff_ids
230    } else {
231        // Artifact configs are media-type-defined blobs and need not be JSON.
232        None
233    };
234
235    for descriptor in manifest.layers().map_err(PullError::Metadata)? {
236        let descriptor = descriptor.map_err(PullError::Metadata)?;
237        let diff_id = match diff_ids.as_mut() {
238            Some(ids) => Some(
239                ids.next()
240                    .ok_or(PullError::DiffIdCount)?
241                    .map_err(PullError::Metadata)?,
242            ),
243            None => None,
244        };
245        let kind = BlobKind::Layer { diff_id };
246        let action = visitor
247            .begin_blob(kind, descriptor)
248            .map_err(PullError::Visitor)?;
249        if action == BlobAction::Skip {
250            continue;
251        }
252        let mut sink = VisitorSink {
253            visitor,
254            kind,
255            error: None,
256        };
257        fetcher
258            .blob(descriptor, &mut sink)
259            .await
260            .map_err(PullError::Fetch)?;
261        if let Some(error) = sink.error {
262            return Err(PullError::Visitor(error));
263        }
264        visitor.end_blob(kind).map_err(PullError::Visitor)?;
265    }
266    if diff_ids.as_mut().and_then(Iterator::next).is_some() {
267        return Err(PullError::DiffIdCount);
268    }
269    Ok(())
270}
271
272struct BufferSink<'a> {
273    buffer: &'a mut [u8],
274    length: usize,
275    overflow: bool,
276}
277
278impl<'a> BufferSink<'a> {
279    fn new(buffer: &'a mut [u8]) -> Self {
280        Self {
281            buffer,
282            length: 0,
283            overflow: false,
284        }
285    }
286
287    fn finish(&self) -> Result<usize, ()> {
288        (!self.overflow).then_some(self.length).ok_or(())
289    }
290}
291
292impl BlobSink for BufferSink<'_> {
293    fn chunk(&mut self, bytes: &[u8]) {
294        self.overflow |= crate::buffer::append(self.buffer, &mut self.length, bytes).is_err();
295    }
296
297    fn cancelled(&self) -> bool {
298        self.overflow
299    }
300}
301
302struct VisitorSink<'a, V: PullVisitor> {
303    visitor: &'a mut V,
304    kind: BlobKind,
305    error: Option<V::Error>,
306}
307
308impl<V: PullVisitor> BlobSink for VisitorSink<'_, V> {
309    fn chunk(&mut self, bytes: &[u8]) {
310        if self.error.is_none() {
311            self.error = self.visitor.blob_data(self.kind, bytes).err();
312        }
313    }
314
315    fn cancelled(&self) -> bool {
316        self.error.is_some()
317    }
318}
319
320#[derive(Debug, derive_more::Display)]
321pub enum PullError<F, V> {
322    #[display("registry fetch failed: {_0}")]
323    Fetch(F),
324    #[display("pull visitor failed: {_0}")]
325    Visitor(V),
326    #[display("OCI metadata failed: {_0}")]
327    Metadata(crate::metadata::MetadataError),
328    #[display("image config exceeds its caller buffer")]
329    ConfigTooLarge,
330    #[display("image config is missing rootfs.diff_ids")]
331    MissingDiffIds,
332    #[display("layer and diff_id counts differ")]
333    DiffIdCount,
334    #[display("nested OCI indexes exceed this pull frame")]
335    NestedIndex,
336    #[display("fetcher returned an out-of-range length")]
337    BufferContract,
338}
339
340#[cfg(test)]
341mod tests {
342    use core::{future::Future, task::Poll};
343    use std::{string::ToString, sync::Arc, task::Wake};
344
345    use super::{
346        pull, BlobAction, BlobKind, BlobSink, BufferSink, Fetcher, ManifestReference, PullBuffers,
347        PullError, PullVisitor, Selection, VisitorSink,
348    };
349    use crate::{metadata::Descriptor, reference::Reference};
350
351    const MANIFEST: &[u8] = br#"{
352        "schemaVersion": 2,
353        "mediaType": "application/vnd.oci.image.manifest.v1+json",
354        "artifactType": "application/example",
355        "config": {
356            "mediaType": "application/vnd.example.config",
357            "digest": "sha256:0000000000000000000000000000000000000000000000000000000000000000",
358            "size": 8
359        },
360        "layers": [{
361            "mediaType": "application/vnd.example.payload",
362            "digest": "sha256:1111111111111111111111111111111111111111111111111111111111111111",
363            "size": 5
364        }]
365    }"#;
366
367    const ESCAPED_MEDIA_TYPE_MANIFEST: &[u8] = br#"{
368        "schemaVersion": 2,
369        "mediaType": "application/vnd.oci.image.manifest.v1+json",
370        "config": {
371            "mediaType": "application\/vnd.oci.image.config.v1+json",
372            "digest": "sha256:0000000000000000000000000000000000000000000000000000000000000000",
373            "size": 115
374        },
375        "layers": [{
376            "mediaType": "application/vnd.oci.image.layer.v1.tar",
377            "digest": "sha256:1111111111111111111111111111111111111111111111111111111111111111",
378            "size": 5
379        }]
380    }"#;
381
382    const IMAGE_CONFIG: &[u8] = br#"{"rootfs":{"type":"layers","diff_ids":["sha256:2222222222222222222222222222222222222222222222222222222222222222"]}}"#;
383
384    const INDEX: &[u8] = br#"{
385        "schemaVersion": 2,
386        "mediaType": "application/vnd.oci.image.index.v1+json",
387        "manifests": [{
388            "mediaType": "application/vnd.oci.image.manifest.v1+json",
389            "digest": "sha256:3333333333333333333333333333333333333333333333333333333333333333",
390            "size": 0
391        }]
392    }"#;
393
394    struct MockFetcher;
395
396    impl Fetcher for MockFetcher {
397        type Error = ();
398
399        async fn manifest(
400            &mut self,
401            _reference: ManifestReference<'_>,
402            destination: &mut [u8],
403        ) -> Result<usize, Self::Error> {
404            destination[..MANIFEST.len()].copy_from_slice(MANIFEST);
405            Ok(MANIFEST.len())
406        }
407
408        async fn blob<S: BlobSink>(
409            &mut self,
410            descriptor: Descriptor<'_>,
411            sink: &mut S,
412        ) -> Result<(), Self::Error> {
413            if descriptor.media_type().as_str() == Some("application/vnd.example.config") {
414                sink.chunk(b"not json");
415            } else {
416                sink.chunk(b"layer");
417            }
418            Ok(())
419        }
420    }
421
422    struct EscapedMediaTypeFetcher;
423
424    impl Fetcher for EscapedMediaTypeFetcher {
425        type Error = ();
426
427        async fn manifest(
428            &mut self,
429            _reference: ManifestReference<'_>,
430            destination: &mut [u8],
431        ) -> Result<usize, Self::Error> {
432            destination[..ESCAPED_MEDIA_TYPE_MANIFEST.len()]
433                .copy_from_slice(ESCAPED_MEDIA_TYPE_MANIFEST);
434            Ok(ESCAPED_MEDIA_TYPE_MANIFEST.len())
435        }
436
437        async fn blob<S: BlobSink>(
438            &mut self,
439            descriptor: Descriptor<'_>,
440            sink: &mut S,
441        ) -> Result<(), Self::Error> {
442            if descriptor.media_type().encoded().contains("\\/") {
443                sink.chunk(IMAGE_CONFIG);
444            } else {
445                sink.chunk(b"layer");
446            }
447            Ok(())
448        }
449    }
450
451    #[derive(Default)]
452    struct IndexFetcher {
453        manifest_calls: usize,
454    }
455
456    impl Fetcher for IndexFetcher {
457        type Error = ();
458
459        async fn manifest(
460            &mut self,
461            _reference: ManifestReference<'_>,
462            destination: &mut [u8],
463        ) -> Result<usize, Self::Error> {
464            let document = if self.manifest_calls == 0 {
465                INDEX
466            } else {
467                MANIFEST
468            };
469            self.manifest_calls += 1;
470            destination[..document.len()].copy_from_slice(document);
471            Ok(document.len())
472        }
473
474        async fn blob<S: BlobSink>(
475            &mut self,
476            descriptor: Descriptor<'_>,
477            sink: &mut S,
478        ) -> Result<(), Self::Error> {
479            if descriptor.media_type().as_str() == Some("application/vnd.example.config") {
480                sink.chunk(b"not json");
481            } else {
482                sink.chunk(b"layer");
483            }
484            Ok(())
485        }
486    }
487
488    #[derive(Default)]
489    struct Visitor {
490        config_bytes: usize,
491        layer_bytes: usize,
492        layer_diff_id_was_none: bool,
493        layer_diff_id_was_some: bool,
494    }
495
496    impl PullVisitor for Visitor {
497        type Error = ();
498
499        fn begin_blob(
500            &mut self,
501            kind: BlobKind,
502            _descriptor: Descriptor<'_>,
503        ) -> Result<BlobAction, Self::Error> {
504            if let BlobKind::Layer { diff_id } = kind {
505                self.layer_diff_id_was_none = diff_id.is_none();
506                self.layer_diff_id_was_some = diff_id.is_some();
507            }
508            Ok(BlobAction::Fetch)
509        }
510
511        fn blob_data(&mut self, kind: BlobKind, bytes: &[u8]) -> Result<(), Self::Error> {
512            match kind {
513                BlobKind::Config => self.config_bytes += bytes.len(),
514                BlobKind::Layer { .. } => self.layer_bytes += bytes.len(),
515            }
516            Ok(())
517        }
518    }
519
520    struct SkipVisitor;
521
522    impl PullVisitor for SkipVisitor {
523        type Error = ();
524
525        fn select_manifest(
526            &mut self,
527            _descriptor: Descriptor<'_>,
528        ) -> Result<Selection, Self::Error> {
529            Ok(Selection::Skip)
530        }
531    }
532
533    #[test]
534    fn pulls_non_json_artifact_configs() {
535        let mut fetcher = MockFetcher;
536        let mut visitor = Visitor::default();
537        let mut root = [0; 1024];
538        let mut child = [0; 1024];
539        let mut config = [0; 64];
540        let future = pull(
541            &mut fetcher,
542            Reference::parse("oci://example.com/artifact:latest").unwrap(),
543            PullBuffers {
544                root_manifest: &mut root,
545                child_manifest: &mut child,
546                config: &mut config,
547            },
548            &mut visitor,
549        );
550        assert!(block_on_ready(future).is_ok());
551        assert_eq!(visitor.config_bytes, 8);
552        assert_eq!(visitor.layer_bytes, 5);
553        assert!(visitor.layer_diff_id_was_none);
554    }
555
556    #[test]
557    fn recognizes_escaped_image_config_media_types() {
558        let mut fetcher = EscapedMediaTypeFetcher;
559        let mut visitor = Visitor::default();
560        let mut root = [0; 1024];
561        let mut child = [0; 1024];
562        let mut config = [0; 256];
563        let future = pull(
564            &mut fetcher,
565            Reference::parse("oci://example.com/image:latest").unwrap(),
566            PullBuffers {
567                root_manifest: &mut root,
568                child_manifest: &mut child,
569                config: &mut config,
570            },
571            &mut visitor,
572        );
573        assert!(block_on_ready(future).is_ok());
574        assert!(visitor.layer_diff_id_was_some);
575    }
576
577    #[test]
578    fn accepts_manifest_buffers_at_their_exact_contract_size() {
579        let mut fetcher = MockFetcher;
580        let mut visitor = Visitor::default();
581        let mut root = [0; MANIFEST.len()];
582        let mut child = [0; 1];
583        let mut config = [0; 8];
584        let result = block_on_ready(pull(
585            &mut fetcher,
586            Reference::parse("oci://example.com/artifact:latest").unwrap(),
587            PullBuffers {
588                root_manifest: &mut root,
589                child_manifest: &mut child,
590                config: &mut config,
591            },
592            &mut visitor,
593        ));
594        assert!(result.is_ok());
595    }
596
597    #[test]
598    fn pulls_index_children_with_exact_and_roomy_buffers() {
599        for exact in [true, false] {
600            let mut fetcher = IndexFetcher::default();
601            let mut visitor = Visitor::default();
602            let mut root = [0; INDEX.len()];
603            let mut child = [0; 1024];
604            let child = if exact {
605                &mut child[..MANIFEST.len()]
606            } else {
607                &mut child[..]
608            };
609            let mut config = [0; 8];
610            let result = block_on_ready(pull(
611                &mut fetcher,
612                Reference::parse("oci://example.com/artifact:latest").unwrap(),
613                PullBuffers {
614                    root_manifest: &mut root,
615                    child_manifest: child,
616                    config: &mut config,
617                },
618                &mut visitor,
619            ));
620            assert!(result.is_ok());
621            assert_eq!(fetcher.manifest_calls, 2);
622        }
623    }
624
625    #[test]
626    fn skips_unselected_index_children() {
627        let mut fetcher = IndexFetcher::default();
628        let mut visitor = SkipVisitor;
629        let mut root = [0; INDEX.len()];
630        let mut child = [0; MANIFEST.len()];
631        let mut config = [0; 8];
632        let result = block_on_ready(pull(
633            &mut fetcher,
634            Reference::parse("oci://example.com/artifact:latest").unwrap(),
635            PullBuffers {
636                root_manifest: &mut root,
637                child_manifest: &mut child,
638                config: &mut config,
639            },
640            &mut visitor,
641        ));
642        assert!(result.is_ok());
643        assert_eq!(fetcher.manifest_calls, 1);
644    }
645
646    #[test]
647    fn reports_sink_cancellation_state() {
648        struct DefaultSink;
649        impl BlobSink for DefaultSink {
650            fn chunk(&mut self, _bytes: &[u8]) {}
651        }
652        assert!(!DefaultSink.cancelled());
653
654        let mut buffer = [0; 1];
655        let mut sink = BufferSink::new(&mut buffer);
656        assert!(!sink.cancelled());
657        sink.chunk(b"ab");
658        assert!(sink.cancelled());
659
660        struct FailingVisitor;
661        impl PullVisitor for FailingVisitor {
662            type Error = &'static str;
663
664            fn blob_data(&mut self, _kind: BlobKind, _bytes: &[u8]) -> Result<(), Self::Error> {
665                Err("stop")
666            }
667        }
668        let mut visitor = FailingVisitor;
669        let mut sink = VisitorSink {
670            visitor: &mut visitor,
671            kind: BlobKind::Config,
672            error: None,
673        };
674        assert!(!sink.cancelled());
675        sink.chunk(b"data");
676        assert!(sink.cancelled());
677    }
678
679    #[test]
680    fn formats_pull_errors() {
681        let error = PullError::<&str, &str>::BufferContract;
682        assert_eq!(error.to_string(), "fetcher returned an out-of-range length");
683    }
684
685    fn block_on_ready<F: Future>(future: F) -> F::Output {
686        struct NoopWake;
687        impl Wake for NoopWake {
688            fn wake(self: Arc<Self>) {}
689        }
690
691        let waker = std::task::Waker::from(Arc::new(NoopWake));
692        let mut context = core::task::Context::from_waker(&waker);
693        let mut future = core::pin::pin!(future);
694        match future.as_mut().poll(&mut context) {
695            Poll::Ready(output) => output,
696            Poll::Pending => panic!("in-memory pull unexpectedly yielded"),
697        }
698    }
699}