oci_zero/
pull.rs

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

oci_zero/
pull.rs
