Client::call

ocipkg/src/distribution/client.rs

@@ -10,22 +10,51 @@ pub struct Client {
     url: Url,
     /// Name of repository
     name: Name,
-    /// Authorization token
+    /// Loaded authentication info from filesystem
+    auth: StoredAuth,
+    /// Cached token
     token: Option<String>,
 }
 
 impl Client {
     pub fn new(url: Url, name: Name) -> Result<Self> {
         let auth = StoredAuth::load_all()?;
-        let token = auth.get_token(&url)?;
         Ok(Client {
             agent: ureq::Agent::new(),
             url,
             name,
-            token,
+            auth,
+            token: None,
         })
     }
 
+    fn call(&mut self, req: ureq::Request) -> Result<ureq::Response> {
+        if let Some(token) = &self.token {
+            return req
+                .set("Authorization", &format!("Bearer {}", token))
+                .call()
+                .check_response();
+        }
+
+        // Try get token
+        let try_req = req.clone();
+        let www_auth = match try_req.call() {
+            Ok(res) => return Ok(res),
+            Err(ureq::Error::Status(status, res)) => {
+                if status == 401 && res.has("www-authenticate") {
+                    res.header("www-authenticate").unwrap().to_string()
+                } else {
+                    let err = res.into_json::<ErrorResponse>()?;
+                    return Err(Error::RegistryError(err));
+                }
+            }
+            Err(ureq::Error::Transport(e)) => return Err(Error::NetworkError(e)),
+        };
+        let challenge = AuthChallenge::from_header(&www_auth)?;
+        self.token = Some(self.auth.challenge(&challenge)?);
+        return self.call(req);
+    }
+
     fn add_auth_header(&self, req: ureq::Request) -> ureq::Request {
         if let Some(token) = &self.token {
             req.set("Authorization", &format!("Bearer {}", token))
@@ -53,9 +82,9 @@ impl Client {
     /// ```
     ///
     /// See [corresponding OCI distribution spec document](https://github.com/opencontainers/distribution-spec/blob/main/spec.md#content-discovery) for detail.
-    pub fn get_tags(&self) -> Result<Vec<String>> {
+    pub fn get_tags(&mut self) -> Result<Vec<String>> {
         let url = self.url.join(&format!("/v2/{}/tags/list", self.name))?;
-        let res = self.get(&url).call().check_response()?;
+        let res = self.call(self.get(&url))?;
         let tag_list = res.into_json::<TagList>()?;
         Ok(tag_list.tags().to_vec())
     }
@@ -67,22 +96,18 @@ impl Client {
     /// ```
     ///
     /// See [corresponding OCI distribution spec document](https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pulling-manifests) for detail.
-    pub fn get_manifest(&self, reference: &Reference) -> Result<ImageManifest> {
+    pub fn get_manifest(&mut self, reference: &Reference) -> Result<ImageManifest> {
         let url = self
             .url
             .join(&format!("/v2/{}/manifests/{}", self.name, reference))?;
-        let res = self
-            .get(&url)
-            .set(
-                "Accept",
-                &format!(
-                    "{}, {}",
-                    MediaType::ImageManifest.to_docker_v2s2().unwrap(),
-                    MediaType::ImageManifest,
-                ),
-            )
-            .call()
-            .check_response()?;
+        let res = self.call(self.get(&url).set(
+            "Accept",
+            &format!(
+                "{}, {}",
+                MediaType::ImageManifest.to_docker_v2s2().unwrap(),
+                MediaType::ImageManifest,
+            ),
+        ))?;
         let manifest = ImageManifest::from_reader(res.into_reader())?;
         Ok(manifest)
     }
@@ -120,11 +145,11 @@ impl Client {
     /// ```
     ///
     /// See [corresponding OCI distribution spec document](https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pulling-blobs) for detail.
-    pub fn get_blob(&self, digest: &Digest) -> Result<Vec<u8>> {
+    pub fn get_blob(&mut self, digest: &Digest) -> Result<Vec<u8>> {
         let url = self
             .url
             .join(&format!("/v2/{}/blobs/{}", self.name.as_str(), digest,))?;
-        let res = self.get(&url).call().check_response()?;
+        let res = self.call(self.get(&url))?;
         let mut bytes = Vec::new();
         res.into_reader().read_to_end(&mut bytes)?;
         Ok(bytes)
@@ -139,11 +164,11 @@ impl Client {
     /// and following `PUT` to URL obtained by `POST`.
     ///
     /// See [corresponding OCI distribution spec document](https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pushing-manifests) for detail.
-    pub fn push_blob(&self, blob: &[u8]) -> Result<Url> {
+    pub fn push_blob(&mut self, blob: &[u8]) -> Result<Url> {
         let url = self
             .url
             .join(&format!("/v2/{}/blobs/uploads/", self.name))?;
-        let res = self.post(&url).call().check_response()?;
+        let res = self.call(self.post(&url))?;
         let loc = res
             .header("Location")
             .expect("Location header is lacked in OCI registry response");
@@ -205,7 +230,7 @@ mod tests {
     #[test]
     #[ignore]
     fn get_tags() -> Result<()> {
-        let client = Client::new(test_url(), test_name())?;
+        let mut client = Client::new(test_url(), test_name())?;
         let mut tags = client.get_tags()?;
         tags.sort_unstable();
         assert_eq!(
@@ -218,7 +243,7 @@ mod tests {
     #[test]
     #[ignore]
     fn get_images() -> Result<()> {
-        let client = Client::new(test_url(), test_name())?;
+        let mut client = Client::new(test_url(), test_name())?;
         for tag in ["tag1", "tag2", "tag3"] {
             let manifest = client.get_manifest(&Reference::new(tag)?)?;
             for layer in manifest.layers() {
@@ -232,7 +257,7 @@ mod tests {
     #[test]
     #[ignore]
     fn push_blob() -> Result<()> {
-        let client = Client::new(test_url(), test_name())?;
+        let mut client = Client::new(test_url(), test_name())?;
         let url = client.push_blob("test string".as_bytes())?;
         dbg!(url);
         Ok(())

ocipkg/src/distribution/mod.rs

@@ -22,7 +22,7 @@ pub fn push_image(path: &Path) -> Result<()> {
     let mut f = fs::File::open(&path)?;
     let mut ar = crate::image::Archive::new(&mut f);
     for (image_name, manifest) in ar.get_manifests()? {
-        let client = Client::new(image_name.registry_url()?, image_name.name)?;
+        let mut client = Client::new(image_name.registry_url()?, image_name.name)?;
         for layer in manifest.layers() {
             let digest = Digest::new(layer.digest())?;
             let mut entry = ar.get_blob(&digest)?;
@@ -45,7 +45,7 @@ pub fn get_image(image_name: &ImageName) -> Result<()> {
     let ImageName {
         name, reference, ..
     } = image_name;
-    let client = Client::new(image_name.registry_url()?, name.clone())?;
+    let mut client = Client::new(image_name.registry_url()?, name.clone())?;
     let manifest = client.get_manifest(reference)?;
     let dest = crate::local::image_dir(image_name)?;
     log::info!("Get {} into {}", image_name, dest.display());