Allow for differentiated cache times

1 files changed, 35 insertions, 21 deletions

diff --git a/src/registry.rs b/src/registry.rs
index 45cf625..b18c1eb 100644
--- a/src/registry.rs
+++ b/src/registry.rs
@@ -6,6 +6,7 @@ use std::collections::HashMap;
 pub trait Registries {
     fn get_tags(&self, registry: &str, image: &str) -> Option<Vec<String>>;
     fn get_digest(&self, registry: &str, image: &str, tag: &str) -> Option<String>;
+    fn get_cache_ttl(&self, registry: &str) -> chrono::Duration;
 }
 
 #[derive(Debug)]
@@ -44,9 +45,10 @@ pub enum Credentials {
     Basic(String),
 }
 
-pub struct AuthInfo {
+pub struct Config {
     pub host: String,
     pub credentials: Credentials,
+    pub cache_ttl: chrono::Duration,
 }
 
 enum AuthStage {
@@ -54,15 +56,15 @@ enum AuthStage {
     Authorized(String),
 }
 
-struct AuthState {
-    info: AuthInfo,
-    stage: AuthStage,
+struct RegistryState {
+    config: Config,
+    auth_stage: AuthStage,
 }
 
 
-impl AuthState {
+impl RegistryState {
     fn add_to_request<T>(&self, req: ureq::RequestBuilder<T>) -> ureq::RequestBuilder<T> {
-        if let AuthStage::Authorized(x) = &self.stage {
+        if let AuthStage::Authorized(x) = &self.auth_stage {
             return req.header("Authorization", x);
         }
 
@@ -70,25 +72,25 @@ impl AuthState {
     }
 
     fn authenticate(&mut self, response: &ureq::http::Response<ureq::Body>) -> Result<(), String> {
-        match self.stage {
+        match self.auth_stage {
             AuthStage::Authorized(_) =>
                 // The token must have expired
-                self.stage = AuthStage::Idle,
+                self.auth_stage = AuthStage::Idle,
             AuthStage::Idle => {}
         }
 
         let auth_header = response.headers().get("www-authenticate").unwrap();
         match AuthMethod::from_header(auth_header) {
             Some(AuthMethod::Basic) => {
-                match &self.info.credentials {
+                match &self.config.credentials {
                     Credentials::Basic(creds) => {
-                        self.stage = AuthStage::Authorized(format!("Basic {}", creds));
+                        self.auth_stage = AuthStage::Authorized(format!("Basic {}", creds));
                         return Ok(());
                     }
                     Credentials::Anonymous => {
                         return Err(format!(
                             "Server {} requires Basic auth but configured as anonymous",
-                            self.info.host
+                            self.config.host
                         ));
                     }
                 }
@@ -101,7 +103,7 @@ impl AuthState {
                     .http_status_as_error(false)
                     .build();
 
-                if let Credentials::Basic(creds) = &self.info.credentials {
+                if let Credentials::Basic(creds) = &self.config.credentials {
                     request = request.header("Authorization", format!("Basic {}", creds));
                 }
 
@@ -111,7 +113,7 @@ impl AuthState {
                     .map_err(|x| {
                         format!(
                             "Server {} authentication request failed: {}",
-                            self.info.host,
+                            self.config.host,
                             x.to_string()
                         )
                     })?
@@ -120,7 +122,7 @@ impl AuthState {
                     .expect(
                         format!(
                             "Server {} responded with something non-string like",
-                            self.info.host
+                            self.config.host
                         )
                         .as_str(),
                     );
@@ -128,13 +130,13 @@ impl AuthState {
                 let body = body.parse::<tinyjson::JsonValue>().unwrap();
 
                 let token = body["token"].get::<String>().unwrap();
-                self.stage = AuthStage::Authorized(format!("Bearer {}", token));
+                self.auth_stage = AuthStage::Authorized(format!("Bearer {}", token));
                 return Ok(());
             }
             None => {
                 return Err(format!(
                     "Server {} provided us with a challenge, but we didn't understand it",
-                    self.info.host
+                    self.config.host
                 ))
             }
         }
@@ -148,16 +150,16 @@ enum RegistryError {
 }
 
 pub struct HttpRegistry {
-    auth: RefCell<HashMap<String, AuthState>>,
+    auth: RefCell<HashMap<String, RegistryState>>,
 }
 
 impl HttpRegistry {
-    pub fn new(infos: Vec<AuthInfo>) -> Self {
+    pub fn new(infos: Vec<Config>) -> Self {
         let mut states = HashMap::new();
         for info in infos {
-            states.insert(info.host.clone(), AuthState {
-                info: info,
-                stage: AuthStage::Idle,
+            states.insert(info.host.clone(), RegistryState {
+                config: info,
+                auth_stage: AuthStage::Idle,
             });
         }
         return HttpRegistry { auth: RefCell::new(states) };
@@ -311,6 +313,14 @@ impl Registries for HttpRegistry {
             .to_string();
         return Some(digest);
     }
+
+    fn get_cache_ttl(&self, registry: &str) -> chrono::Duration {
+        let auth = self.auth.borrow();
+        if let Some(state) = auth.get(registry) {
+            return state.config.cache_ttl;
+        }
+        return chrono::Duration::minutes(1440);
+    }
 }
 
 #[cfg(test)]
@@ -384,6 +394,10 @@ impl Registries for StubRegistry {
         }
         return None;
     }
+
+    fn get_cache_ttl(&self, _registry: &str) -> chrono::Duration {
+        return chrono::Duration::minutes(1440);
+    }
 }
 
 #[cfg(test)]