Merge pull request #145 from SuperTux88/partial-pull

Partial Requests and some other pull refactorings

Author: thomastaylor312

src/client.rs

@@ -19,13 +19,13 @@ use crate::Reference;
 use crate::errors::{OciDistributionError, Result};
 use crate::token_cache::{RegistryOperation, RegistryToken, RegistryTokenType, TokenCache};
 use futures_util::future;
-use futures_util::stream::{self, StreamExt, TryStreamExt};
-use futures_util::Stream;
-use http::HeaderValue;
+use futures_util::stream::{self, BoxStream, StreamExt, TryStreamExt};
+use http::header::RANGE;
+use http::{HeaderValue, StatusCode};
 use http_auth::{parser::ChallengeParser, ChallengeRef};
 use olpc_cjson::CanonicalFormatter;
 use reqwest::header::HeaderMap;
-use reqwest::{RequestBuilder, Url};
+use reqwest::{RequestBuilder, Response, Url};
 use serde::Deserialize;
 use serde::Serialize;
 use sha2::Digest;
@@ -173,6 +173,22 @@ impl ImageLayer {
     }
 }
 
+/// Stream response of a blob with optional content length if available
+pub struct SizedStream {
+    /// The length of the stream if the upstream registry sent a `Content-Length` header
+    pub content_length: Option<u64>,
+    /// The stream of bytes
+    pub stream: BoxStream<'static, std::result::Result<bytes::Bytes, std::io::Error>>,
+}
+
+/// The response of a partial blob request
+pub enum BlobResponse {
+    /// The response is a full blob (for example when partial requests aren't supported)
+    Full(SizedStream),
+    /// The response is a partial blob as requested
+    Partial(SizedStream),
+}
+
 /// The data and media type for a configuration object
 #[derive(Clone)]
 pub struct Config {
@@ -1021,39 +1037,7 @@ impl Client {
         layer: impl AsLayerDescriptor,
         mut out: T,
     ) -> Result<()> {
-        let layer = layer.as_layer_descriptor();
-        let url = self.to_v2_blob_url(image.resolve_registry(), image.repository(), layer.digest);
-
-        let mut response = RequestBuilderWrapper::from_client(self, |client| client.get(&url))
-            .apply_accept(MIME_TYPES_DISTRIBUTION_MANIFEST)?
-            .apply_auth(image, RegistryOperation::Pull)
-            .await?
-            .into_request_builder()
-            .send()
-            .await?;
-
-        if let Some(urls) = &layer.urls {
-            for url in urls {
-                if response.error_for_status_ref().is_ok() {
-                    break;
-                }
-
-                let url = Url::parse(url)
-                    .map_err(|e| OciDistributionError::UrlParseError(e.to_string()))?;
-
-                if url.scheme() == "http" || url.scheme() == "https" {
-                    // NOTE: we must not authenticate on additional URLs as those
-                    // can be abused to leak credentials or tokens.  Please
-                    // refer to CVE-2020-15157 for more information.
-                    response =
-                        RequestBuilderWrapper::from_client(self, |client| client.get(url.clone()))
-                            .apply_accept(MIME_TYPES_DISTRIBUTION_MANIFEST)?
-                            .into_request_builder()
-                            .send()
-                            .await?
-                }
-            }
-        }
+        let response = self.pull_blob_response(image, layer, None).await?;
 
         let mut stream = response.error_for_status()?.bytes_stream();
 
@@ -1072,17 +1056,58 @@ impl Client {
         &self,
         image: &Reference,
         layer: impl AsLayerDescriptor,
-    ) -> Result<impl Stream<Item = std::result::Result<bytes::Bytes, std::io::Error>>> {
+    ) -> Result<SizedStream> {
+        stream_from_response(self.pull_blob_response(image, layer, None).await?)
+    }
+
+    /// Stream a single layer from an OCI registry starting with a byte offset.
+    /// This can be used to continue downloading a layer after a network error.
+    ///
+    /// Returns [`Stream`](futures_util::Stream).
+    pub async fn pull_blob_stream_partial(
+        &self,
+        image: &Reference,
+        layer: impl AsLayerDescriptor,
+        offset: u64,
+    ) -> Result<BlobResponse> {
+        let response = self.pull_blob_response(image, layer, Some(offset)).await?;
+
+        let status = response.status();
+        match status {
+            StatusCode::OK => Ok(BlobResponse::Full(stream_from_response(response)?)),
+            StatusCode::PARTIAL_CONTENT => {
+                Ok(BlobResponse::Partial(stream_from_response(response)?))
+            }
+            _ => Err(OciDistributionError::ServerError {
+                code: status.as_u16(),
+                url: response.url().to_string(),
+                message: response.text().await?,
+            }),
+        }
+    }
+
+    /// Pull a single layer from an OCI registry.
+    async fn pull_blob_response(
+        &self,
+        image: &Reference,
+        layer: impl AsLayerDescriptor,
+        offset: Option<u64>,
+    ) -> Result<Response> {
         let layer = layer.as_layer_descriptor();
         let url = self.to_v2_blob_url(image.resolve_registry(), image.repository(), layer.digest);
 
-        let mut response = RequestBuilderWrapper::from_client(self, |client| client.get(&url))
+        let mut request = RequestBuilderWrapper::from_client(self, |client| client.get(&url))
             .apply_accept(MIME_TYPES_DISTRIBUTION_MANIFEST)?
             .apply_auth(image, RegistryOperation::Pull)
             .await?
-            .into_request_builder()
-            .send()
-            .await?;
+            .into_request_builder();
+        if let Some(offset) = offset {
+            request = request.header(
+                RANGE,
+                HeaderValue::from_str(&format!("bytes={offset}-")).unwrap(),
+            );
+        }
+        let mut response = request.send().await?;
 
         if let Some(urls) = &layer.urls {
             for url in urls {
@@ -1097,22 +1122,22 @@ impl Client {
                     // NOTE: we must not authenticate on additional URLs as those
                     // can be abused to leak credentials or tokens.  Please
                     // refer to CVE-2020-15157 for more information.
-                    response =
+                    request =
                         RequestBuilderWrapper::from_client(self, |client| client.get(url.clone()))
                             .apply_accept(MIME_TYPES_DISTRIBUTION_MANIFEST)?
-                            .into_request_builder()
-                            .send()
-                            .await?
+                            .into_request_builder();
+                    if let Some(offset) = offset {
+                        request = request.header(
+                            RANGE,
+                            HeaderValue::from_str(&format!("bytes={offset}-")).unwrap(),
+                        );
+                    }
+                    response = request.send().await?
                 }
             }
         }
 
-        let stream = response
-            .error_for_status()?
-            .bytes_stream()
-            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e));
-
-        Ok(stream)
+        Ok(response)
     }
 
     /// Begins a session to push an image to registry in a monolithical way
@@ -1511,6 +1536,19 @@ fn validate_registry_response(status: reqwest::StatusCode, body: &[u8], url: &st
     }
 }
 
+/// Converts a response into a stream
+fn stream_from_response(response: Response) -> Result<SizedStream> {
+    let content_length = response.content_length();
+    let stream = response
+        .error_for_status()?
+        .bytes_stream()
+        .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e));
+    Ok(SizedStream {
+        content_length,
+        stream: Box::pin(stream),
+    })
+}
+
 /// The request builder wrapper allows to be instantiated from a
 /// `Client` and allows composable operations on the request builder,
 /// to produce a `RequestBuilder` object that can be executed.
@@ -2567,7 +2605,8 @@ mod test {
                 .await
                 .expect("failed to pull blob stream");
 
-            AsyncReadExt::read_to_end(&mut StreamReader::new(layer_stream), &mut file)
+            assert_eq!(layer_stream.content_length, Some(layer0.size as u64));
+            AsyncReadExt::read_to_end(&mut StreamReader::new(layer_stream.stream), &mut file)
                 .await
                 .unwrap();