Agent: Extract SSM logic to agent_utils

Author: ecpullen

Cargo.lock

@@ -10,6 +10,8 @@ agent-common = { version = "0.0.5", path = "../../agent/agent-common" }
 aws-config = "0.54"
 aws-credential-types = "0.54"
 aws-types = "0.54"
+aws-sdk-iam = "0.24"
+aws-sdk-ssm = "0.24"
 aws-sdk-sts = "0.24"
 aws-smithy-types = "0.54"
 base64 = "0.20"

agent/utils/Cargo.toml

@@ -1,3 +1,5 @@
+use aws_sdk_iam::error::{AttachRolePolicyError, CreateRoleError, GetRoleError};
+use aws_sdk_ssm::error::{CreateActivationError, DescribeInstanceInformationError};
 use aws_sdk_sts::error::AssumeRoleError;
 use aws_sdk_sts::types::SdkError;
 use snafu::Snafu;
@@ -12,18 +14,61 @@ pub enum Error {
         source: SdkError<AssumeRoleError>,
     },
 
+    #[snafu(display(
+        "Failed to attach policy '{}' to role '{}': {}",
+        policy_arn,
+        role_name,
+        source
+    ))]
+    AttachRolePolicy {
+        role_name: String,
+        policy_arn: String,
+        source: SdkError<AttachRolePolicyError>,
+    },
+
     #[snafu(display("Failed to decode base64 blob: {}", source))]
     Base64Decode { source: base64::DecodeError },
 
-    #[snafu(display("Failed to setup environment variables: {}", what))]
-    EnvSetup { what: String },
-
     #[snafu(display("Could not convert '{}' secret to string: {}", what, source))]
     Conversion { what: String, source: FromUtf8Error },
 
+    #[snafu(display("Failed to send create SSM command: {}", source))]
+    CreateSsmActivation {
+        source: SdkError<CreateActivationError>,
+    },
+
+    #[snafu(display(
+        "Unable to create role '{}' with policy '{}': {}",
+        role_name,
+        role_policy,
+        source
+    ))]
+    CreateRole {
+        role_name: String,
+        role_policy: String,
+        source: SdkError<CreateRoleError>,
+    },
+
     #[snafu(display("Credentials were missing for assumed role '{}'", role_arn))]
     CredentialsMissing { role_arn: String },
 
+    #[snafu(display("Failed to setup environment variables: {}", what))]
+    EnvSetup { what: String },
+
+    #[snafu(display("Unable to get managed instance information: {}", source))]
+    GetManagedInstanceInfo {
+        source: SdkError<DescribeInstanceInformationError>,
+    },
+
+    #[snafu(display("Unable to get SSM role '{}': {}", role_name, source))]
+    GetSSMRole {
+        role_name: String,
+        source: SdkError<GetRoleError>,
+    },
+
+    #[snafu(display("{} was missing from {}", what, from))]
+    Missing { what: String, from: String },
+
     #[snafu(display("Secret was missing: {}", source))]
     SecretMissing {
         source: agent_common::secrets::Error,
@@ -35,3 +80,5 @@ pub enum Error {
         source: std::io::Error,
     },
 }
+
+pub type Result<T> = std::result::Result<T, Error>;

agent/utils/src/error.rs

@@ -19,6 +19,7 @@ use std::{env, fs};
 pub mod aws;
 pub mod constants;
 mod error;
+pub mod ssm;
 
 /// Decode base64 blob and write to a file at the specified path
 pub async fn base64_decode_write_file(

agent/utils/src/lib.rs

@@ -1,8 +1,9 @@
+use crate::error::{self, Result};
 use aws_sdk_iam::types::SdkError;
 use aws_sdk_ssm::model::{InstanceInformation, InstanceInformationStringFilter, Tag};
 use log::info;
-use resource_agent::provider::{IntoProviderError, ProviderError, ProviderResult, Resources};
 use serde_json::json;
+use snafu::{OptionExt, ResultExt};
 use std::thread::sleep;
 use std::time::Duration;
 
@@ -12,9 +13,7 @@ const SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME: &str = "BR-SSMServiceRole";
 const SSM_MANAGED_INSTANCE_POLICY_ARN: &str =
     "arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore";
 
-pub(crate) async fn ensure_ssm_service_role(
-    iam_client: &aws_sdk_iam::Client,
-) -> ProviderResult<()> {
+pub async fn ensure_ssm_service_role(iam_client: &aws_sdk_iam::Client) -> Result<()> {
     let get_role_result = iam_client
         .get_role()
         .role_name(SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME)
@@ -38,23 +37,19 @@ pub(crate) async fn ensure_ssm_service_role(
                 iam_client
                     .create_role()
                     .role_name(SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME)
-                    .assume_role_policy_document(assume_role_doc)
+                    .assume_role_policy_document(&assume_role_doc)
                     .send()
                     .await
-                    .context(
-                        Resources::Clear,
-                        "Could not create SSM managed instance service role",
-                    )?;
+                    .context(error::CreateRoleSnafu {
+                        role_name: SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME,
+                        role_policy: assume_role_doc,
+                    })?;
             }
-            _ => {
-                return Err(ProviderError::new_with_source_and_context(
-                    Resources::Clear,
-                    format!(
-                        "Failed to determine whether SSM '{}' service role exists",
-                        SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME
-                    ),
-                    sdk_err,
-                ));
+            e => {
+                return Err(error::Error::GetSSMRole {
+                    role_name: SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME.to_string(),
+                    source: e,
+                });
             }
         }
     }
@@ -66,20 +61,19 @@ pub(crate) async fn ensure_ssm_service_role(
         .policy_arn(SSM_MANAGED_INSTANCE_POLICY_ARN)
         .send()
         .await
-        .context(
-            Resources::Clear,
-            "Could not attach SSM managed instance policy to the service role",
-        )?;
+        .context(error::AttachRolePolicySnafu {
+            role_name: SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME,
+            policy_arn: SSM_MANAGED_INSTANCE_POLICY_ARN,
+        })?;
 
     Ok(())
 }
 
-pub(crate) async fn create_ssm_activation(
-    resources: Resources,
+pub async fn create_ssm_activation(
     cluster_name: &str,
     num_registration: i32,
     ssm_client: &aws_sdk_ssm::Client,
-) -> ProviderResult<(String, String)> {
+) -> Result<(String, String)> {
     let activations = ssm_client
         .create_activation()
         .iam_role(SSM_MANAGED_INSTANCE_SERVICE_ROLE_NAME)
@@ -92,23 +86,24 @@ pub(crate) async fn create_ssm_activation(
         )
         .send()
         .await
-        .context(resources, "Failed to send SSM create activation command")?;
-    let activation_id = activations
-        .activation_id
-        .context(resources, "Unable to fetch SSM activation ID")?;
-    let activation_code = activations
-        .activation_code
-        .context(resources, "Unable to generate SSM activation code")?;
+        .context(error::CreateSsmActivationSnafu {})?;
+    let activation_id = activations.activation_id.context(error::MissingSnafu {
+        what: "activation id",
+        from: "activations",
+    })?;
+    let activation_code = activations.activation_code.context(error::MissingSnafu {
+        what: "activation code",
+        from: "activations",
+    })?;
     Ok((activation_id, activation_code))
 }
 
 // Waits for the SSM agent to be ready for a particular instance, returns the instance information
-pub(crate) async fn wait_for_ssm_ready(
-    resources: Resources,
+pub async fn wait_for_ssm_ready(
     ssm_client: &aws_sdk_ssm::Client,
     activation_id: &str,
     ip: &str,
-) -> ProviderResult<InstanceInformation> {
+) -> Result<InstanceInformation> {
     let seconds_between_checks = Duration::from_secs(5);
     loop {
         let instance_info = ssm_client
@@ -121,10 +116,7 @@ pub(crate) async fn wait_for_ssm_ready(
             )
             .send()
             .await
-            .context(
-                resources,
-                "Failed to get registered managed instance information",
-            )?;
+            .context(error::GetManagedInstanceInfoSnafu {})?;
         if let Some(info) = instance_info.instance_information_list().and_then(|list| {
             list.iter()
                 .find(|info| info.ip_address == Some(ip.to_string()))

bottlerocket/agents/src/bin/vsphere-vm-resource-agent/aws.rs renamed to agent/utils/src/ssm.rs

@@ -4,7 +4,6 @@ Provides Bottlerocket VMWare vSphere VMs to serve as Kubernetes nodes via `govc`
 
 !*/
 
-mod aws;
 mod vsphere_vm_provider;
 
 use crate::vsphere_vm_provider::{VMCreator, VMDestroyer};

bottlerocket/agents/src/bin/vsphere-vm-resource-agent/main.rs

@@ -1,6 +1,6 @@
-use crate::aws::{create_ssm_activation, ensure_ssm_service_role, wait_for_ssm_ready};
 use agent_utils::aws::aws_config;
 use agent_utils::base64_decode_write_file;
+use agent_utils::ssm::{create_ssm_activation, ensure_ssm_service_role, wait_for_ssm_ready};
 use bottlerocket_agents::constants::TEST_CLUSTER_KUBECONFIG_PATH;
 use bottlerocket_agents::tuf::{download_target, tuf_repo_urls};
 use bottlerocket_agents::userdata::{decode_to_string, merge_values};
@@ -139,7 +139,9 @@ impl Create for VMCreator {
             .context(resources, "Error sending cluster creation message")?;
 
         // Ensure we have a SSM service role we can attach to the VMs
-        ensure_ssm_service_role(&iam_client).await?;
+        ensure_ssm_service_role(&iam_client)
+            .await
+            .context(Resources::Clear, "Unable to check for SSM service role")?;
 
         info!("Getting vSphere secret");
         memo.current_status = "Getting vSphere secret".to_string();
@@ -343,8 +345,9 @@ impl Create for VMCreator {
 
         let vm_count = spec.configuration.vm_count.unwrap_or(DEFAULT_VM_COUNT);
         // Generate SSM activation codes and IDs
-        let activation =
-            create_ssm_activation(resources, &vsphere_cluster.name, vm_count, &ssm_client).await?;
+        let activation = create_ssm_activation(&vsphere_cluster.name, vm_count, &ssm_client)
+            .await
+            .context(resources, "Unable to create SSM activation")?;
         memo.ssm_activation_id = activation.0.to_owned();
         let control_host_ctr_userdata = json!({"ssm":{"activation-id": activation.0.to_string(), "activation-code":activation.1.to_string(),"region":"us-west-2"}});
         debug!(
@@ -461,13 +464,14 @@ impl Create for VMCreator {
 
             let instance_info = tokio::time::timeout(
                 Duration::from_secs(60),
-                wait_for_ssm_ready(resources, &ssm_client, &memo.ssm_activation_id, ip),
+                wait_for_ssm_ready(&ssm_client, &memo.ssm_activation_id, ip),
             )
             .await
             .context(
                 resources,
                 format!("Timed out waiting for SSM agent to be ready on VM '{}'", ip),
-            )??;
+            )?
+            .context(resources, "Unable to determine if SSM activation is ready")?;
 
             memo.vms.push(VSphereVM {
                 name: node_name,