From 486961f3d3d956f3666bcd72338bd0c5c8d56234 Mon Sep 17 00:00:00 2001
From: Gerald Pinder <gmpinder@gmail.com>
Date: Sat, 21 Sep 2024 18:57:57 -0400
Subject: [PATCH] fix: Make sigstore driver more resilient to network errors

---
 process/drivers/sigstore_driver.rs | 50 +++++++++++++++++-------------
 utils/src/lib.rs                   |  9 +++---
 2 files changed, 34 insertions(+), 25 deletions(-)

diff --git a/process/drivers/sigstore_driver.rs b/process/drivers/sigstore_driver.rs
index 65ba986..5b93a93 100644
--- a/process/drivers/sigstore_driver.rs
+++ b/process/drivers/sigstore_driver.rs
@@ -13,6 +13,7 @@ use super::{
 use blue_build_utils::{
     constants::{COSIGN_PRIV_PATH, COSIGN_PUB_PATH},
     credentials::Credentials,
+    retry,
 };
 use log::{debug, trace};
 use miette::{bail, miette, Context, IntoDiagnostic};
@@ -133,10 +134,11 @@ impl SigningDriver for SigstoreDriver {
         let auth = Auth::Basic(username.clone(), password.clone());
         debug!("Credentials retrieved");
 
-        let (cosign_signature_image, source_image_digest) = RT
-            .block_on(client.triangulate(&image_digest, &auth))
-            .into_diagnostic()
-            .with_context(|| format!("Failed to triangulate image {image_digest}"))?;
+        let (cosign_signature_image, source_image_digest) = retry(2, 5, || {
+            RT.block_on(client.triangulate(&image_digest, &auth))
+                .into_diagnostic()
+                .with_context(|| format!("Failed to triangulate image {image_digest}"))
+        })?;
         debug!("Triangulating image");
         trace!("{cosign_signature_image}, {source_image_digest}");
 
@@ -148,15 +150,19 @@ impl SigningDriver for SigstoreDriver {
         debug!("Created signing layer");
 
         debug!("Pushing signature");
-        RT.block_on(client.push_signature(
-            None,
-            &auth,
-            &cosign_signature_image,
-            vec![signature_layer],
-        ))
-        .into_diagnostic()
-        .with_context(|| {
-            format!("Failed to push signature {cosign_signature_image} for image {image_digest}")
+        retry(2, 5, || {
+            RT.block_on(client.push_signature(
+                None,
+                &auth,
+                &cosign_signature_image,
+                vec![signature_layer.clone()],
+            ))
+            .into_diagnostic()
+            .with_context(|| {
+                format!(
+                    "Failed to push signature {cosign_signature_image} for image {image_digest}"
+                )
+            })
         })?;
         debug!("Successfully pushed signature");
 
@@ -187,21 +193,23 @@ impl SigningDriver for SigstoreDriver {
             PublicKeyVerifier::new(pub_key.as_bytes(), &signing_scheme).into_diagnostic()?;
         let verification_constraints: VerificationConstraintVec = vec![Box::new(verifier)];
 
-        let auth = Auth::Anonymous;
-        let (cosign_signature_image, source_image_digest) = RT
-            .block_on(client.triangulate(&image_digest, &auth))
-            .into_diagnostic()
-            .with_context(|| format!("Failed to triangulate image {image_digest}"))?;
         debug!("Triangulating image");
+        let auth = Auth::Anonymous;
+        let (cosign_signature_image, source_image_digest) = retry(2, 5, || {
+            RT.block_on(client.triangulate(&image_digest, &auth))
+                .into_diagnostic()
+                .with_context(|| format!("Failed to triangulate image {image_digest}"))
+        })?;
         trace!("{cosign_signature_image}, {source_image_digest}");
 
-        let trusted_layers = RT
-            .block_on(client.trusted_signature_layers(
+        let trusted_layers = retry(2, 5, || {
+            RT.block_on(client.trusted_signature_layers(
                 &auth,
                 &source_image_digest,
                 &cosign_signature_image,
             ))
-            .into_diagnostic()?;
+            .into_diagnostic()
+        })?;
 
         sigstore::cosign::verify_constraints(&trusted_layers, verification_constraints.iter())
             .map_err(
diff --git a/utils/src/lib.rs b/utils/src/lib.rs
index 8e52701..775a8cf 100644
--- a/utils/src/lib.rs
+++ b/utils/src/lib.rs
@@ -21,7 +21,7 @@ use blake2::{
 };
 use chrono::Local;
 use format_serde_error::SerdeError;
-use log::trace;
+use log::{trace, warn};
 use miette::{miette, Context, IntoDiagnostic, Result};
 
 use crate::constants::CONTAINER_FILE;
@@ -72,16 +72,17 @@ pub fn serde_yaml_err(contents: &str) -> impl Fn(serde_yaml::Error) -> SerdeErro
 ///
 /// # Errors
 /// Will error when retries have been expended.
-pub fn retry<V, F>(mut retries: u8, delay_secs: u64, f: F) -> miette::Result<V>
+pub fn retry<V, F>(mut retries: u8, delay_secs: u64, mut f: F) -> miette::Result<V>
 where
-    F: Fn() -> miette::Result<V>,
+    F: FnMut() -> miette::Result<V>,
 {
     loop {
         match f() {
             Ok(v) => return Ok(v),
             Err(e) if retries == 0 => return Err(e),
-            _ => {
+            Err(e) => {
                 retries -= 1;
+                warn!("Failed operation, will retry {retries} more time(s). Error:\n{e:?}");
                 thread::sleep(Duration::from_secs(delay_secs));
             }
         };