From 29b64781969f944a4aa273d181d85689e156dc52 Mon Sep 17 00:00:00 2001
From: Phuong Nguyen <ChaoticTempest@users.noreply.github.com>
Date: Thu, 3 Aug 2023 16:38:24 -0700
Subject: [PATCH] Add 5 retries before moving on to rotate keys (#239)

---
 integration-tests/src/containers.rs     |  8 ++++----
 integration-tests/tests/mpc/positive.rs | 17 ++++++++++++++++-
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/integration-tests/src/containers.rs b/integration-tests/src/containers.rs
index e61d6e27a..6b78efb4d 100644
--- a/integration-tests/src/containers.rs
+++ b/integration-tests/src/containers.rs
@@ -382,7 +382,7 @@ pub struct SignerNodeApi {
     pub address: String,
     pub node_id: usize,
     pub sk_share: ExpandedKeyPair,
-    pub cipher_key: Aes256Gcm,
+    pub cipher_key: GenericArray<u8, U32>,
     pub gcp_project_id: String,
     pub gcp_datastore_local_url: String,
 }
@@ -464,7 +464,7 @@ impl<'a> SignerNode<'a> {
             address: self.local_address.clone(),
             node_id: self.node_id,
             sk_share: self.sk_share.clone(),
-            cipher_key: Aes256Gcm::new(&self.cipher_key),
+            cipher_key: self.cipher_key,
             gcp_project_id: self.gcp_project_id.clone(),
             gcp_datastore_local_url: self.gcp_datastore_local_url.clone(),
         }
@@ -492,12 +492,12 @@ impl SignerNodeApi {
         .await?;
 
         let new_cipher = Aes256Gcm::new(new_cipher_key);
-        let old_cipher = &self.cipher_key;
+        let old_cipher = Aes256Gcm::new(&self.cipher_key);
 
         // Do inplace rotation of node key
         mpc_recovery::sign_node::migration::rotate_cipher(
             self.node_id,
-            old_cipher,
+            &old_cipher,
             &new_cipher,
             &gcp_service,
             &gcp_service,
diff --git a/integration-tests/tests/mpc/positive.rs b/integration-tests/tests/mpc/positive.rs
index aaad68964..e12513f4e 100644
--- a/integration-tests/tests/mpc/positive.rs
+++ b/integration-tests/tests/mpc/positive.rs
@@ -218,7 +218,22 @@ async fn test_rotate_node_keys() -> anyhow::Result<()> {
                 .collect::<HashMap<_, _>>();
 
             // Generate a new set of ciphers to rotate out each node:
-            let mpc_recovery::GenerateResult { secrets, .. } = mpc_recovery::generate(3);
+            let mut counter = 0;
+            let mpc_recovery::GenerateResult { secrets, .. } = loop {
+                let result = mpc_recovery::generate(3);
+                let all_diff = result.secrets.iter().zip(ctx.signer_nodes.iter()).all(|((_, new_cipher), signer_node)| {
+                    signer_node.cipher_key != *new_cipher
+                });
+
+                if all_diff {
+                    break result;
+                }
+
+                counter += 1;
+                if counter == 5 {
+                    panic!("Failed to generate a new set of ciphers after 5 tries");
+                }
+            };
 
             let mut ciphers = HashMap::new();
             // Rotate out with new the cipher.