Build challenge urls safely according to www-authenticate header

src/registry.rs

@@ -13,34 +13,81 @@ impl Default for Registry {
 }
 
 impl Registry {
-    pub fn pull(&mut self, image: &ImageIdentifier, destination: &str) {
+    fn build_challenge_url(auth_header: &str) -> reqwest::Url {
-        // Perform the little auth dance
+        let challenge_part = auth_header.split(" ").nth(1).unwrap();
-        let auth_url = format!("https://auth.docker.io/token?service=registry.docker.io&scope=repository%3A{}%2F{}%3Apull", image.author, image.name);
+        let mut field_map = std::collections::HashMap::<&str, &str>::new();
-        let auth: serde_json::Value = self
+        for field in challenge_part.split(",") {
-            .http_client
+            let mut field_split = field.split("=");
-            .get(&auth_url)
+            let key = field_split.next().unwrap();
-            .send()
+            let value = field_split.next().unwrap();
-            .unwrap()
+            field_map.insert(key, &value[1..value.len() - 1]);
-            .json()
+        }
-            .unwrap();
+        let realm = field_map.get("realm").unwrap().to_owned();
-        let access_token = auth["token"].as_str().unwrap();
+        field_map.remove("realm").unwrap();
+        let mut url = reqwest::Url::parse(realm).unwrap();
 
-        // Download the image manifest
+        {
-        let auth_header = format!("Bearer {}", access_token);
+            let mut query_pairs = url.query_pairs_mut();
+            for (key, value) in field_map {
+                query_pairs.append_pair(key, value);
+            }
+        }
+
+        return url;
+    }
+
+    fn fetch_manifest(
+        &self,
+        image: &ImageIdentifier,
+        additional_headers: Option<reqwest::header::HeaderMap>,
+    ) -> reqwest::blocking::Response {
         let image_url = format!(
             "https://registry.hub.docker.com/v2/{}/{}/manifests/{}",
             image.author, image.name, image.tag
         );
-        let image_manifest: serde_json::Value = self
+        let image_manifest = self
             .http_client
             .get(&image_url)
-            .header(reqwest::header::AUTHORIZATION, auth_header.to_owned())
+            .headers(additional_headers.unwrap_or_default())
             .send()
-            .unwrap()
-            .json()
             .unwrap();
 
-        // Download the image layers and extracts them
+        return image_manifest;
+    }
+
+    pub fn pull(&mut self, image: &ImageIdentifier, destination: &str) {
+        let mut header_map = reqwest::header::HeaderMap::new();
+
+        let mut manifest_resp = self.fetch_manifest(image, None);
+
+        // Perform the little auth dance if necessary
+        if manifest_resp.status() != reqwest::StatusCode::OK {
+            let auth_header = manifest_resp
+                .headers()
+                .get(reqwest::header::WWW_AUTHENTICATE)
+                .unwrap()
+                .to_str()
+                .unwrap();
+
+            let challenge_url = Registry::build_challenge_url(auth_header);
+            let challenge_body: serde_json::Value = self
+                .http_client
+                .get(challenge_url)
+                .send()
+                .unwrap()
+                .json()
+                .unwrap();
+            let access_token = challenge_body["token"].as_str().unwrap();
+            header_map.append(
+                reqwest::header::AUTHORIZATION,
+                format!("Bearer {}", access_token).parse().unwrap(),
+            );
+
+            manifest_resp = self.fetch_manifest(image, Some(header_map.to_owned()));
+        }
+
+        let image_manifest: serde_json::Value = manifest_resp.json().unwrap();
+
         let temp_path = std::env::temp_dir();
         for layer in image_manifest["fsLayers"].as_array().unwrap() {
             let digest = layer["blobSum"].as_str().unwrap();
@@ -51,7 +98,7 @@ impl Registry {
             let blob = self
                 .http_client
                 .get(&blob_url)
-                .header(reqwest::header::AUTHORIZATION, auth_header.to_owned())
+                .headers(header_map.to_owned())
                 .send()
                 .unwrap()
                 .bytes()
@@ -112,7 +159,7 @@ mod tests {
             ImageIdentifier {
                 author: "library".to_string(),
                 name: "alpine".to_string(),
-                tag: "latest".to_string()
+                tag: "latest".to_string(),
             }
         );