impl From<ureq::Error> for ocipkg::Error, drop CheckResponse

termoshtt · Sep 10, 2022 · faf4c33 · faf4c33
1 parent 75f47f2
commit faf4c33
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 42 deletions.
diff --git a/ocipkg/src/distribution/client.rs b/ocipkg/src/distribution/client.rs
@@ -30,10 +30,9 @@ impl Client {
 
     fn call(&mut self, req: ureq::Request) -> Result<ureq::Response> {
         if let Some(token) = &self.token {
-            return req
+            return Ok(req
                 .set("Authorization", &format!("Bearer {}", token))
-                .call()
-                .check_response();
+                .call()?);
         }
 
         // Try get token
@@ -55,24 +54,16 @@ impl Client {
         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))
-        } else {
-            req
-        }
-    }
-
     fn get(&self, url: &Url) -> ureq::Request {
-        self.add_auth_header(self.agent.get(url.as_str()))
+        self.agent.get(url.as_str())
     }
 
     fn put(&self, url: &Url) -> ureq::Request {
-        self.add_auth_header(self.agent.put(url.as_str()))
+        self.agent.put(url.as_str())
     }
 
     fn post(&self, url: &Url) -> ureq::Request {
-        self.add_auth_header(self.agent.post(url.as_str()))
+        self.agent.post(url.as_str())
     }
 
     /// Get tags of `<name>` repository.
@@ -130,8 +121,12 @@ impl Client {
         let res = self
             .put(&url)
             .set("Content-Type", &MediaType::ImageManifest.to_string())
-            .send_bytes(&buf)
-            .check_response()?;
+            .set(
+                "Authorization",
+                &format!("Bearer {}", self.token.as_ref().unwrap()),
+                // Authorization must be done while blobs push
+            )
+            .send_bytes(&buf)?;
         let loc = res
             .header("Location")
             .expect("Location header is lacked in OCI registry response");
@@ -180,37 +175,19 @@ impl Client {
             .query("digest", &digest.to_string())
             .set("Content-Length", &blob.len().to_string())
             .set("Content-Type", "application/octet-stream")
-            .send_bytes(blob)
-            .check_response()?;
+            .set(
+                "Authorization",
+                &format!("Bearer {}", self.token.as_ref().unwrap()),
+                // Authorization must be done while the first POST
+            )
+            .send_bytes(blob)?;
         let loc = res
             .header("Location")
             .expect("Location header is lacked in OCI registry response");
         Ok(Url::parse(loc).or_else(|_| self.url.join(loc))?)
     }
 }
 
-trait CheckResponse {
-    fn check_response(self) -> Result<ureq::Response>;
-}
-
-impl CheckResponse for std::result::Result<ureq::Response, ureq::Error> {
-    fn check_response(self) -> Result<ureq::Response> {
-        match self {
-            Ok(res) => Ok(res),
-            Err(ureq::Error::Status(status, res)) => {
-                if status == 401 {
-                    if let Some(msg) = res.header("www-authenticate") {
-                        log::error!("Server returns WWW-Authenticate header: {}", msg);
-                    }
-                }
-                let err = res.into_json::<ErrorResponse>()?;
-                Err(Error::RegistryError(err))
-            }
-            Err(ureq::Error::Transport(e)) => Err(Error::NetworkError(e)),
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/ocipkg/src/error.rs b/ocipkg/src/error.rs
@@ -1,5 +1,5 @@
 use crate::Digest;
-use oci_spec::OciSpecError;
+use oci_spec::{distribution::ErrorResponse, OciSpecError};
 use std::path::PathBuf;
 
 #[derive(Debug, thiserror::Error)]
@@ -46,7 +46,7 @@ pub enum Error {
     #[error(transparent)]
     NetworkError(#[from] ureq::Transport),
     #[error(transparent)]
-    RegistryError(#[from] oci_spec::distribution::ErrorResponse),
+    RegistryError(#[from] ErrorResponse),
     #[error("Authorization failed: {0}")]
     AuthorizationFailed(url::Url),
     #[error("Unsupported WWW-Authentication header: {0}")]
@@ -81,3 +81,15 @@ impl From<walkdir::Error> for Error {
         Self::UnknownIo(e.into())
     }
 }
+
+impl From<ureq::Error> for Error {
+    fn from(e: ureq::Error) -> Self {
+        match e {
+            ureq::Error::Status(_status, res) => match res.into_json::<ErrorResponse>() {
+                Ok(err) => Error::RegistryError(err),
+                Err(e) => Error::UnknownIo(e),
+            },
+            ureq::Error::Transport(e) => Error::NetworkError(e),
+        }
+    }
+}