1use 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
36pub trait BlobSink {
38 fn chunk(&mut self, bytes: &[u8]);
39
40 fn cancelled(&self) -> bool {
41 false
42 }
43}
44
45pub 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
65pub 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
108pub 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 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}