Merge pull request #892 from gthao313/ephemeral-cluster

sonobuoy: re-set the assume role credentials if it expires
bottlerocket-os · Mar 18, 2024 · c778fb6 · c778fb6
2 parents f1b58de + 2a84be3
commit c778fb6
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 9 deletions.
diff --git a/bottlerocket/agents/src/bin/k8s-workload-agent/main.rs b/bottlerocket/agents/src/bin/k8s-workload-agent/main.rs
@@ -99,6 +99,7 @@ where
             &self.config,
             &self.results_dir,
             info_client,
+            &self.aws_secret_name.as_ref(),
         )
         .await
     }
@@ -115,7 +116,14 @@ where
             .await?;
         info!("Stored kubeconfig in {}", TEST_CLUSTER_KUBECONFIG_PATH);
 
-        rerun_failed_workload(TEST_CLUSTER_KUBECONFIG_PATH, &self.results_dir, info_client).await
+        rerun_failed_workload(
+            TEST_CLUSTER_KUBECONFIG_PATH,
+            &self.results_dir,
+            info_client,
+            &self.config,
+            &self.aws_secret_name.as_ref(),
+        )
+        .await
     }
 
     async fn terminate(&mut self) -> Result<(), Self::E> {

diff --git a/bottlerocket/agents/src/bin/sonobuoy-test-agent/main.rs b/bottlerocket/agents/src/bin/sonobuoy-test-agent/main.rs
@@ -46,8 +46,8 @@ use test_agent::{
 };
 use testsys_model::{SecretName, TestResults};
 
-// Default Sonobuoy agents assume role duration to 4 hours.
-const DEFAULT_ASSUME_ROLE_SESSION_DURATION: i32 = 14400;
+// Default Sonobuoy agents assume role duration to 1 hour.
+const DEFAULT_ASSUME_ROLE_SESSION_DURATION: i32 = 3600;
 
 struct SonobuoyTestRunner {
     config: SonobuoyConfig,
@@ -111,6 +111,7 @@ where
             &self.config,
             &self.results_dir,
             info_client,
+            &self.aws_secret_name.as_ref(),
         )
         .await
     }
@@ -153,6 +154,7 @@ where
             &self.config,
             &self.results_dir,
             info_client,
+            &self.aws_secret_name.as_ref(),
         )
         .await
     }

diff --git a/bottlerocket/agents/src/sonobuoy.rs b/bottlerocket/agents/src/sonobuoy.rs
@@ -1,4 +1,5 @@
 use crate::error;
+use agent_utils::aws::aws_config;
 use bottlerocket_types::agent_config::{SonobuoyConfig, SONOBUOY_RESULTS_FILENAME};
 use log::{error, info, trace};
 use serde_json::Value;
@@ -9,7 +10,7 @@ use std::path::Path;
 use std::process::Command;
 use std::time::Duration;
 use test_agent::InfoClient;
-use testsys_model::{Outcome, TestResults};
+use testsys_model::{Outcome, SecretName, TestResults};
 
 /// Timeout for sonobuoy status to become available (seconds)
 const SONOBUOY_STATUS_TIMEOUT: u64 = 900;
@@ -22,6 +23,7 @@ pub async fn run_sonobuoy<I>(
     sonobuoy_config: &SonobuoyConfig,
     results_dir: &Path,
     info_client: &I,
+    aws_secret_name: &Option<&SecretName>,
 ) -> Result<TestResults, error::Error>
 where
     I: InfoClient,
@@ -102,7 +104,14 @@ where
     .await
     .context(error::SonobuoyTimeoutSnafu)??;
     info!("Sonobuoy status is available, waiting for test to complete");
-    wait_for_sonobuoy_results(kubeconfig_path, None, info_client).await?;
+    wait_for_sonobuoy_results(
+        kubeconfig_path,
+        None,
+        info_client,
+        &sonobuoy_config.assume_role,
+        aws_secret_name,
+    )
+    .await?;
     info!("Sonobuoy testing has completed, checking results");
 
     results_sonobuoy(kubeconfig_path, results_dir)
@@ -115,6 +124,7 @@ pub async fn rerun_failed_sonobuoy<I>(
     sonobuoy_config: &SonobuoyConfig,
     results_dir: &Path,
     info_client: &I,
+    aws_secret_name: &Option<&SecretName>,
 ) -> Result<TestResults, error::Error>
 where
     I: InfoClient,
@@ -175,7 +185,14 @@ where
     .await
     .context(error::SonobuoyTimeoutSnafu)??;
     info!("Sonobuoy status is available, waiting for test to complete");
-    wait_for_sonobuoy_results(kubeconfig_path, None, info_client).await?;
+    wait_for_sonobuoy_results(
+        kubeconfig_path,
+        None,
+        info_client,
+        &sonobuoy_config.assume_role,
+        aws_secret_name,
+    )
+    .await?;
     info!("Sonobuoy testing has completed, checking results");
 
     results_sonobuoy(kubeconfig_path, results_dir)
@@ -220,6 +237,8 @@ pub async fn wait_for_sonobuoy_results<I>(
     kubeconfig_path: &str,
     namespace: Option<&str>,
     info_client: &I,
+    assume_role: &Option<String>,
+    aws_secret_name: &Option<&SecretName>,
 ) -> Result<(), error::Error>
 where
     I: InfoClient,
@@ -230,10 +249,20 @@ where
         ..Default::default()
     };
     let mut retries = 0;
+    // Max duration for assume role credential is 3600 seconds, and we refresh every 50 loops * 30 seconds (loop sleep time) before it expires.
+    let mut credential_refresh_countdown = 50;
+
     loop {
         if retries > 5 {
             return Err(error::Error::SonobuoyStatus { retries });
         }
+
+        // Refresh the credentials if the countdown is 0
+        if credential_refresh_countdown == 0 {
+            aws_config(aws_secret_name, assume_role, &None, &None, &None, true).await?;
+            credential_refresh_countdown = 50;
+        }
+
         let kubeconfig_arg = vec!["--kubeconfig", kubeconfig_path];
         let namespace_arg = namespace
             .map(|namespace| vec!["--namespace", namespace])
@@ -300,6 +329,7 @@ where
             .for_each(|e| error!("Unable to send test update: {}", e));
 
         tokio::time::sleep(Duration::from_secs(30)).await;
+        credential_refresh_countdown -= 1;
     }
 }
 

diff --git a/bottlerocket/agents/src/workload.rs b/bottlerocket/agents/src/workload.rs
@@ -11,7 +11,7 @@ use std::path::{Path, PathBuf};
 use std::process::Command;
 use std::time::Duration;
 use test_agent::InfoClient;
-use testsys_model::TestResults;
+use testsys_model::{SecretName, TestResults};
 
 /// Timeout for sonobuoy status to become available (seconds)
 const SONOBUOY_STATUS_TIMEOUT: u64 = 900;
@@ -24,6 +24,7 @@ pub async fn run_workload<I>(
     workload_config: &WorkloadConfig,
     results_dir: &Path,
     info_client: &I,
+    aws_secret_name: &Option<&SecretName>,
 ) -> Result<TestResults, error::Error>
 where
     I: InfoClient,
@@ -92,7 +93,14 @@ where
     .await
     .context(error::SonobuoyTimeoutSnafu)??;
     info!("Workload status is available, waiting for test to complete");
-    wait_for_sonobuoy_results(kubeconfig_path, Some("testsys-workload"), info_client).await?;
+    wait_for_sonobuoy_results(
+        kubeconfig_path,
+        Some("testsys-workload"),
+        info_client,
+        &workload_config.assume_role,
+        aws_secret_name,
+    )
+    .await?;
     info!("Workload testing has completed, checking results");
 
     results_workload(kubeconfig_path, results_dir)
@@ -103,6 +111,8 @@ pub async fn rerun_failed_workload<I>(
     kubeconfig_path: &str,
     results_dir: &Path,
     info_client: &I,
+    workload_config: &WorkloadConfig,
+    aws_secret_name: &Option<&SecretName>,
 ) -> Result<TestResults, error::Error>
 where
     I: InfoClient,
@@ -138,7 +148,14 @@ where
     .await
     .context(error::SonobuoyTimeoutSnafu)??;
     info!("Workload status is available, waiting for test to complete");
-    wait_for_sonobuoy_results(kubeconfig_path, Some("testsys-workload"), info_client).await?;
+    wait_for_sonobuoy_results(
+        kubeconfig_path,
+        Some("testsys-workload"),
+        info_client,
+        &workload_config.assume_role,
+        aws_secret_name,
+    )
+    .await?;
     info!("Workload testing has completed, checking results");
 
     results_workload(kubeconfig_path, results_dir)