From def2efa0d112d1e64387485b9ba70982bb32dee2 Mon Sep 17 00:00:00 2001
From: Sergio <77530549+sergi0g@users.noreply.github.com>
Date: Sat, 31 Aug 2024 21:24:56 +0300
Subject: [PATCH] Add better error handling (mostly because I have bad
 internet)

---
 src/registry.rs | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/registry.rs b/src/registry.rs
index 289b7fd..da70655 100644
--- a/src/registry.rs
+++ b/src/registry.rs
@@ -2,7 +2,7 @@ use std::sync::Mutex;
 
 use json::JsonValue;
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
-use ureq::Error;
+use ureq::{Error, ErrorKind};
 
 use http_auth::parse_challenges;
 
@@ -56,8 +56,20 @@ pub fn get_latest_digest(image: &Image, token: Option<&String>, config: &JsonVal
                 digest: None,
                 ..image.clone()
             }
-        }
-        Err(ureq::Error::Transport(e)) => error!("Failed to send request!\n{}", e),
+        },
+        Err(Error::Transport(error)) => {
+            match error.kind() {
+                ErrorKind::Dns => {
+                    warn!("Failed to lookup the IP of the registry, retrying.");
+                    return get_latest_digest(image, token, config)
+                }, // If something goes really wrong, this can get stuck in a loop
+                ErrorKind::ConnectionFailed => {
+                    warn!("Connection probably timed out, retrying.");
+                    return get_latest_digest(image, token, config)
+                }, // Same here
+                _ => error!("Failed to retrieve image digest\n{}!", error)
+            }
+        },
     };
     match raw_response.header("docker-content-digest") {
         Some(digest) => Image {
@@ -83,7 +95,7 @@ pub fn get_latest_digests(images: Vec<&Image>, token: Option<&String>, config: &
 
 pub fn get_token(images: Vec<&Image>, auth_url: &str, credentials: &Option<String>) -> String {
     let mut final_url = auth_url.to_owned();
-    for image in images {
+    for image in &images {
         final_url = format!("{}&scope=repository:{}:pull", final_url, image.repository);
     }
     let mut base_request = ureq::get(&final_url).set("Accept", "application/vnd.oci.image.index.v1+json"); // Seems to be unnecesarry. Will probably remove in the future
@@ -99,6 +111,19 @@ pub fn get_token(images: Vec<&Image>, auth_url: &str, credentials: &Option<Strin
                 error!("Failed to parse response into string!\n{}", e)
             }
         },
+        Err(Error::Transport(error)) => {
+            match error.kind() {
+                ErrorKind::Dns => {
+                    warn!("Failed to lookup the IP of the registry, retrying.");
+                    return get_token(images, auth_url, credentials)
+                }, // If something goes really wrong, this can get stuck in a loop
+                ErrorKind::ConnectionFailed => {
+                    warn!("Connection probably timed out, retrying.");
+                    return get_token(images, auth_url, credentials)
+                }, // Same here
+                _ => error!("Token request failed\n{}!", error)
+            }
+        },
         Err(e) => {
             error!("Token request failed!\n{}", e)
         }