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().as_str().unwrap_or("");
222    let image_config = matches!(config_media_type, OCI_IMAGE_CONFIG | DOCKER_IMAGE_CONFIG);
223    let mut diff_ids = if image_config {
224        let config_document =
225            ImageConfig::parse(&config_buffer[..config_length]).map_err(PullError::Metadata)?;
226        let diff_ids = config_document.diff_ids().map_err(PullError::Metadata)?;
227        if diff_ids.is_none() {
228            return Err(PullError::MissingDiffIds);
229        }
230        diff_ids
231    } else {
232        // Artifact configs are media-type-defined blobs and need not be JSON.
233        None
234    };
235
236    for descriptor in manifest.layers().map_err(PullError::Metadata)? {
237        let descriptor = descriptor.map_err(PullError::Metadata)?;
238        let diff_id = match diff_ids.as_mut() {
239            Some(ids) => Some(
240                ids.next()
241                    .ok_or(PullError::DiffIdCount)?
242                    .map_err(PullError::Metadata)?,
243            ),
244            None => None,
245        };
246        let kind = BlobKind::Layer { diff_id };
247        let action = visitor
248            .begin_blob(kind, descriptor)
249            .map_err(PullError::Visitor)?;
250        if action == BlobAction::Skip {
251            continue;
252        }
253        let mut sink = VisitorSink {
254            visitor,
255            kind,
256            error: None,
257        };
258        fetcher
259            .blob(descriptor, &mut sink)
260            .await
261            .map_err(PullError::Fetch)?;
262        if let Some(error) = sink.error {
263            return Err(PullError::Visitor(error));
264        }
265        visitor.end_blob(kind).map_err(PullError::Visitor)?;
266    }
267    if diff_ids.as_mut().and_then(Iterator::next).is_some() {
268        return Err(PullError::DiffIdCount);
269    }
270    Ok(())
271}
272
273struct BufferSink<'a> {
274    buffer: &'a mut [u8],
275    length: usize,
276    overflow: bool,
277}
278
279impl<'a> BufferSink<'a> {
280    fn new(buffer: &'a mut [u8]) -> Self {
281        Self {
282            buffer,
283            length: 0,
284            overflow: false,
285        }
286    }
287
288    fn finish(&self) -> Result<usize, ()> {
289        (!self.overflow).then_some(self.length).ok_or(())
290    }
291}
292
293impl BlobSink for BufferSink<'_> {
294    fn chunk(&mut self, bytes: &[u8]) {
295        let Some(end) = self.length.checked_add(bytes.len()) else {
296            self.overflow = true;
297            return;
298        };
299        let Some(output) = self.buffer.get_mut(self.length..end) else {
300            self.overflow = true;
301            return;
302        };
303        output.copy_from_slice(bytes);
304        self.length = end;
305    }
306
307    fn cancelled(&self) -> bool {
308        self.overflow
309    }
310}
311
312struct VisitorSink<'a, V: PullVisitor> {
313    visitor: &'a mut V,
314    kind: BlobKind,
315    error: Option<V::Error>,
316}
317
318impl<V: PullVisitor> BlobSink for VisitorSink<'_, V> {
319    fn chunk(&mut self, bytes: &[u8]) {
320        if self.error.is_none() {
321            self.error = self.visitor.blob_data(self.kind, bytes).err();
322        }
323    }
324
325    fn cancelled(&self) -> bool {
326        self.error.is_some()
327    }
328}
329
330#[derive(Debug)]
331pub enum PullError<F, V> {
332    Fetch(F),
333    Visitor(V),
334    Metadata(crate::metadata::MetadataError),
335    ConfigTooLarge,
336    MissingDiffIds,
337    DiffIdCount,
338    NestedIndex,
339    BufferContract,
340}
341
342impl<F: fmt::Display, V: fmt::Display> fmt::Display for PullError<F, V> {
343    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
344        match self {
345            Self::Fetch(error) => write!(formatter, "registry fetch failed: {error}"),
346            Self::Visitor(error) => write!(formatter, "pull visitor failed: {error}"),
347            Self::Metadata(error) => write!(formatter, "OCI metadata failed: {error}"),
348            Self::ConfigTooLarge => formatter.write_str("image config exceeds its caller buffer"),
349            Self::MissingDiffIds => formatter.write_str("image config is missing rootfs.diff_ids"),
350            Self::DiffIdCount => formatter.write_str("layer and diff_id counts differ"),
351            Self::NestedIndex => formatter.write_str("nested OCI indexes exceed this pull frame"),
352            Self::BufferContract => formatter.write_str("fetcher returned an out-of-range length"),
353        }
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use core::{future::Future, task::Poll};
360    use std::{sync::Arc, task::Wake};
361
362    use super::{
363        pull, BlobAction, BlobKind, BlobSink, Fetcher, ManifestReference, PullBuffers, PullVisitor,
364    };
365    use crate::{metadata::Descriptor, reference::Reference};
366
367    const MANIFEST: &[u8] = br#"{
368        "schemaVersion": 2,
369        "mediaType": "application/vnd.oci.image.manifest.v1+json",
370        "artifactType": "application/example",
371        "config": {
372            "mediaType": "application/vnd.example.config",
373            "digest": "sha256:0000000000000000000000000000000000000000000000000000000000000000",
374            "size": 8
375        },
376        "layers": [{
377            "mediaType": "application/vnd.example.payload",
378            "digest": "sha256:1111111111111111111111111111111111111111111111111111111111111111",
379            "size": 5
380        }]
381    }"#;
382
383    struct MockFetcher;
384
385    impl Fetcher for MockFetcher {
386        type Error = ();
387
388        async fn manifest(
389            &mut self,
390            _reference: ManifestReference<'_>,
391            destination: &mut [u8],
392        ) -> Result<usize, Self::Error> {
393            destination[..MANIFEST.len()].copy_from_slice(MANIFEST);
394            Ok(MANIFEST.len())
395        }
396
397        async fn blob<S: BlobSink>(
398            &mut self,
399            descriptor: Descriptor<'_>,
400            sink: &mut S,
401        ) -> Result<(), Self::Error> {
402            if descriptor.media_type().as_str() == Some("application/vnd.example.config") {
403                sink.chunk(b"not json");
404            } else {
405                sink.chunk(b"layer");
406            }
407            Ok(())
408        }
409    }
410
411    #[derive(Default)]
412    struct Visitor {
413        config_bytes: usize,
414        layer_bytes: usize,
415        layer_diff_id_was_none: bool,
416    }
417
418    impl PullVisitor for Visitor {
419        type Error = ();
420
421        fn begin_blob(
422            &mut self,
423            kind: BlobKind,
424            _descriptor: Descriptor<'_>,
425        ) -> Result<BlobAction, Self::Error> {
426            if let BlobKind::Layer { diff_id } = kind {
427                self.layer_diff_id_was_none = diff_id.is_none();
428            }
429            Ok(BlobAction::Fetch)
430        }
431
432        fn blob_data(&mut self, kind: BlobKind, bytes: &[u8]) -> Result<(), Self::Error> {
433            match kind {
434                BlobKind::Config => self.config_bytes += bytes.len(),
435                BlobKind::Layer { .. } => self.layer_bytes += bytes.len(),
436            }
437            Ok(())
438        }
439    }
440
441    #[test]
442    fn pulls_non_json_artifact_configs() {
443        let mut fetcher = MockFetcher;
444        let mut visitor = Visitor::default();
445        let mut root = [0; 1024];
446        let mut child = [0; 1024];
447        let mut config = [0; 64];
448        let future = pull(
449            &mut fetcher,
450            Reference::parse("oci://example.com/artifact:latest").unwrap(),
451            PullBuffers {
452                root_manifest: &mut root,
453                child_manifest: &mut child,
454                config: &mut config,
455            },
456            &mut visitor,
457        );
458        assert!(block_on_ready(future).is_ok());
459        assert_eq!(visitor.config_bytes, 8);
460        assert_eq!(visitor.layer_bytes, 5);
461        assert!(visitor.layer_diff_id_was_none);
462    }
463
464    fn block_on_ready<F: Future>(future: F) -> F::Output {
465        struct NoopWake;
466        impl Wake for NoopWake {
467            fn wake(self: Arc<Self>) {}
468        }
469
470        let waker = std::task::Waker::from(Arc::new(NoopWake));
471        let mut context = core::task::Context::from_waker(&waker);
472        let mut future = core::pin::pin!(future);
473        match future.as_mut().poll(&mut context) {
474            Poll::Ready(output) => output,
475            Poll::Pending => panic!("in-memory pull unexpectedly yielded"),
476        }
477    }
478}
oci_zero/pull.rs

oci_zero/
pull.rs
