Add configurable retries for the pin

Support for an optional field for setting the number of retries. This
can be either a positive number or 'inifinity'. If no value is specify,
then the pin tries 10 times all the servers.

Signed-off-by: Alice Frosi <[email protected]>

Author: alicefr

cli/src/main.rs

@@ -111,15 +111,18 @@ struct ClevisHeader {
     servers: Vec<Server>,
     path: String,
     initdata: Option<String>,
+    #[serde(default)]
+    num_retries: Option<NumRetries>,
 }
 
 fn fetch_and_prepare_jwk<E: CommandExecutor>(
     servers: &[Server],
     path: &str,
     initdata: Option<String>,
+    num_retries: &NumRetries,
     executor: &E,
 ) -> Result<Jwk> {
-    let key = fetch_luks_key(servers, path, initdata, executor)?;
+    let key = fetch_luks_key(servers, path, initdata, num_retries, executor)?;
     let key = String::from_utf8(
         general_purpose::STANDARD
             .decode(&key)
@@ -160,7 +163,8 @@ fn encrypt(config: &str) -> Result<()> {
     io::stdin().read_to_end(&mut input)?;
 
     let executor = RealCommandExecutor;
-    let jwk = fetch_and_prepare_jwk(&config.servers, &config.path, initdata.clone(), &executor)?;
+    let num_retries = config.num_retries.as_ref().unwrap_or(&NumRetries::Finite(10));
+    let jwk = fetch_and_prepare_jwk(&config.servers, &config.path, initdata.clone(), num_retries, &executor)?;
 
     eprintln!("{}", jwk);
     let encrypter = Dir
@@ -172,6 +176,7 @@ fn encrypt(config: &str) -> Result<()> {
         servers: config.servers.clone(),
         path: config.path,
         initdata,
+        num_retries: config.num_retries,
     };
 
     let mut hdr = josekit::jwe::JweHeader::new();
@@ -207,10 +212,12 @@ fn decrypt() -> Result<()> {
     eprintln!("Decrypt with header: {:?}", hdr_clevis);
 
     let executor = RealCommandExecutor;
+    let num_retries = hdr_clevis.num_retries.as_ref().unwrap_or(&NumRetries::Finite(10));
     let decrypter_jwk = fetch_and_prepare_jwk(
         &hdr_clevis.servers,
         &hdr_clevis.path,
         hdr_clevis.initdata,
+        num_retries,
         &executor,
     )?;
 
@@ -227,55 +234,87 @@ fn decrypt() -> Result<()> {
     Ok(())
 }
 
+fn try_fetch_from_servers<E: CommandExecutor>(
+    servers: &[Server],
+    path: &str,
+    initdata: &Option<String>,
+    executor: &E,
+) -> Option<String> {
+    for (index, server) in servers.iter().enumerate() {
+        eprintln!("Trying URL {}/{}: {}", index + 1, servers.len(), server.url);
+        match executor.try_fetch_luks_key(&server.url, path, &server.cert, initdata.clone()) {
+            Ok(key) => {
+                eprintln!("Successfully fetched LUKS key from URL: {}", server.url);
+                return Some(key);
+            }
+            Err(e) => {
+                eprintln!("Error with URL {}: {}", server.url, e);
+            }
+        }
+    }
+    None
+}
+
 fn fetch_luks_key<E: CommandExecutor>(
     servers: &[Server],
     path: &str,
     initdata: Option<String>,
+    num_retries: &NumRetries,
     executor: &E,
 ) -> Result<String> {
-    const MAX_ATTEMPTS: u32 = 3;
     const DELAY: Duration = Duration::from_secs(5);
 
     if servers.is_empty() {
         return Err(anyhow!("No URLs provided"));
     }
 
-    (1..=MAX_ATTEMPTS)
-        .find_map(|attempt| {
-            eprintln!(
-                "Attempting to fetch LUKS key (attempt {}/{})",
-                attempt, MAX_ATTEMPTS
-            );
-
-            for (index, server) in servers.iter().enumerate() {
-                eprintln!("Trying URL {}/{}: {}", index + 1, servers.len(), server.url);
-                match executor.try_fetch_luks_key(&server.url, path, &server.cert, initdata.clone())
-                {
-                    Ok(key) => {
-                        eprintln!("Successfully fetched LUKS key from URL: {}", server.url);
+    match num_retries {
+        NumRetries::Finite(max_attempts) => {
+            (1..=*max_attempts)
+                .find_map(|attempt| {
+                    eprintln!(
+                        "Attempting to fetch LUKS key (attempt {}/{})",
+                        attempt, max_attempts
+                    );
+
+                    if let Some(key) = try_fetch_from_servers(servers, path, &initdata, executor) {
                         return Some(Ok(key));
                     }
-                    Err(e) => {
-                        eprintln!("Error with URL {}: {}", server.url, e);
+
+                    if attempt < *max_attempts {
+                        eprintln!(
+                            "All URLs failed for attempt {}. Retrying in {:?} seconds...",
+                            attempt, DELAY
+                        );
+                        thread::sleep(DELAY);
                     }
+                    None
+                })
+                .unwrap_or_else(|| {
+                    Err(anyhow!(
+                        "Failed to fetch the LUKS key from all URLs after {} attempts",
+                        max_attempts
+                    ))
+                })
+        }
+        NumRetries::Infinity => {
+            let mut attempt = 0;
+            loop {
+                attempt += 1;
+                eprintln!("Attempting to fetch LUKS key (attempt {})", attempt);
+
+                if let Some(key) = try_fetch_from_servers(servers, path, &initdata, executor) {
+                    return Ok(key);
                 }
-            }
 
-            if attempt < MAX_ATTEMPTS {
                 eprintln!(
                     "All URLs failed for attempt {}. Retrying in {:?} seconds...",
                     attempt, DELAY
                 );
                 thread::sleep(DELAY);
             }
-            None
-        })
-        .unwrap_or_else(|| {
-            Err(anyhow!(
-                "Failed to fetch the LUKS key from all URLs after {} attempts",
-                MAX_ATTEMPTS
-            ))
-        })
+        }
+    }
 }
 
 /// Clevis PIN for Trustee
@@ -323,7 +362,8 @@ mod tests {
             cert: String::new(),
         }];
 
-        let result = fetch_luks_key(&servers, "/test/path", None, &mock);
+        let num_retries = NumRetries::Finite(3);
+        let result = fetch_luks_key(&servers, "/test/path", None, &num_retries, &mock);
 
         assert!(result.is_ok());
         assert_eq!(result.unwrap(), "test_luks_key_12345");
@@ -340,7 +380,8 @@ mod tests {
             cert: String::new(),
         }];
 
-        let result = fetch_luks_key(&servers, "/test/path", None, &mock);
+        let num_retries = NumRetries::Finite(3);
+        let result = fetch_luks_key(&servers, "/test/path", None, &num_retries, &mock);
 
         assert!(result.is_err());
         assert_eq!(

lib/src/lib.rs

@@ -3,9 +3,83 @@
 //
 // SPDX-License-Identifier: MIT
 
-use serde::{Deserialize, Serialize};
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use std::collections::HashMap;
 
+#[derive(Debug, Clone, PartialEq)]
+pub enum NumRetries {
+    Finite(u32),
+    Infinity,
+}
+
+impl Serialize for NumRetries {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        match self {
+            NumRetries::Finite(n) => serializer.serialize_u32(*n),
+            NumRetries::Infinity => serializer.serialize_str("infinity"),
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for NumRetries {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        struct NumRetriesVisitor;
+
+        impl<'de> serde::de::Visitor<'de> for NumRetriesVisitor {
+            type Value = NumRetries;
+
+            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
+                formatter.write_str("a positive number (>= 1) or the string 'infinity'")
+            }
+
+            fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                if value == 0 {
+                    return Err(E::custom("number must be at least 1, got: 0"));
+                }
+                if value > u32::MAX as u64 {
+                    return Err(E::custom(format!("number too large: {}", value)));
+                }
+                Ok(NumRetries::Finite(value as u32))
+            }
+
+            fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                if value <= 0 {
+                    return Err(E::custom(format!("number must be at least 1, got: {}", value)));
+                }
+                if value > u32::MAX as i64 {
+                    return Err(E::custom(format!("number too large: {}", value)));
+                }
+                Ok(NumRetries::Finite(value as u32))
+            }
+
+            fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                if value == "infinity" {
+                    Ok(NumRetries::Infinity)
+                } else {
+                    Err(E::custom(format!("expected 'infinity', got: '{}'", value)))
+                }
+            }
+        }
+
+        deserializer.deserialize_any(NumRetriesVisitor)
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct Server {
     pub url: String,
@@ -17,6 +91,7 @@ pub struct Config {
     pub servers: Vec<Server>,
     pub path: String,
     pub initdata: Option<String>,
+    pub num_retries: Option<NumRetries>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]